Source code for mdtraj.geometry.dihedral

##############################################################################
# MDTraj: A Python Library for Loading, Saving, and Manipulating
#         Molecular Dynamics Trajectories.
# Copyright 2012-2013 Stanford University and the Authors
#
# Authors: Robert McGibbon
# Contributors: Kyle A Beauchamp, Ravi Ramanathan, Lee-Ping Wang
#
# MDTraj is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as
# published by the Free Software Foundation, either version 2.1
# of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with MDTraj. If not, see <http://www.gnu.org/licenses/>.
##############################################################################


import warnings

import numpy as np

from mdtraj.geometry import _geometry, distance
from mdtraj.utils import ensure_type

__all__ = [
    "compute_dihedrals",
    "compute_phi",
    "compute_psi",
    "compute_omega",
    "compute_chi1",
    "compute_chi2",
    "compute_chi3",
    "compute_chi4",
    "compute_chi5",
    "indices_phi",
    "indices_psi",
    "indices_omega",
    "indices_chi1",
    "indices_chi2",
    "indices_chi3",
    "indices_chi4",
    "indices_chi5",
]


def _dihedral(traj, indices, periodic, out=None):
    """Compute the dihedral angles of traj for the atom indices in indices.

    Parameters
    ----------
    xyz : np.ndarray, shape=(num_frames, num_atoms, 3), dtype=float
        The XYZ coordinates of a trajectory
    indices : np.ndarray, shape=(num_dihedrals, 4), dtype=int
        Atom indices to compute dihedrals.
    periodic : bool, default=True
        If `periodic` is True and the trajectory contains unitcell
        information, we will treat dihedrals that cross periodic images
        using the minimum image convention.

    Returns
    -------
    dih : np.ndarray, shape=(num_dihedrals), dtype=float
        dih[i,j] gives the dihedral angle at traj[i] correponding to indices[j].

    """
    ix10 = indices[:, [0, 1]]
    ix21 = indices[:, [1, 2]]
    ix32 = indices[:, [2, 3]]

    b1 = distance.compute_displacements(traj, ix10, periodic=periodic, opt=False)
    b2 = distance.compute_displacements(traj, ix21, periodic=periodic, opt=False)
    b3 = distance.compute_displacements(traj, ix32, periodic=periodic, opt=False)

    c1 = np.cross(b2, b3)
    c2 = np.cross(b1, b2)

    p1 = (b1 * c1).sum(-1)
    p1 *= (b2 * b2).sum(-1) ** 0.5
    p2 = (c1 * c2).sum(-1)

    return np.arctan2(p1, p2, out)


[docs] def compute_dihedrals(traj, indices, periodic=True, opt=True): """Compute the dihedral angles between the supplied quartets of atoms in each frame in a trajectory. Parameters ---------- traj : Trajectory An mtraj trajectory. indices : np.ndarray, shape=(n_dihedrals, 4), dtype=int Each row gives the indices of four atoms which together make a dihedral angle. The angle is between the planes spanned by the first three atoms and the last three atoms, a torsion around the bond between the middle two atoms. periodic : bool, default=True If `periodic` is True and the trajectory contains unitcell information, we will treat dihedrals that cross periodic images using the minimum image convention. opt : bool, default=True Use an optimized native library to calculate angles. Returns ------- dihedrals : np.ndarray, shape=(n_frames, n_dihedrals), dtype=float The output array gives, in each frame from the trajectory, each of the `n_dihedrals` torsion angles. The angles are measured in **radians**. """ xyz = ensure_type( traj.xyz, dtype=np.float32, ndim=3, name="traj.xyz", shape=(None, None, 3), warn_on_cast=False, ) quartets = ensure_type( indices, dtype=np.int32, ndim=2, name="indices", shape=(None, 4), warn_on_cast=False, ) if not np.all(np.logical_and(quartets < traj.n_atoms, quartets >= 0)): raise ValueError("indices must be between 0 and %d" % traj.n_atoms) if len(quartets) == 0: return np.zeros((len(xyz), 0), dtype=np.float32) out = np.zeros((xyz.shape[0], quartets.shape[0]), dtype=np.float32) if periodic and traj._have_unitcell: box = ensure_type( traj.unitcell_vectors, dtype=np.float32, ndim=3, name="unitcell_vectors", shape=(len(xyz), 3, 3), ) if opt: orthogonal = np.allclose(traj.unitcell_angles, 90) _geometry._dihedral_mic( xyz, quartets, box.transpose(0, 2, 1).copy(), out, orthogonal, ) return out else: _dihedral(traj, quartets, periodic, out) return out if opt: _geometry._dihedral(xyz, quartets, out) else: _dihedral(traj, quartets, periodic, out) return out
def _construct_atom_dict(topology): """Create dictionary to lookup indices by atom name, residue_id, and chain index. Parameters ---------- topology : Topology The topology to parse Returns ------- atom_dict : dict Tree of nested dictionaries such that `atom_dict[chain_index][residue_index][atom_name] = atom_index` """ atom_dict = {} for chain in topology.chains: residue_dict = {} for residue in chain.residues: local_dict = {} for atom in residue.atoms: local_dict[atom.name] = atom.index residue_dict[residue.index] = local_dict atom_dict[chain.index] = residue_dict return atom_dict 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 dihedral angles (like phi or psi). It looks for the given pattern of atoms in each residue of a given chain. See the example for details. Parameters ---------- 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' Array of integer offsets for each atom. These are used to refer to atoms forward or backward in the chain relative to the current residue Notes ----- In addition to finding dihedral atoms, this function could be used to match *general* sequences of atoms and residue_id offsets. Examples -------- Here we calculate the phi torsion angles by specifying the correct atom names and the residue_id offsets (e.g. forward or backward in chain) for each atom. >>> traj = mdtraj.load("native.pdb") >>> atom_names = ["C" ,"N" , "CA", "C"] >>> residue_offsets = [-1, 0, 0, 0] >>> found_residue_ids, indices = _atom_sequence(traj, atom_names, residue_offsets) """ if residue_offsets is None: residue_offsets = parse_offsets(atom_names) atom_names = _strip_offsets(atom_names) 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 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]): # 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], ): # 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], ) 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=int) return found_residue_ids, atom_indices def parse_offsets(atom_names): """Convert a list of atom+offset strings into lists offsets. Parameters ---------- atom_names : list The names of the atoms to parse for their offsets. Returns ------- offsets : list The offsets of the atoms, giving whether they refer to atoms in the previous residue (-1), current residue (0) or next residue (+1) Notes ----- For example, ["-C", "N", "CA", "C"] will be parsed as [-1, 0, 0, 0] """ offsets = [] for atom in atom_names: if atom[0] == "-": offsets.append(-1) elif atom[0] == "+": offsets.append(+1) else: offsets.append(0) return offsets def _strip_offsets(atom_names): """Convert a list of atom + offset strings into lists of atoms. Parameters ---------- atom_names : list The names of the atoms, whose offset prexifs you want to strip Notes ----- For example, ["-C", "N", "CA", "C"] will be parsed as ["C","N","CA","C"] """ atoms = [] for atom in atom_names: if atom[0] == "-": atoms.append(atom[1:]) elif atom[0] == "+": atoms.append(atom[1:]) else: atoms.append(atom) return atoms PHI_ATOMS = ["-C", "N", "CA", "C"] PSI_ATOMS = ["N", "CA", "C", "+N"] OMEGA_ATOMS = ["CA", "C", "+N", "+CA"] CHI1_ATOMS = [ ["N", "CA", "CB", "CG"], ["N", "CA", "CB", "CG1"], ["N", "CA", "CB", "SG"], ["N", "CA", "CB", "OG"], ["N", "CA", "CB", "OG1"], ] CHI2_ATOMS = [ ["CA", "CB", "CG", "CD"], ["CA", "CB", "CG", "CD1"], ["CA", "CB", "CG1", "CD1"], ["CA", "CB", "CG", "OD1"], ["CA", "CB", "CG", "ND1"], ["CA", "CB", "CG", "SD"], ] CHI3_ATOMS = [ ["CB", "CG", "CD", "NE"], ["CB", "CG", "CD", "CE"], ["CB", "CG", "CD", "OE1"], ["CB", "CG", "SD", "CE"], ] CHI4_ATOMS = [ ["CG", "CD", "NE", "CZ"], ["CG", "CD", "CE", "NZ"], ] CHI5_ATOMS = [["CD", "NE", "CZ", "NH1"]] 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=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 indices_chi5(top): """Calculate indices for chi5 dihedral angles Parameters ---------- top : Topology Topology for which you want dihedral indices Returns ------- indices : np.ndarray, shape=(n_chi5, 4) The indices of the atoms involved in each of ths chi4 angles """ return _indices_chi(top, CHI5_ATOMS)
[docs] def compute_phi(traj, periodic=True, opt=True): """Calculate the phi torsions of a trajectory. Parameters ---------- traj : Trajectory Trajectory for which you want dihedrals. periodic : bool, default=True If `periodic` is True and the trajectory contains unitcell information, we will treat dihedrals that cross periodic images using the minimum image convention. opt : bool, default=True Use an optimized native library to calculate angles. Returns ------- indices : np.ndarray, shape=(n_phi, 4) The indices of the atoms involved in each of the phi dihedral angles angles : np.ndarray, shape=(n_frames, n_phi) The value of the dihedral angle for each of the angles in each of the frames. """ indices = indices_phi(traj.topology) if len(indices) == 0: return indices, np.empty(shape=(len(traj), 0), dtype=np.float32) return indices, compute_dihedrals(traj, indices, periodic=periodic, opt=opt)
[docs] def compute_psi(traj, periodic=True, opt=True): """Calculate the psi torsions of a trajectory. Parameters ---------- traj : Trajectory Trajectory for which you want dihedrals. periodic : bool, default=True If `periodic` is True and the trajectory contains unitcell information, we will treat dihedrals that cross periodic images using the minimum image convention. opt : bool, default=True Use an optimized native library to calculate angles. Returns ------- indices : np.ndarray, shape=(n_psi, 4) The indices of the atoms involved in each of the psi dihedral angles angles : np.ndarray, shape=(n_frames, n_psi) The value of the dihedral angle for each of the angles in each of the frames. """ indices = indices_psi(traj.topology) if len(indices) == 0: return indices, np.empty(shape=(len(traj), 0), dtype=np.float32) return indices, compute_dihedrals(traj, indices, periodic=periodic, opt=opt)
[docs] def compute_chi1(traj, periodic=True, opt=True): """Calculate the chi1 torsions of a trajectory. chi1 is the first side chain torsion angle formed between the 4 atoms over the CA-CB axis. Parameters ---------- traj : Trajectory Trajectory for which you want dihedrals. periodic : bool, default=True If `periodic` is True and the trajectory contains unitcell information, we will treat dihedrals that cross periodic images using the minimum image convention. opt : bool, default=True Use an optimized native library to calculate angles. Returns ------- indices : np.ndarray, shape=(n_chi, 4) The indices of the atoms involved in each of the chi1 dihedral angles angles : np.ndarray, shape=(n_frames, n_chi) The value of the dihedral angle for each of the angles in each of the frames. """ 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
[docs] def compute_chi2(traj, periodic=True, opt=True): """Calculate the chi2 torsions of a trajectory. chi2 is the second side chain torsion angle formed between the corresponding 4 atoms over the CB-CG axis. Parameters ---------- traj : Trajectory Trajectory for which you want dihedrals. periodic : bool, default=True If `periodic` is True and the trajectory contains unitcell information, we will treat dihedrals that cross periodic images using the minimum image convention. opt : bool, default=True Use an optimized native library to calculate angles. Returns ------- indices : np.ndarray, shape=(n_chi, 4) The indices of the atoms involved in each of the chi dihedral angles angles : np.ndarray, shape=(n_frames, n_chi) The value of the dihedral angle for each of the angles in each of the frames. """ 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
[docs] def compute_chi3(traj, periodic=True, opt=True): """Calculate the chi3 torsions of a trajectory. chi3 is the third side chain torsion angle formed between the corresponding 4 atoms over the CG-CD axis (only the residues ARG, GLN, GLU, LYS & MET have these atoms) Parameters ---------- traj : Trajectory Trajectory for which you want dihedrals. periodic : bool, default=True If `periodic` is True and the trajectory contains unitcell information, we will treat dihedrals that cross periodic images using the minimum image convention. opt : bool, default=True Use an optimized native library to calculate angles. Returns ------- indices : np.ndarray, shape=(n_chi, 4) The indices of the atoms involved in each of the chi dihedral angles angles : np.ndarray, shape=(n_frames, n_chi) The value of the dihedral angle for each of the angles in each of the frames. """ 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
[docs] def compute_chi4(traj, periodic=True, opt=True): """Calculate the chi4 torsions of a trajectory. chi4 is the fourth side chain torsion angle formed between the corresponding 4 atoms over the CD-CE or CD-NE axis (only ARG & LYS residues have these atoms) Parameters ---------- traj : Trajectory Trajectory for which you want dihedrals. periodic : bool, default=True If `periodic` is True and the trajectory contains unitcell information, we will treat dihedrals that cross periodic images using the minimum image convention. opt : bool, default=True Use an optimized native library to calculate angles. Returns ------- indices : np.ndarray, shape=(n_chi, 4) The indices of the atoms involved in each of the chi dihedral angles angles : np.ndarray, shape=(n_frames, n_chi) The value of the dihedral angle for each of the angles in each of the frames. """ indices = indices_chi4(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_chi5(traj, periodic=True, opt=True): """Calculate the chi5 torsions of a trajectory. chi5 is the fiths side chain torsion angle formed between the corresponding 4 atoms over the NE-CZ axis (only ARG residues have these atoms) Parameters ---------- traj : Trajectory Trajectory for which you want dihedrals. periodic : bool, default=True If `periodic` is True and the trajectory contains unitcell information, we will treat dihedrals that cross periodic images using the minimum image convention. opt : bool, default=True Use an optimized native library to calculate angles. Returns ------- indices : np.ndarray, shape=(n_chi, 4) The indices of the atoms involved in each of the chi dihedral angles angles : np.ndarray, shape=(n_frames, n_chi) The value of the dihedral angle for each of the angles in each of the frames. """ indices = indices_chi5(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
[docs] def compute_omega(traj, periodic=True, opt=True): """Calculate the omega torsions of a trajectory. Parameters ---------- traj : Trajectory Trajectory for which you want dihedrals. periodic : bool, default=True If `periodic` is True and the trajectory contains unitcell information, we will treat dihedrals that cross periodic images using the minimum image convention. opt : bool, default=True Use an optimized native library to calculate angles. Returns ------- indices : np.ndarray, shape=(n_omega, 4) The indices of the atoms involved in each of the omega dihedral angles angles : np.ndarray, shape=(n_frames, n_omega) The value of the dihedral angle for each of the angles in each of the frames. """ indices = indices_omega(traj.topology) if len(indices) == 0: return indices, np.empty(shape=(len(traj), 0), dtype=np.float32) return indices, compute_dihedrals(traj, indices, periodic=periodic, opt=opt)