Merge pull request mdtraj#972 from mpharrigan/dihinds

Compute dihedral indices separately

mdtraj/geometry/dihedral.py

@@ -32,7 +32,9 @@
 import warnings
 
 __all__ = ['compute_dihedrals', 'compute_phi', 'compute_psi', 'compute_omega',
-           'compute_chi1','compute_chi2','compute_chi3','compute_chi4']
+           'compute_chi1', 'compute_chi2', 'compute_chi3', 'compute_chi4',
+           'indices_phi', 'indices_psi', 'indices_omega',
+           'indices_chi1', 'indices_chi2', 'indices_chi3', 'indices_chi4']
 
 ##############################################################################
 # Functions
@@ -164,7 +166,7 @@ def _construct_atom_dict(topology):
     return atom_dict
 
 
-def _atom_sequence(traj, atom_names, residue_offsets=None):
+def _atom_sequence(top, atom_names, residue_offsets=None):
     """Find sequences of atom indices corresponding to desired atoms.
 
     This method can be used to find sets of atoms corresponding to specific
@@ -173,8 +175,10 @@ def _atom_sequence(traj, atom_names, residue_offsets=None):
 
     Parameters
     ----------
-    traj : Trajectory
-        Trajectory for which you want dihedrals.
+    top : Topology or Trajectory (deprecated)
+        Topology for which you want dihedrals. If you pass an object
+        with a "topology" attribute, we will use that topology, but this
+        behavior is deprecated.
     atom_names : np.ndarray, shape=(4), dtype='str'
         Array of atoms to in each dihedral angle.
     residue_offsets : np.ndarray, optional, shape=(4), dtype='int'
@@ -184,7 +188,7 @@ def _atom_sequence(traj, atom_names, residue_offsets=None):
 
     Notes
     -----
-    In additional finding dihedral atoms, this function could be used to
+    In addition to finding dihedral atoms, this function could be used to
     match *general* sequences of atoms and residue_id offsets.
 
     Examples
@@ -202,28 +206,39 @@ def _atom_sequence(traj, atom_names, residue_offsets=None):
         residue_offsets = parse_offsets(atom_names)
     atom_names = _strip_offsets(atom_names)
 
-    atom_dict = _construct_atom_dict(traj.top)
+    if hasattr(top, "topology"):
+        warnings.warn("Passing a Trajectory object to _atom_sequence is"
+                      "deprecated. Please pass a Topology object",
+                      DeprecationWarning)
+        top = top.topology
+    atom_dict = _construct_atom_dict(top)
 
     atom_indices = []
     found_residue_ids = []
     # py3k criticial list(zip(, not just zip(, since we iterate multiple
     # times through it
     atoms_and_offsets = list(zip(atom_names, residue_offsets))
-    for chain in traj.top.chains:
+    for chain in top.chains:
         cid = chain.index
         for residue in chain.residues:
             rid = residue.index
             # Check that desired residue_IDs are in dict
-            if all([rid + offset in atom_dict[cid] for offset in residue_offsets]):
+            if all([rid + offset in atom_dict[cid]
+                    for offset in residue_offsets]):
                 # Check that we find all atom names in dict
-                if all([atom in atom_dict[cid][rid + offset] for atom, offset in atoms_and_offsets]):
+                if all([atom in atom_dict[cid][rid + offset]
+                        for atom, offset in atoms_and_offsets]):
                     # Lookup desired atom indices and and add to list.
-                    atom_indices.append([atom_dict[cid][rid + offset][atom] for atom, offset in atoms_and_offsets])
+                    atom_indices.append([atom_dict[cid][rid + offset][atom]
+                                         for atom, offset in atoms_and_offsets])
                     found_residue_ids.append(rid)
 
     atom_indices = np.array(atom_indices)
     found_residue_ids = np.array(found_residue_ids)
 
+    if len(atom_indices) == 0:
+        atom_indices = np.empty(shape=(0, 4), dtype=np.int)
+
     return found_residue_ids, atom_indices
 
 
@@ -281,6 +296,7 @@ def _strip_offsets(atom_names):
             atoms.append(atom)
     return atoms
 
+
 PHI_ATOMS = ["-C", "N", "CA", "C"]
 PSI_ATOMS = ["N", "CA", "C", "+N"]
 OMEGA_ATOMS = ["CA", "C", "+N", "+CA"]
@@ -306,9 +322,126 @@ def _strip_offsets(atom_names):
 CHI4_ATOMS = [["CG", "CD", "NE", "CZ"],
               ["CG", "CD", "CE", "NZ"]]
 
-_get_indices_omega = lambda traj: _atom_sequence(traj, OMEGA_ATOMS)
-_get_indices_phi = lambda traj: _atom_sequence(traj, PHI_ATOMS)
-_get_indices_psi = lambda traj: _atom_sequence(traj, PSI_ATOMS)
+
+def indices_phi(top):
+    """Calculate indices for phi dihedral angles
+
+    Parameters
+    ----------
+    top : Topology
+        Topology for which you want dihedral indices
+
+    Returns
+    -------
+    indices : np.ndarray, shape=(n_phi, 4)
+        The indices of the atoms involved in each of ths phi angles
+    """
+    return _atom_sequence(top, PHI_ATOMS)[1]
+
+
+def indices_psi(top):
+    """Calculate indices for psi dihedral angles
+
+    Parameters
+    ----------
+    top : Topology
+        Topology for which you want dihedral indices
+
+    Returns
+    -------
+    indices : np.ndarray, shape=(n_psi, 4)
+        The indices of the atoms involved in each of ths psi angles
+    """
+    return _atom_sequence(top, PSI_ATOMS)[1]
+
+
+def indices_omega(top):
+    """Calculate indices for omega dihedral angles
+
+    Parameters
+    ----------
+    top : Topology
+        Topology for which you want dihedral indices
+
+    Returns
+    -------
+    indices : np.ndarray, shape=(n_omega, 4)
+        The indices of the atoms involved in each of ths omega angles
+    """
+    return _atom_sequence(top, OMEGA_ATOMS)[1]
+
+
+def _indices_chi(top, chi_atoms):
+    rids, indices = zip(*(_atom_sequence(top, atoms) for atoms in chi_atoms))
+    id_sort = np.argsort(np.concatenate(rids))
+    if not any(x.size for x in indices):
+        return np.empty(shape=(0, 4), dtype=np.int)
+    indices = np.vstack(x for x in indices if x.size)[id_sort]
+    return indices
+
+
+def indices_chi1(top):
+    """Calculate indices for chi1 dihedral angles
+
+    Parameters
+    ----------
+    top : Topology
+        Topology for which you want dihedral indices
+
+    Returns
+    -------
+    indices : np.ndarray, shape=(n_chi1, 4)
+        The indices of the atoms involved in each of ths chi1 angles
+    """
+    return _indices_chi(top, CHI1_ATOMS)
+
+
+def indices_chi2(top):
+    """Calculate indices for chi2 dihedral angles
+
+    Parameters
+    ----------
+    top : Topology
+        Topology for which you want dihedral indices
+
+    Returns
+    -------
+    indices : np.ndarray, shape=(n_chi2, 4)
+        The indices of the atoms involved in each of ths chi2 angles
+    """
+    return _indices_chi(top, CHI2_ATOMS)
+
+
+def indices_chi3(top):
+    """Calculate indices for chi3 dihedral angles
+
+    Parameters
+    ----------
+    top : Topology
+        Topology for which you want dihedral indices
+
+    Returns
+    -------
+    indices : np.ndarray, shape=(n_chi3, 4)
+        The indices of the atoms involved in each of ths chi3 angles
+    """
+    return _indices_chi(top, CHI3_ATOMS)
+
+
+def indices_chi4(top):
+    """Calculate indices for chi4 dihedral angles
+
+    Parameters
+    ----------
+    top : Topology
+        Topology for which you want dihedral indices
+
+    Returns
+    -------
+    indices : np.ndarray, shape=(n_chi4, 4)
+        The indices of the atoms involved in each of ths chi4 angles
+    """
+    return _indices_chi(top, CHI4_ATOMS)
 
 
 def compute_phi(traj, periodic=True, opt=True):
@@ -333,9 +466,9 @@ def compute_phi(traj, periodic=True, opt=True):
         The value of the dihedral angle for each of the angles in each of
         the frames.
     """
-    rid, indices = _get_indices_phi(traj)
+    indices = indices_phi(traj.topology)
     if len(indices) == 0:
-        return np.empty(shape=(0, 4), dtype=np.int), np.empty(shape=(len(traj), 0), dtype=np.float32)
+        return indices, np.empty(shape=(len(traj), 0), dtype=np.float32)
     return indices, compute_dihedrals(traj, indices, periodic=periodic, opt=opt)
 
 
@@ -361,9 +494,9 @@ def compute_psi(traj, periodic=True, opt=True):
         The value of the dihedral angle for each of the angles in each of
         the frames.
     """
-    rid, indices = _get_indices_psi(traj)
+    indices = indices_psi(traj.topology)
     if len(indices) == 0:
-        return np.empty(shape=(0, 4), dtype=np.int), np.empty(shape=(len(traj), 0), dtype=np.float32)
+        return indices, np.empty(shape=(len(traj), 0), dtype=np.float32)
     return indices, compute_dihedrals(traj, indices, periodic=periodic, opt=opt)
 
 
@@ -390,14 +523,11 @@ def compute_chi1(traj, periodic=True, opt=True):
         The value of the dihedral angle for each of the angles in each of
         the frames.
     """
-    rids, indices = zip(*(_atom_sequence(traj, atoms) for atoms in CHI1_ATOMS))
-    id_sort = np.argsort(np.concatenate(rids))
-    if not any(x.size for x in indices):
-        return np.empty(shape=(0, 4), dtype=np.int), np.empty(shape=(len(traj), 0), dtype=np.float32)
-
-    indices = np.vstack(x for x in indices if x.size)[id_sort]
-    all_chi1 = compute_dihedrals(traj, indices, periodic=periodic, opt=opt)
-    return indices, all_chi1
+    indices = indices_chi1(traj.topology)
+    if len(indices) == 0:
+        return indices, np.empty(shape=(len(traj), 0), dtype=np.float32)
+    all_chi = compute_dihedrals(traj, indices, periodic=periodic, opt=opt)
+    return indices, all_chi
 
 
 def compute_chi2(traj, periodic=True, opt=True):
@@ -423,14 +553,12 @@ def compute_chi2(traj, periodic=True, opt=True):
         The value of the dihedral angle for each of the angles in each of
         the frames.
     """
-    rids, indices = zip(*(_atom_sequence(traj, atoms) for atoms in CHI2_ATOMS))
-    id_sort = np.argsort(np.concatenate(rids))
-    if not any(x.size for x in indices):
-        return np.empty(shape=(0, 4), dtype=np.int), np.empty(shape=(len(traj), 0), dtype=np.float32)
 
-    indices = np.vstack(x for x in indices if x.size)[id_sort]
-    all_chi1 = compute_dihedrals(traj, indices, periodic=periodic, opt=opt)
-    return indices, all_chi1
+    indices = indices_chi2(traj.topology)
+    if len(indices) == 0:
+        return indices, np.empty(shape=(len(traj), 0), dtype=np.float32)
+    all_chi = compute_dihedrals(traj, indices, periodic=periodic, opt=opt)
+    return indices, all_chi
 
 
 def compute_chi3(traj, periodic=True, opt=True):
@@ -457,14 +585,11 @@ def compute_chi3(traj, periodic=True, opt=True):
         The value of the dihedral angle for each of the angles in each of
         the frames.
     """
-    rids, indices = zip(*(_atom_sequence(traj, atoms) for atoms in CHI3_ATOMS))
-    id_sort = np.argsort(np.concatenate(rids))
-    if not any(x.size for x in indices):
-        return np.empty(shape=(0, 4), dtype=np.int), np.empty(shape=(len(traj), 0), dtype=np.float32)
-
-    indices = np.vstack(x for x in indices if x.size)[id_sort]
-    all_chi1 = compute_dihedrals(traj, indices, periodic=periodic, opt=opt)
-    return indices, all_chi1
+    indices = indices_chi3(traj.topology)
+    if len(indices) == 0:
+        return indices, np.empty(shape=(len(traj), 0), dtype=np.float32)
+    all_chi = compute_dihedrals(traj, indices, periodic=periodic, opt=opt)
+    return indices, all_chi
 
 
 def compute_chi4(traj, periodic=True, opt=True):
@@ -491,14 +616,11 @@ def compute_chi4(traj, periodic=True, opt=True):
         The value of the dihedral angle for each of the angles in each of
         the frames.
     """
-    rids, indices = zip(*(_atom_sequence(traj, atoms) for atoms in CHI4_ATOMS))
-    id_sort = np.argsort(np.concatenate(rids))
-    if not any(x.size for x in indices):
-        return np.empty(shape=(0, 4), dtype=np.int), np.empty(shape=(len(traj), 0), dtype=np.float32)
-
-    indices = np.vstack(x for x in indices if x.size)[id_sort]
-    all_chi1 = compute_dihedrals(traj, indices, periodic=periodic, opt=opt)
-    return indices, all_chi1
+    indices = indices_chi3(traj.topology)
+    if len(indices) == 0:
+        return indices, np.empty(shape=(len(traj), 0), dtype=np.float32)
+    all_chi = compute_dihedrals(traj, indices, periodic=periodic, opt=opt)
+    return indices, all_chi
 
 
 def compute_omega(traj, periodic=True, opt=True):
@@ -523,7 +645,7 @@ def compute_omega(traj, periodic=True, opt=True):
         The value of the dihedral angle for each of the angles in each of
         the frames.
     """
-    rid, indices = _get_indices_omega(traj)
+    indices = indices_omega(traj.topology)
     if len(indices) == 0:
-        return np.empty(shape=(0,4), dtype=np.int), np.empty(shape=(len(traj), 0), dtype=np.float32)
+        return indices, np.empty(shape=(len(traj), 0), dtype=np.float32)
     return indices, compute_dihedrals(traj, indices, periodic=periodic, opt=opt)

mdtraj/geometry/tests/test_geometry.py

@@ -106,27 +106,24 @@ def test_dihedral_indices():
     phi0_ind = np.array([3, 12, 13, 14]) - 1  # Taken from PDB, so subtract 1
     psi0_ind = np.array([1, 2,  3, 12]) - 1  # Taken from PDB, so subtract 1
 
-    rid, ind = mdtraj.geometry.dihedral._get_indices_phi(traj)
+    ind = mdtraj.geometry.dihedral.indices_phi(traj)
     eq(ind[0], phi0_ind)
-    eq(int(rid[0]), 1)
 
-    rid, ind = mdtraj.geometry.dihedral._get_indices_psi(traj)
+    ind = mdtraj.geometry.dihedral.indices_psi(traj)
     eq(ind[0], psi0_ind)
-    eq(int(rid[0]), 0)
 
 
 def test_dihedral_index_offset_generation():
     traj = md.load(get_fn('1bpi.pdb'))
+    top = traj.topology
 
-    result = np.array([2, 11, 12, 13])  # The atom indices of the first phi angle
+    # The atom indices of the first phi angle
+    result = np.array([2, 11, 12, 13])
     print([e.name for e in traj.topology.atoms])
-    rid1, ind1 = mdtraj.geometry.dihedral._get_indices_phi(traj)
-    rid2, ind2 = mdtraj.geometry.dihedral._atom_sequence(traj, ["-C","N","CA","C"])
-    rid3, ind3 = mdtraj.geometry.dihedral._atom_sequence(traj, ["-C","N","CA","C"], [-1, 0, 0, 0])
-    rid4, ind4 = mdtraj.geometry.dihedral._atom_sequence(traj, ["C" ,"N","CA","C"], [-1, 0, 0, 0])
-    eq(rid1, rid2)
-    eq(rid1, rid3)
-    eq(rid1, rid4)
+    ind1 = mdtraj.geometry.dihedral.indices_phi(top)
+    rid2, ind2 = mdtraj.geometry.dihedral._atom_sequence(top, ["-C","N","CA","C"])
+    rid3, ind3 = mdtraj.geometry.dihedral._atom_sequence(top, ["-C","N","CA","C"], [-1, 0, 0, 0])
+    rid4, ind4 = mdtraj.geometry.dihedral._atom_sequence(top, ["C" ,"N","CA","C"], [-1, 0, 0, 0])
 
     eq(ind1, ind2)
     eq(ind1, ind3)