Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

filter molecule for templates #339

Merged
merged 14 commits into from
Jun 21, 2024
3 changes: 3 additions & 0 deletions bin/polyply
Original file line number Diff line number Diff line change
Expand Up @@ -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 '
Expand Down
2 changes: 1 addition & 1 deletion polyply/src/backmap.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
33 changes: 33 additions & 0 deletions polyply/src/check_residue_equivalence.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import networkx as nx
from vermouth.molecule import attributes_match
from collections import defaultdict

def check_residue_equivalence(topology):
"""
Expand Down Expand Up @@ -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
fgrunewald marked this conversation as resolved.
Show resolved Hide resolved

return unique_graphs
11 changes: 8 additions & 3 deletions polyply/src/gen_coords.py
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,7 @@ def gen_coords(toppath,
grid_spacing=0.2,
grid=None,
maxiter=800,
skip_filter=False,
start=[],
density=None,
box=None,
Expand Down Expand Up @@ -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:
fgrunewald marked this conversation as resolved.
Show resolved Hide resolved
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)
Expand Down
64 changes: 42 additions & 22 deletions polyply/src/generate_templates.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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):
fgrunewald marked this conversation as resolved.
Show resolved Hide resolved
"""
Given a `vermouth.molecule` and a `resname`
extract the information of a block from the
Expand All @@ -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
Expand All @@ -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"]
fgrunewald marked this conversation as resolved.
Show resolved Hide resolved

for inter_type in molecule.interactions:
for interaction in molecule.interactions[inter_type]:
Expand All @@ -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):
Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
3 changes: 2 additions & 1 deletion polyply/src/nonbond_engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
6 changes: 3 additions & 3 deletions polyply/tests/test_backmap.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading
Loading