From 081807050fa6501e23e6e79306bfc23d26b52795 Mon Sep 17 00:00:00 2001
From: Pegerto Fernandez <pegerto.fernandez@creditkarma.com>
Date: Tue, 4 Jun 2024 14:49:29 +0100
Subject: [PATCH] add polarity case for unk residues

---
 src/graphpro/util/polarity.py     | 7 +++++--
 test/graphpro/annotations_test.py | 2 +-
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/src/graphpro/util/polarity.py b/src/graphpro/util/polarity.py
index b8d99b5..d5831ea 100644
--- a/src/graphpro/util/polarity.py
+++ b/src/graphpro/util/polarity.py
@@ -4,6 +4,7 @@
     a: apolar
     nc: negative charged
     pc: positive charged
+    unk: unknown
 """
 POLARITY = {
     'CYS': 'p', 
@@ -26,8 +27,10 @@
     'ASP': 'nc', 
     'LYS': 'nc', 
     'ARG': 'pc',
-    'HSD': 'pc'}
+    'HSD': 'pc',
+    '': 'unk' }
 
 def residue_polarity(resname):
-    return POLARITY[resname]
+    return POLARITY.get(resname, 'unk')
+
 POLARITY_CLASSES = list(set(POLARITY.values()))
diff --git a/test/graphpro/annotations_test.py b/test/graphpro/annotations_test.py
index 2475398..662192b 100644
--- a/test/graphpro/annotations_test.py
+++ b/test/graphpro/annotations_test.py
@@ -54,7 +54,7 @@ def test_polarity_encoding():
     G = md_analisys(u1).generate(ContactMap(cutoff=6), [Polarity()])
     data = G.to_data(node_encoders=[Polarity()])
     
-    assert data.x.size() == (214, 4)
+    assert data.x.size() == (214, 5)
     
 def test_encode_multiple_attributes():
     G = md_analisys(u1).generate(ContactMap(cutoff=6), [ResidueType(), SASAResArea(chain_id='')])