Merge pull request #339 from marrink-lab/master_copy

filter molecule for templates

bin/polyply

@@ -138,6 +138,9 @@ def main(): # pylint: disable=too-many-locals,too-many-statements
                               help='file with grid-points', default=None)
     system_group.add_argument('-mi', dest='maxiter', type=int,
                               help='max number of trys to grow a molecule', default=800)
+    system_group.add_argument('-skip_filter', action='store_true', dest='skip_filter',
+                              help='do not group residues by isomophism when making '
+                                   'templates but just resname')
     system_group.add_argument('-start', dest='start', nargs='+', type=str,
                               default=[],
                               help=('Specify which residue to build first. The syntax is '

polyply/src/backmap.py

@@ -169,7 +169,7 @@ def _place_init_coords(self, meta_molecule):
         built_nodes = []
         for node in meta_molecule.nodes:
             if  meta_molecule.nodes[node]["backmap"]:
-                resname = meta_molecule.nodes[node]["resname"]
+                resname = meta_molecule.nodes[node]["template"]
                 cg_coord = meta_molecule.nodes[node]["position"]
                 resid = meta_molecule.nodes[node]["resid"]
                 high_res_atoms = meta_molecule.nodes[node]["graph"].nodes

polyply/src/check_residue_equivalence.py

@@ -1,5 +1,6 @@
 import networkx as nx
 from vermouth.molecule import attributes_match
+from collections import defaultdict
 
 def check_residue_equivalence(topology):
     """
@@ -42,3 +43,35 @@ def check_residue_equivalence(topology):
                 visited_residues[resname] = graph
                 molnames[resname] = mol_name
                 resids[resname] = molecule.nodes[node]["resid"]
+
+def group_residues_by_hash(meta_molecule, template_graphs={}):
+    """
+    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
+    ----------
+    meta_molecule: `:class:polyply.meta_molecule.MetaMolecule`
+    template_graphs: dict[`:class:nx.Graph`]
+
+    Returns
+    -------
+    dict[`:class:nx.Graph`]
+        keys are the hash of the graph
+    """
+    unique_graphs = {}
+    for graph in template_graphs.values():
+        graph_hash = nx.algorithms.graph_hashing.weisfeiler_lehman_graph_hash(graph, node_attr='atomname')
+        unique_graphs[graph_hash] = graph
+
+    for node in meta_molecule.nodes:
+        graph = meta_molecule.nodes[node]["graph"]
+        graph_hash = nx.algorithms.graph_hashing.weisfeiler_lehman_graph_hash(graph, node_attr='atomname')
+        if graph_hash not in unique_graphs:
+            unique_graphs[graph_hash] = graph
+        meta_molecule.nodes[node]["template"] = graph_hash
+
+    return unique_graphs

polyply/src/gen_coords.py

@@ -108,6 +108,7 @@ def gen_coords(toppath,
                grid_spacing=0.2,
                grid=None,
                maxiter=800,
+               skip_filter=False,
                start=[],
                density=None,
                box=None,
@@ -254,11 +255,15 @@ def gen_coords(toppath,
             LOGGER.warning(msg,  type="warning")
 
     # do a sanity check
-    LOGGER.info("checking residue integrity",  type="step")
-    check_residue_equivalence(topology)
+    if skip_filter:
+        LOGGER.info("checking residue integrity",  type="step")
+        check_residue_equivalence(topology)
+
     # Build polymer structure
     LOGGER.info("generating templates",  type="step")
-    GenerateTemplates(topology=topology, max_opt=10).run_system(topology)
+    GenerateTemplates(topology=topology,
+                      max_opt=10,
+                      skip_filter=skip_filter).run_system(topology)
     LOGGER.info("annotating ligands",  type="step")
     ligand_annotator = AnnotateLigands(topology, ligands)
     ligand_annotator.run_system(topology)

polyply/src/generate_templates.py

@@ -21,13 +21,30 @@
                                radius_of_gyration)
 from .topology import replace_defined_interaction
 from .linalg_functions import dih
+from .check_residue_equivalence import group_residues_by_hash
+from tqdm import tqdm
 """
 Processor generating coordinates for all residues of a meta_molecule
 matching those in the meta_molecule.molecule attribute.
 """
 
 LOGGER = StyleAdapter(get_logger(__name__))
 
+def _extract_template_graphs(meta_molecule, template_graphs={}, skip_filter=False):
+    if skip_filter:
+        for node in meta_molecule.nodes:
+            resname = meta_molecule.nodes[node]["resname"]
+            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_hash(meta_molecule, template_graphs)
+    return template_graphs
+
 def find_atoms(molecule, attr, value):
     """
     Find all nodes of a `vermouth.molecule.Molecule` that have the
@@ -239,7 +256,7 @@ def _relabel_interaction_atoms(interaction, mapping):
     new_interaction = interaction._replace(atoms=new_atoms)
     return new_interaction
 
-def extract_block(molecule, resname, defines):
+def extract_block(molecule, template_graph, defines):
     """
     Given a `vermouth.molecule` and a `resname`
     extract the information of a block from the
@@ -249,16 +266,15 @@ def extract_block(molecule, resname, 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
 
     Returns
     -------
     :class:vermouth.molecule.Block
     """
-    nodes = find_atoms(molecule, "resname", resname)
-    resid = molecule.nodes[nodes[0]]["resid"]
     block = vermouth.molecule.Block()
 
     # select all nodes with the same first resid and
@@ -267,11 +283,10 @@ def extract_block(molecule, resname, defines):
     # label in the molecule and in the block for
     # relabeling the interactions
     mapping = {}
-    for node in nodes:
+    for node in template_graph.nodes:
         attr_dict = molecule.nodes[node]
-        if attr_dict["resid"] == resid:
-            block.add_node(attr_dict["atomname"], **attr_dict)
-            mapping[node] = attr_dict["atomname"]
+        block.add_node(attr_dict["atomname"], **attr_dict)
+        mapping[node] = attr_dict["atomname"]
 
     for inter_type in molecule.interactions:
         for interaction in molecule.interactions[inter_type]:
@@ -284,12 +299,6 @@ def extract_block(molecule, resname, defines):
                        "virtual_sites2", "virtual_sites3", "virtual_sites4"]:
         block.make_edges_from_interaction_type(inter_type)
 
-    if not nx.is_connected(block):
-        msg = ('\n Residue {} with id {} consistes of two disconnected parts. '
-               'Make sure all atoms/particles in a residue are connected by bonds,'
-               ' constraints or virual-sites.')
-        raise IOError(msg.format(resname, resid))
-
     return block
 
 class GenerateTemplates(Processor):
@@ -300,14 +309,15 @@ class GenerateTemplates(Processor):
     in the templates attribute. The processor also stores the volume
     of each block in the volume attribute.
     """
-    def __init__(self, topology, max_opt, *args, **kwargs):
+    def __init__(self, topology, max_opt, skip_filter, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.max_opt = max_opt
         self.topology = topology
         self.volumes = self.topology.volumes
         self.templates = {}
+        self.skip_filter = skip_filter
 
-    def gen_templates(self, meta_molecule):
+    def gen_templates(self, meta_molecule, template_graphs):
         """
         Generate blocks for each unique residue by extracting the
         block information, placing initial coordinates, and geometry
@@ -318,17 +328,18 @@ class variable.
         Parameters
         ----------
         meta_molecule: :class:`polyply.src.meta_molecule.MetaMolecule`
+        template_graphs: dict
+            a dict of graphs corresponding to the templates to be generated
 
         Updates
         ---------
         self.templates
         self.volumes
         """
-        resnames = set(nx.get_node_attributes(meta_molecule, "resname").values())
-
-        for resname in resnames:
+        for resname, template_graph in tqdm(template_graphs.items()):
             if resname not in self.templates:
-                block = extract_block(meta_molecule.molecule, resname,
+                block = extract_block(meta_molecule.molecule,
+                                      template_graph,
                                       self.topology.defines)
 
                 opt_counter = 0
@@ -366,7 +377,16 @@ def run_molecule(self, meta_molecule):
         and volume attribute.
         """
         if hasattr(meta_molecule, "templates"):
-            self.templates.update(meta_molecule.templates)
-        self.gen_templates(meta_molecule)
+            template_graphs = {res: None for res in meta_molecule.templates}
+        else:
+            template_graphs = {}
+            meta_molecule.templates = {}
+
+        template_graphs = _extract_template_graphs(meta_molecule,
+                                                   skip_filter=self.skip_filter,
+                                                   template_graphs=template_graphs)
+        self.templates.update(meta_molecule.templates)
+
+        self.gen_templates(meta_molecule, template_graphs)
         meta_molecule.templates = self.templates
         return meta_molecule

polyply/src/nonbond_engine.py

@@ -384,7 +384,8 @@ def from_topology(cls, molecules, topology, box):
                                 "Make sure all coordiantes are wrapped inside the box.")
                         raise IOError(msg)
 
-                resname = molecule.nodes[node]["resname"]
+                # try getting the template name; if no template name is given use resname
+                resname = molecule.nodes[node].get("template", molecule.nodes[node]["resname"])
                 atom_types.append(resname)
                 nodes_to_gndx[(mol_count, node)] = idx
                 idx += 1

polyply/tests/test_backmap.py

@@ -43,13 +43,13 @@ def test_backmapping():
     meta_molecule = MetaMolecule(graph)
     meta_molecule.molecule = molecule
 
-    nx.set_node_attributes(meta_molecule, {0: {"resname": "test",
+    nx.set_node_attributes(meta_molecule, {0: {"resname": "test", "template": "test",
                                                "position": np.array([0, 0, 0]),
                                                "resid": 1, "backmap": True},
-                                           1: {"resname": "test",
+                                           1: {"resname": "test", "template": "test",
                                                "position": np.array([0, 0, 1.0]),
                                                "resid": 2, "backmap": True},
-                                           2: {"resname": "test",
+                                           2: {"resname": "test", "template": "test",
                                                "position": np.array([0, 0, 2.0]),
                                                "resid": 3, "backmap": False}})
     # test if disordered template works