add polarity case for unk residues

pegerto · Jun 4, 2024 · 0818070 · 0818070
1 parent 32ac386
commit 0818070
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git 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
@@ -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='')])