Mapping trajectory

htf/simmodel.py

@@ -480,6 +480,7 @@ def wrap_vector(r, box):
 
 @tf.function
 def box_size(box):
+    # Note that this function is compatible with HOOMD box and not the one from GROMACS.
     return box[1, :] - box[0, :]
 
 

htf/test-py/test_utils.py

@@ -1,4 +1,3 @@
-# Copyright (c) 2020 HOOMD-TF Developers
 import shutil
 import tempfile
 import build_examples
@@ -43,6 +42,47 @@ def test_find_molecules(self):
         assert len(mapping) == 9
         assert len(mapping[0]) == 10
 
+    def test_find_molecules_from_topology(self):
+        try:
+            import MDAnalysis as mda
+        except ImportError:
+            self.skipTest(
+                "MDAnalysis not available; skipping test_find_molecules_from_topology")
+        # Loading inputs
+        import os
+        TPR = os.path.join(os.path.dirname(__file__),
+                           'CG_mapping/test_nvt_prod.tpr')
+        TRAJECTORY = os.path.join(os.path.dirname(__file__),
+                                  'CG_mapping/test_traj.trr')
+        selection = "resname PHE"
+        u = mda.Universe(TPR, TRAJECTORY)
+        protein_FF = u.select_atoms("resname PHE and resid 0:1")
+        atoms_in_molecule_list = [protein_FF.names]
+        molecule_mapping_index = hoomd.htf.find_molecules_from_topology(
+            u, atoms_in_molecule_list, selection=selection)
+        Total_number_of_atoms = u.select_atoms(selection).n_atoms
+        assert Total_number_of_atoms - 1 == molecule_mapping_index[-1][-1]
+
+    def test_matrix_mapping(self):
+        try:
+            import MDAnalysis as mda
+        except ImportError:
+            self.skipTest(
+                "MDAnalysis not available; skipping test_matrix_mapping")
+        # Loading inputs
+        import os
+        TPR = os.path.join(os.path.dirname(__file__),
+                           'CG_mapping/test_nvt_prod.tpr')
+        TRAJECTORY = os.path.join(os.path.dirname(__file__),
+                                  'CG_mapping/test_traj.trr')
+        u = mda.Universe(TPR, TRAJECTORY)
+        # Generating Mapping Matrix for Water
+        water = u.select_atoms("resname SOL and resid 500")
+        Beads_distribution = [['OW', 'HW1', 'HW2']]
+        mapping_water = hoomd.htf.matrix_mapping(water, Beads_distribution)
+        np.testing.assert_array_equal(np.round(mapping_water, 9), np.array([
+                                      [0.88809574, 0.05595213, 0.05595213]]))
+
     def test_bad_sparse_mapping(self):
         mapping_matrix = np.array([
             [1, 0, 0],
@@ -145,35 +185,23 @@ def test_com(self):
 
     def test_compute_nlist(self):
         N = 10
-        positions = tf.cast(
-            tf.tile(tf.reshape(tf.range(N), [-1, 1]), [1, 3]), tf.float32)
+        positions = tf.tile(tf.reshape(tf.range(N), [-1, 1]), [1, 3])
         box_size = [100., 100., 100.]
         nlist = hoomd.htf.compute_nlist(
-            positions,
+            tf.cast(
+                positions,
+                tf.float32),
             100.,
             9,
             box_size,
-            return_types=False,
-            sorted=True)
+            True)
         nlist = nlist.numpy()
         # particle 1 is closest to 0
         np.testing.assert_array_almost_equal(nlist[0, 0, :], [1, 1, 1, 1])
         # particle 0 is -9 away from 9
         np.testing.assert_array_almost_equal(nlist[-1, -1, :],
                                              [-9, -9, -9, 0])
 
-        extended_positions = tf.concat([positions, tf.zeros((N, 1))], axis=1)
-        nlist = hoomd.htf.compute_nlist(
-            extended_positions,
-            100.,
-            9,
-            box_size,
-            return_types=True,
-            sorted=True)
-        nlist = nlist.numpy()
-        # particle 1 is closest to type 0
-        np.testing.assert_array_almost_equal(nlist[0, 0, :], [1, 1, 1, 0])
-
     def test_compute_nlist_cut(self):
         N = 10
         positions = tf.tile(tf.reshape(tf.range(N), [-1, 1]), [1, 3])
@@ -185,7 +213,7 @@ def test_compute_nlist_cut(self):
             5.5,
             9,
             box_size,
-            sorted=True)
+            True)
         nlist = nlist.numpy()
         # particle 1 is closest to 0
         np.testing.assert_array_almost_equal(nlist[0, 0, :], [1, 1, 1, 1])

htf/utils.py

@@ -89,6 +89,49 @@ def find_molecules(system):
     return mapping
 
 
+def find_molecules_from_topology(universe, atoms_in_molecule_list, selection='all'):
+    R""" Given a universe from MDAnaylis and list of atoms in every molecule type in the system, return a mapping from molecule index to particle index.
+        Depnding on the size of your system, this fuction might be slow to run.
+
+    :param universe: Use MDAnalysis universe to read the tpr topology file from GROMACS.
+    :type universe: MDAnalysis Universe object
+    :param selection: The atom groups to extract from universe
+    :param atoms_in_molecule_list: This is a list of atoms lists in every molecule type in the system.
+
+    :return: A list of length L (number of molecules) whose elements are lists of atom indices.
+
+    Here's an example:
+
+        .. code:: python
+                TPR = 'nvt_prod.tpr'
+                TRAJECTORY = 'Molecules_CG_Mapping/traj.trr'
+                u = mda.Universe(TPR, TRAJECTORY)
+                atoms_in_molecule_list = [u.select_atoms("resname PHE and resid 0:1").names]
+                find_molecules_from_topology(u, atoms_in_molecule_list, selection = "resname PHE")
+    """
+
+    # Getting total number of atoms in selection from topology
+    Total_number_of_atoms = universe.select_atoms(selection).n_atoms
+
+    # Getting the list of every redidue type name in topology
+    _, idx = np.unique(universe.atoms.resnames, return_index=True)
+    resname_list = universe.atoms.resnames[np.sort(idx)].tolist()
+
+    molecule_list_indexed = []
+    molecule_to_be_added = []
+    for i in range(Total_number_of_atoms):
+        resname_type_index = resname_list.index(universe.atoms.resnames[i])
+        molecule_length = len(atoms_in_molecule_list[resname_type_index])
+        if len(molecule_to_be_added) < molecule_length:
+            molecule_to_be_added.append(i)
+        if len(molecule_to_be_added) == molecule_length:
+            molecule_list_indexed.append(molecule_to_be_added)
+            molecule_to_be_added = []
+    if molecule_list_indexed[-1][-1] != Total_number_of_atoms - 1:
+        raise Exception(
+            "Mismatch found between the number of atoms in the system and the final index value. Check your atoms_in_molecule_list input.")
+    return molecule_list_indexed
+
 def matrix_mapping(molecule, beads_distribution):
     R''' This will create a M x N mass weighted mapping matrix where M is the number
         of atoms in the molecule and N is the number of mapping beads.
@@ -114,6 +157,7 @@ def matrix_mapping(molecule, beads_distribution):
     return CG_matrix
 
 
+
 def sparse_mapping(molecule_mapping, molecule_mapping_index,
                    system=None):
     ''' This will create the necessary indices and values for
@@ -240,6 +284,9 @@ def iter_from_trajectory(nneighbor_cutoff, universe, selection='all', r_cut=10.,
     # box_size = [box[0], box[1], box[2]]
     nlist = compute_nlist(atom_group.positions, r_cut=r_cut,
                           NN=nneighbor_cutoff, box_size=[box[0], box[1], box[2]])
+    # Modifying the universe for non 'all' atom selections.
+    if selection != 'all':
+        universe = mda.Merge(universe.select_atoms(selection))
     # Run the model at every nth frame, where n = period
     for i, ts in enumerate(universe.trajectory):
         if i % period == 0: