return 0 for not know sasa encodings

pegerto · Jan 3, 2024 · 1ea470f · 1ea470f
1 parent 3365a27
commit 1ea470f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/graphpro/annotations.py b/src/graphpro/annotations.py
@@ -75,5 +75,5 @@ def generate(self, G: Graph, atom_group: AtomGroup):
             G.node_attr_add(node_id, self.attr_name, chain_res[k].total)
 
     def encode(self, G: Graph) -> torch.tensor:
-        total_area = [G.node_attr(n)[self.attr_name] for n in G.nodes()]
+        total_area = [G.node_attr(n)[self.attr_name] if self.attr_name in G.node_attr(n) else 0 for n in G.nodes()]
         return torch.tensor([total_area], dtype=torch.float).T