filter molecule for templates

polyply/src/check_residue_equivalence.py

@@ -44,12 +44,13 @@ def check_residue_equivalence(topology):
                 molnames[resname] = mol_name
                 resids[resname] = molecule.nodes[node]["resid"]
 
-def group_residues_by_isomorphism(meta_molecule, template_graphs={}):
+def group_residues_by_hash(meta_molecule, template_graphs={}):
     """
-    Collect all unique residue graphs. If the same resname matches
-    multiple graphs the resname is appended by a number. If required
-    template_graphs can be given that are used for matching rather
-    than the first founds residue.
+    Collect all unique residue graphs using the Weisfeiler-Lehman has.
+    A dict of unique graphs with the hash as key is returned. The
+    `meta_molecule` nodes are annotated with the hash using the template
+    keyword. If required template_graphs can be given that are used for
+    matching rather than the first founds residue.
 
     Parameters
     ----------

polyply/src/generate_templates.py

@@ -21,7 +21,7 @@
                                radius_of_gyration)
 from .topology import replace_defined_interaction
 from .linalg_functions import dih
-from .check_residue_equivalence import group_residues_by_isomorphism
+from .check_residue_equivalence import group_residues_by_hash
 from tqdm import tqdm
 """
 Processor generating coordinates for all residues of a meta_molecule
@@ -34,10 +34,15 @@ def _extract_template_graphs(meta_molecule, template_graphs={}, skip_filter=Fals
     if skip_filter:
         for node in meta_molecule.nodes:
             resname = meta_molecule.nodes[node]["resname"]
-            if resname not in template_graphs:
-                template_graphs[resname] = meta_molecule.nodes[node]["graph"]
+            graph = meta_molecule.nodes[node]["graph"]
+            graph_hash = nx.algorithms.graph_hashing.weisfeiler_lehman_graph_hash(graph, node_attr='atomname')
+            if resname in template_graphs:
+                template_graphs[graph_hash] = graph
+                del template_graphs[resname]
+            elif resname not in template_graphs and graph_hash not in template_graphs:
+                template_graphs[graph_hash] = graph
     else:
-        template_graphs = group_residues_by_isomorphism(meta_molecule, template_graphs)
+        template_graphs = group_residues_by_hash(meta_molecule, template_graphs)
     return template_graphs
 
 def find_atoms(molecule, attr, value):
@@ -261,7 +266,8 @@ def extract_block(molecule, template_graph, defines):
     Parameters
     ----------
     molecule:  :class:vermouth.molecule.Molecule
-    resname:   str
+    template_graph: :class:`nx.Graph`
+        the graph of the template reisdue
     defines:   dict
       dict of type define: value
 

polyply/tests/test_generate_templates.py

@@ -31,7 +31,9 @@
                                             _relabel_interaction_atoms,
                                             compute_volume, map_from_CoG,
                                             extract_block, GenerateTemplates,
-					    find_interaction_involving)
+					                        find_interaction_involving,
+                                            _extract_template_graphs)
+from .example_fixtures import example_meta_molecule
 
 class TestGenTemps:
 
@@ -307,3 +309,36 @@ def test_compute_volume(lines, coords, volume):
     assert np.isclose(new_vol, volume, atol=0.000001)
 
 
+@pytest.mark.parametrize('resnames, gen_template_graphs, skip_filter', (
+                        # two different residues no template_graphs
+                        (['A', 'B', 'A'], [], False),
+                        # two different residues no template_graphs
+                        (['A', 'B', 'A'], [], True),
+                        # two different residues one template_graphs
+                        (['A', 'B', 'A'], [1], True),
+                        # two different residues one template_graphs
+                        (['A', 'B', 'A'], [1], False),
+))
+def test_extract_template_graphs(example_meta_molecule, resnames, gen_template_graphs, skip_filter):
+    # set the residue names
+    for resname, node in zip(resnames, example_meta_molecule.nodes):
+        example_meta_molecule.nodes[node]['resname'] = resname
+        nx.set_node_attributes(example_meta_molecule.nodes[node]['graph'], resname, 'resname')
+
+    # extract template graphs if needed
+    template_graphs = {}
+    for node in gen_template_graphs:
+        graph = example_meta_molecule.nodes[node]['graph']
+        nx.set_node_attributes(graph, True, 'template')
+        template_graphs[example_meta_molecule.nodes[node]['resname']] = graph
+
+    # perfrom the grouping
+    unique_graphs = _extract_template_graphs(example_meta_molecule, template_graphs, skip_filter)
+
+    # check the outcome
+    assert len(unique_graphs) == 2
+
+    for graph in template_graphs.values():
+        graph_hash = nx.algorithms.graph_hashing.weisfeiler_lehman_graph_hash(graph, node_attr='atomname')
+        templated = list(nx.get_node_attributes(unique_graphs[graph_hash], 'template').values())
+        assert all(templated)