Order from rdkit

bin/ff_map

@@ -139,10 +139,10 @@ def main():
 
     if cgsmiles_strs:
         print("INFO - Mapping CGsmiles to universe.")
-        mapped_atoms, bead_idxs, mappings, weights = cgsmiles_to_mapping(init_universe,
-                                                                         cgsmiles_strs,
-                                                                         args.mol_names,
-                                                                         mol_matching=args.cgsorder)
+        mapped_atoms, bead_idxs, mappings, weights, res_iter = cgsmiles_to_mapping(init_universe,
+                                                                                  cgsmiles_strs,
+                                                                                  args.mol_names,
+                                                                                  mol_matching=args.cgsorder)
     n_frames = len(init_universe.trajectory)
     print("INFO - Mapping universe - positions")
     # extract the position array from universe
@@ -165,7 +165,11 @@ def main():
                                               treated_atoms)
 
     print("INFO - Mapping universe - building pos-array")
-    cg_universe = create_new_universe(init_universe, mapped_trajectory, mappings)
+    cg_universe = create_new_universe(init_universe,
+                                      mapped_trajectory,
+                                      mappings,
+                                      n_residues=None,
+                                      res_iter=res_iter)
     # write coordinate
     print("INFO - Writing CG trajectory")
     if args.traj:

fast_forward/cgsmiles.py

@@ -16,8 +16,21 @@
 import networkx as nx
 from cgsmiles.resolve import MoleculeResolver
 from .map_file_parers import Mapping
+from .universe_handler import ResidueIter
 
 def load_cgsmiles_library(filepath):
+    """
+    Given a file that lists cgsmiles strings,
+    read and return a list.
+
+    Parameters
+    ----------
+    filepath: str
+
+    Returns
+    -------
+    list[str]
+    """
     cgsmiles_strs = []
     with open(filepath) as _file:
         for line in _file.readlines():
@@ -27,7 +40,18 @@ def load_cgsmiles_library(filepath):
 def find_one_graph_match(graph1, graph2):
     """
     Returns one ismags match when graphs are isomorphic
-    otherwise [].
+    otherwise []. Matches are done based on element and
+    bond order. That should be sufficent to also account
+    for charges implicitly via the orders.
+
+    Parameters
+    ----------
+    graph1: networkx.Graph
+    graph2: networkx.Graph
+
+    Returns
+    -------
+    abc.iteratable
     """
 
     def node_match(n1, n2):
@@ -60,6 +84,23 @@ def get_mappings(cg, univ, _match, mappings):
     based on the all-atom univ. If a bead is split between
     residues we simply assing the resname and id of the
     majority of atoms.
+
+    Parameters
+    ----------
+    cg: networkx.graph
+        the coarse molecule
+    univ: :class:fast_forward.universe_handler.UniverseHandler
+        the universe stroing the fine grained molecule information
+    _match: dict
+        a dict mapping the coarse atoms to fine grained atoms
+    mappings:
+        dict[:class:fast_forward.map_file_parser.Mapping]
+        a dict of mapping objects
+
+    Return
+    ------
+    dict[:class:fast_forward.map_file_parser.Mapping]
+        updated mappings
     """
     target_resids = {}
     for bead in cg.nodes:
@@ -68,20 +109,40 @@ def get_mappings(cg, univ, _match, mappings):
         resname = _most_common(resnames)
         resids = [univ.atoms[_match[atom]].resid for atom in atoms]
         resid = _most_common(resids)
+        cg.nodes[bead]['resname'] = resname
+        cg.nodes[bead]['resid'] = resid
         mapping = mappings.get(resname, Mapping(resname, resname))
         target_resid = target_resids.get(resname, resid)
         target_resids[resname] = target_resid
         if resid != target_resids[resname]:
             continue
         for adx, atom in enumerate(atoms):
             weight = np.float32(cg.nodes[bead]['graph'].nodes[atom].get('weight', 1))
-            mapping.add_atom(cg.nodes[bead]["fragname"]+f"{bead}",
+            mapping.add_atom(cg.nodes[bead]["fragname"], #+f"{bead}",
                              _match[atom],
                              atom=univ.atoms[_match[atom]].name,
                              weight=weight)
         mappings[resname] = mapping
     return mappings
 
+def _annotate_vs(cg_graph):
+    """
+    If a node in the cg_graph is not mapped from atom positions,
+    it is a virtual site. Thus, we will compose the mapped position
+    as the union framgment graphs of neihgboring atoms. This appraoch
+    should be equivalent to a virtual_sitesn mapping.
+
+    Parameters
+    ----------
+    cg_graph: netwrokx.Graph
+        the graph describing a coarse molecule
+    """
+    for node in cg_graph.nodes:
+        if len(cg_graph.nodes[node].get('graph',[])) == 0:
+            g = nx.Graph()
+            for neigh in cg_graph.neighbors(node):
+                g.add_nodes_from(list(cg_graph.nodes[neigh]['graph'].nodes))
+            cg_graph.nodes[node]['graph'] = g
 
 def cgsmiles_to_mapping(univ, cgsmiles_strs, mol_names, mol_matching=True):
     """
@@ -94,7 +155,9 @@ def cgsmiles_to_mapping(univ, cgsmiles_strs, mol_names, mol_matching=True):
     univ: :class:`fast_forward.UniverseHandler`
     cgsmiles_strs: list[str]
     mol_names: list[str]
-
+    mol_matching: bool
+        if is False the order of cgsmiles strings is equivalent
+        to the order of molecules
 
     Retunrs
     -------
@@ -114,6 +177,7 @@ def cgsmiles_to_mapping(univ, cgsmiles_strs, mol_names, mol_matching=True):
             # read the cgsmiles string
             resolver = MoleculeResolver.from_string(cgs_str, last_all_atom=True)
             cg, aa = resolver.resolve_all()
+            _annotate_vs(cg)
             _match = find_one_graph_match(aa, mol_graph)
             if _match:
                 break
@@ -126,19 +190,24 @@ def cgsmiles_to_mapping(univ, cgsmiles_strs, mol_names, mol_matching=True):
 
         offset=0
         target_resids = {}
+        res_iter = ResidueIter()
+        prev_resid = -1
         for mol_idx in univ.mol_idxs_by_name[mol_name]:
             for bead in cg.nodes:
+                resid = cg.nodes[bead]['resid']
+                resname = cg.nodes[bead]['resname']
+                if prev_resid != resid:
+                    prev_resid = resid
+                    res_iter.add_residue(resid=resid, resname=resname)
                 atoms = cg.nodes[bead]['graph'].nodes
-                # catches VS
-                if len(atoms) == 0:
-                    continue
                 mapped_atoms.append(numba.typed.List([_match[atom]+offset for atom in atoms]))
                 set_weights = nx.get_node_attributes(cg.nodes[bead]['graph'], 'weight')
                 weights.append(np.array([set_weights.get(node, 1.0) for node in atoms], dtype=np.float32))
                 bead_idxs.append(bead_count)
                 bead_count += 1
             offset += len(mol_graph)
+
     mapped_atoms = numba.typed.List(mapped_atoms)
     bead_idxs = numba.typed.List(bead_idxs)
     weights = numba.typed.List(weights)
-    return mapped_atoms, bead_idxs, mappings, weights
+    return mapped_atoms, bead_idxs, mappings, weights, res_iter

fast_forward/itp_to_ag.py

@@ -7,7 +7,7 @@
 import networkx as nx
 from fast_forward.universe_handler import res_as_mol
 
-def find_mol_indices(universe, atoms, moltype):
+def find_mol_indices(universe, atoms, moltype, atomnames=None, match_names=False):
     """
     Given a universe select all atoms that belong to molecules 
     of the given `moltype`. Subsequently, return indices of all
@@ -26,13 +26,21 @@ def find_mol_indices(universe, atoms, moltype):
     """
     mol_atoms = universe.select_atoms(f'moltype {moltype}')
     n_mols = len(np.unique(mol_atoms.molnums))
+
     try:
         mol_atom_indices = mol_atoms.indices.reshape(n_mols, -1)
     except ValueError:
         msg = ("The target molecules passed to find_mol_indices "
                "do not seem to all have the same number of atoms.")
         raise IndexError(msg) from None
-    return list(mol_atom_indices[:, atoms])
+
+    if match_names:
+        target_atoms = universe.atoms.names[mol_atoms].reshape(n_mols, 1)
+        order = [np.where(target_atoms == r)[0][0] for r in ref]
+        mapped_indices = mol_atom_indices[:, atoms][order]
+        return mapped_indices
+    else:
+        return list(mol_atom_indices[:, atoms])
 
 class ITPInteractionMapper:
     """
@@ -65,8 +73,8 @@ def get_interactions_group(self, molname, itp_mode=False):
         for inter_type in block.interactions:
             for inter in block.interactions[inter_type]:
                 atoms = inter.atoms
+                atomnames=[block.nodes[atom]['atomname'] for atom in atoms]
                 if itp_mode == "all":
-                    atomnames=[block.nodes[atom]['atomname'] for atom in atoms]
                     group = "_".join(atomnames)
                     inter.meta["comment"] = group
                 else:
@@ -93,7 +101,7 @@ def get_pairwise_interaction(self, molname):
         block = self.blocks[molname]
 
         indices_dict = defaultdict(dict)
-        
+
         for node1, name1 in block.nodes(data='atomname'):
             for node2, name2 in list(block.nodes(data='atomname'))[node1+1:]:
                 atoms = np.array([node1, node2])

fast_forward/mapping.py

@@ -21,7 +21,11 @@ def _selector(atomgroup, indices, names):
     atoms = atomgroup.select_atoms("name " + names_string)
     return atoms
 
-def create_new_universe(universe, mapped_trajectory, mappings):
+def create_new_universe(universe,
+                        mapped_trajectory,
+                        mappings,
+                        n_residues=None,
+                        res_iter=None):
     """
     Create a new universe according to the definitions in mappings
     the old universe, and the mapped trajectory.
@@ -47,15 +51,21 @@ def create_new_universe(universe, mapped_trajectory, mappings):
 
     # create the universe to be returend
     # initalize some attributes
-    n_residues = universe.n_residues
+    if n_residues is None:
+        n_residues = universe.n_residues
+
     n_atoms = mapped_trajectory.shape[1]
     res_seg = np.array([1] * n_residues)
     # to read out these we have to iterate over universe again
     atom_resindex = []
     atomnames = []
     resids = []
     resnames = []
-    for idx, residue in enumerate(universe.res_iter()):
+    # either we loop over pre-defined residues or the universe ones
+    if res_iter is None:
+        res_iter = universe
+    for idx, residue in enumerate(res_iter.res_iter()):
+        print(idx, residue.resname, residue.resid)
         for bead in mappings[residue.resname].beads:
             atomnames.append(bead)
             atom_resindex.append(idx)
@@ -120,18 +130,18 @@ def forward_map_positions(mapped_atoms, bead_idxs, weights, positions, n_frames,
             # we need to first average the non-treated atoms
             # and then take the average of the these atoms with
             # the treated atoms
-            pre_count = 0
-            treat_count = 0
+            pre_count = 0.0
+            treat_count = 0.0
             for kdx in prange(len(atom_idxs)):
                 weight = atom_weights[kdx]
                 atom_idx = atom_idxs[kdx]
                 vector = positions[fdx, atom_idx, :]
                 if atom_idx not in treated_atoms:
                     pre_pos = pre_pos + vector * weight
-                    pre_count += 1
+                    pre_count += 1.0*weight
                 else:
                     treat_pos = treat_pos + vector * weight
-                    treat_count += 1
+                    treat_count += 1.0*weight
 
             if pre_count != 0:
                 pre_pos = pre_pos / pre_count