fix-137, remove node functionality editted. (#157)

* fix-137, remove node functionality editted. * Fixing the bugs found after adding tests. * Adding test cases, adding a test for remove node and separating common code into a new method. * Changes in the test cases for remove node. Changed the common code method name and added a docstring. * Indenting properly * Changes in molecules code.
marrink-lab · Nov 14, 2018 · 6136310 · 6136310
1 parent 5d835c5
commit 6136310
Show file tree

Hide file tree

Showing 2 changed files with 125 additions and 2 deletions.
diff --git a/vermouth/molecule.py b/vermouth/molecule.py
@@ -16,7 +16,6 @@
 from collections import defaultdict, OrderedDict, namedtuple
 import copy
 from functools import partial
-import itertools
 
 import networkx as nx
 import numpy as np
@@ -337,7 +336,7 @@ def force_field(self):
     def atoms(self):
         """
         All atoms in this molecule. Alias for `nodes`.
-        
+
         See Also
         --------
         :attr:`networkx.Graph.nodes`
@@ -674,6 +673,42 @@ def edges_between(self, n_bunch1, n_bunch2, data=False):
                 else:
                     yield (node1, node2, self.edges[node1, node2])
 
+    def _remove_interactions_with_node(self, node):
+        """
+        We iterate through the different interactions we have and
+        remove the interactions where the atoms to be deleted are present.
+        Further we also delete the entire interaction_type if it is
+        empty after all the necessary interactions have been deleted.
+        """
+        for name, interactions in self.interactions.items():
+            for interaction in interactions:
+                if node in interaction.atoms:
+                    self.interactions[name].remove(interaction)
+
+        for interaction_type in list(self.interactions):
+            if not self.interactions[interaction_type]:
+                self.interactions.pop(interaction_type)
+
+    def remove_node(self, node):
+        """
+        Overriding the remove_node method of networkx
+        as we have to delete the interaction from the interactions list
+        separately which is not a part of the graph and hence does not
+        get deleted.
+        """
+        super().remove_node(node)
+        self._remove_interactions_with_node(node)
+
+    def remove_nodes_from(self, nodes):
+        """
+        Overriding the remove_nodes_from method of networkx
+        as we have to delete the interaction from the
+        interactions list separately which is not a part of
+        the graph and hence does not get deleted.
+        """
+        super().remove_nodes_from(nodes)
+        for node in nodes:
+            self._remove_interactions_with_node(node)
 
 class Block(Molecule):
     """

diff --git a/vermouth/tests/test_molecule.py b/vermouth/tests/test_molecule.py
@@ -47,6 +47,94 @@ def molecule_copy(molecule):
     return molecule.copy()
 
 
+@pytest.mark.parametrize('atoms, bonds, interactions, removed, expected', [
+
+    # empty molecule
+    ([], [], [], [], {}),
+
+    # Nodes but no interactions
+    ([1, 2, 3], [(1, 2), (2, 3)], [], [2], {}),
+
+    # interactions that all need to be removed
+    ([1, 2], [(1, 2)], [('bond', (1, 2), {})], [2], {}),
+
+    # Molecule with interactions of which some need to be removed
+    ([1, 2, 3, 4, 5, 6, 7], [(1, 2), (3, 4), (5, 7)], [('bond', (5, 4), {}), ('bond', (1, 2), {})], [5]
+     , {'bond': [vermouth.molecule.Interaction(atoms=(1, 2), meta={}, parameters={})]}),
+
+    # Molecule with interactions of which none need to be removed
+    ([1, 2, 3, 4, 5, 6, 7], [(1, 2), (3, 4)], [('bond', (5, 4), {}), ('bond', (1, 2), {})], [6, 7]
+     , {'bond': [vermouth.molecule.Interaction(atoms=(5, 4), meta={}, parameters={})
+                 , vermouth.molecule.Interaction(atoms=(1, 2), meta={}, parameters={})]}),
+
+    # Molecule with interactions of different types of which all need to be removed
+    ([1, 2, 3, 4, 5, 6, 7], [(1, 2), (3, 4), (5, 7), (6, 7)], [('bond', (1, 6), {}), ('bond', (5, 4), {})
+                                                               , ('angle', (1, 2), {})], [1, 5], {}),
+
+    # Molecule with interactions of different types of which some need to be removed
+    ([1, 2, 3, 4, 5, 6, 7], [(1, 2), (3, 4), (5, 7), (6, 7)]
+     , [('bond', (1, 6), {}), ('bond', (5, 4), {}), ('angle', (1, 2), {})], [5]
+     , {'angle': [vermouth.molecule.Interaction(atoms=(1, 2), meta={}, parameters={})]
+                 , 'bond': [vermouth.molecule.Interaction(atoms=(1, 6), meta={}, parameters={})]}),
+
+    # Molecule with interactions of different types of which none need to be removed
+    ([1, 2, 3, 4, 5, 6, 7], [(1, 2), (3, 4), (5, 7), (6, 7)]
+     , [('bond', (1, 6), {}), ('bond', (5, 4), {}), ('angle', (1, 2), {})], [3]
+     , {'angle': [vermouth.molecule.Interaction(atoms=(1, 2), meta={}, parameters={})]
+                 , 'bond': [vermouth.molecule.Interaction(atoms=(1, 6), meta={}, parameters={})
+                            , vermouth.molecule.Interaction(atoms=(5, 4), meta={}, parameters={})]}),
+])
+def test_remove_nodes_from(atoms, bonds, interactions, removed, expected):
+    molecule = vermouth.molecule.Molecule()
+    molecule.add_nodes_from(atoms)
+    molecule.add_edges_from(bonds)
+    for type_, atoms, params in interactions:
+        molecule.add_interaction(type_, atoms, params)
+
+    molecule.remove_nodes_from(removed)
+    assert molecule.interactions == expected
+
+
+@pytest.mark.parametrize('atoms, bonds, interactions, removed, expected', [
+
+    # empty molecule
+    ([], [], [], None, {}),
+
+    # Nodes but no interactions
+    ([1, 2, 3], [(1, 2), (2, 3)], [], 2, {}),
+
+    # interactions that all need to be removed
+    ([1, 2], [(1, 2)], [('bond', (1, 2), {})], 2, {}),
+
+    # Molecule with interactions of which some need to be removed
+    ([1, 2, 3, 4, 5, 6, 7], [(1, 2), (3, 4), (5, 7)], [('bond', (5, 4), {}), ('bond', (1, 2), {})], 5
+     , {'bond': [vermouth.molecule.Interaction(atoms=(1, 2), meta={}, parameters={})]}),
+
+    # Molecule with interactions of different types of which some need to be removed
+    ([1, 2, 3, 4, 5, 6, 7], [(1, 2), (3, 4), (5, 7), (6, 7)]
+     , [('bond', (1, 6), {}), ('bond', (5, 4), {}), ('angle', (1, 2), {})], 5
+     , {'angle': [vermouth.molecule.Interaction(atoms=(1, 2), meta={}, parameters={})],
+        'bond': [vermouth.molecule.Interaction(atoms=(1, 6), meta={}, parameters={})]}),
+
+    # Molecule with interactions of different types of which none need to be removed
+    ([1, 2, 3, 4, 5, 6, 7], [(1, 2), (3, 4), (5, 7), (6, 7)]
+     , [('bond', (1, 6), {}), ('bond', (5, 4), {}), ('angle', (1, 2), {})], 3
+     , {'angle': [vermouth.molecule.Interaction(atoms=(1, 2), meta={}, parameters={})],
+        'bond': [vermouth.molecule.Interaction(atoms=(1, 6), meta={}, parameters={})
+                 , vermouth.molecule.Interaction(atoms=(5, 4), meta={}, parameters={})]}),
+])
+def test_remove_node(atoms, bonds, interactions, removed, expected):
+    molecule = vermouth.molecule.Molecule()
+    molecule.add_nodes_from(atoms)
+    molecule.add_edges_from(bonds)
+    for type_, atoms, params in interactions:
+        molecule.add_interaction(type_, atoms, params)
+
+    if removed:
+        molecule.remove_node(removed)
+
+    assert molecule.interactions == expected
+
 @pytest.fixture
 def molecule_subgraph(molecule):
     return molecule.subgraph([2, 0])