Repository: marrink-lab/vermouth-martinize
Issue: DoAverageBead breaks if a node has no constructing atoms

In case e.g. an amino acid sidechain is completely missing there are no atoms constructing the SC1 bead (in Martini). `do_average_bead` does not like this, since `np.stack([])` raises an error.

The fix should be a fairly straightforward if/else statement that sets the position to [nan, nan, nan] if there are no constructing atoms.
Related to #17 

Current work-around for users: make sure there is at least 1 atom per bead present in the input.

Found with help from Sebastian.

Changed files:
vermouth/processors/average_beads.py, vermouth/tests/test_average_beads.py

diff --git a/vermouth/processors/average_beads.py b/vermouth/processors/average_beads.py
index b126dd7..8c57f2d 100644
--- a/vermouth/processors/average_beads.py
+++ b/vermouth/processors/average_beads.py
@@ -77,7 +77,7 @@ def do_average_bead(molecule, ignore_missing_graphs=False, weight=None):
 
     for node in molecule.nodes.values():
         if 'graph' in node:
-            positions = np.stack([
+            positions = np.array([
                 subnode['position']
                 for subnode in node['graph'].nodes().values()
                 if subnode.get('position') is not None
@@ -87,7 +87,14 @@ def do_average_bead(molecule, ignore_missing_graphs=False, weight=None):
                 for subnode_key, subnode in node['graph'].nodes.items()
                 if subnode.get('position') is not None
             ])
-            node['position'] = np.average(positions, axis=0, weights=weights)
+            try:
+                ndim = positions.shape[1]
+            except IndexError:
+                ndim = 3
+            if abs(sum(weights)) < 1e-7:
+                node['position'] = np.array([np.nan]*ndim, dtype=float)
+            else:
+                node['position'] = np.average(positions, axis=0, weights=weights)
 
     return molecule
 
diff --git a/vermouth/tests/test_average_beads.py b/vermouth/tests/test_average_beads.py
index 4720fea..07fdc12 100644
--- a/vermouth/tests/test_average_beads.py
+++ b/vermouth/tests/test_average_beads.py
@@ -59,6 +59,30 @@ def mol_with_subgraph():
         "target False": np.array([2.6, 3.6, 4.6]),
     })
 
+    subgraph = nx.Graph()
+    mol.add_node(2, **{
+        "graph": subgraph,
+        "mapping_weights": {},
+        "target mass": np.array([np.nan, np.nan, np.nan]),
+        "target not mass": np.array([np.nan, np.nan, np.nan]),
+        "target None": np.array([np.nan, np.nan, np.nan]),
+        "target False": np.array([np.nan, np.nan, np.nan]),
+    })
+
+    subgraph = nx.Graph()
+    subgraph.add_nodes_from((
+        (0, {'mass': 1.2, 'not mass': 2.2, 'position': np.array([2, 3, 4], dtype=float),}),
+        (1, {'mass': 1.3, 'not mass': 2.3, 'position': np.array([3, 4, 5], dtype=float),}),
+    ))
+    mol.add_node(3, **{
+        "graph": subgraph,
+        "mapping_weights": {0: 0, 1: 0},
+        "target mass": np.array([np.nan, np.nan, np.nan]),
+        "target not mass": np.array([np.nan, np.nan, np.nan]),
+        "target None": np.array([np.nan, np.nan, np.nan]),
+        "target False": np.array([np.nan, np.nan, np.nan]),
+    })
+
     return mol
 
 @pytest.fixture(params=(None, 'mass', 'not mass'))
@@ -89,7 +113,7 @@ def test_do_average_bead(mol_with_subgraph, weight):
     target_key = 'target {}'.format(weight)
     target_positions = np.stack([node[target_key] for node in mol_with_subgraph.nodes.values()])
     positions = np.stack([node['position'] for node in mol_with_subgraph.nodes.values()])
-    assert np.allclose(positions, target_positions)
+    assert np.allclose(positions, target_positions, equal_nan=True)
 
 
 @pytest.mark.parametrize('weight', ('mass', 'not mass'))
@@ -120,7 +144,7 @@ def test_processor_variable(mol_with_variable):
     target_key = 'target {}'.format(weight)
     target_positions = np.stack([node[target_key] for node in mol_with_variable.nodes.values()])
     positions = np.stack([node['position'] for node in mol_with_variable.nodes.values()])
-    assert np.allclose(positions, target_positions)
+    assert np.allclose(positions, target_positions, equal_nan=True)
 
 
 @pytest.mark.parametrize('weight', (False, 'mass', 'not mass'))
@@ -130,4 +154,4 @@ def test_processor_weight(mol_with_variable, weight):
     target_key = 'target {}'.format(weight)
     target_positions = np.stack([node[target_key] for node in mol_with_variable.nodes.values()])
     positions = np.stack([node['position'] for node in mol_with_variable.nodes.values()])
-    assert np.allclose(positions, target_positions)
+    assert np.allclose(positions, target_positions, equal_nan=True)