Source code for mdtraj.geometry.distance

##############################################################################
# MDTraj: A Python Library for Loading, Saving, and Manipulating
#         Molecular Dynamics Trajectories.
# Copyright 2012-2015 Stanford University and the Authors
#
# Authors: Robert McGibbon
# Contributors: Kyle A Beauchamp, Jason Swails
#
# 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 numpy as np

from mdtraj.utils import ensure_type
from mdtraj.utils.unitcell import box_vectors_to_lengths_and_angles

from . import _geometry

__all__ = [
    "compute_distances_core",
    "compute_distances",
    "compute_distances_t",
    "compute_displacements",
    "compute_center_of_mass",
    "compute_center_of_geometry",
    "find_closest_contact",
]


def compute_distances_core(
    positions,
    atom_pairs,
    unitcell_vectors=None,
    periodic=True,
    opt=True,
):
    """Compute the distances between pairs of atoms in each frame.

    Parameters
    ----------
    positions : np.ndarray of shape=(n_frames, n_atoms, 3), dtype=float
        The positions of all atoms for a given trajectory.

    atom_pairs : np.ndarray of shape=(num_pairs, 2), dtype=int
        Each row gives the indices of two atoms involved in the interaction.

    unitcell_vectors : None or np.ndarray of shape(n_frames, 3 x 3), default=None
        A numpy array that specifies the box vectors for all frames for a trajectory.

    periodic : bool, default=True
        If `periodic` is True and the trajectory contains unitcell
        information, we will compute distances under the minimum image
        convention.

    opt : bool, default=True
        Use an optimized native library to calculate distances. Our optimized
        SSE minimum image convention calculation implementation is over 1000x
        faster than the naive numpy implementation.

    Returns
    -------

    distances : np.ndarray, shape=(n_frames, num_pairs), dtype=float
        The distance, in each frame, between each pair of atoms.
    """

    xyz = ensure_type(
        positions,
        dtype=np.float32,
        ndim=3,
        name="traj.xyz",
        shape=(None, None, 3),
        warn_on_cast=False,
    )
    pairs = ensure_type(
        atom_pairs,
        dtype=np.int32,
        ndim=2,
        name="atom_pairs",
        shape=(None, 2),
        warn_on_cast=False,
    )
    if not np.all(np.logical_and(pairs < positions.shape[1], pairs >= 0)):
        raise ValueError(f"atom_pairs must be between 0 and {positions.shape[0]}")

    if len(pairs) == 0:
        return np.zeros((len(xyz), 0), dtype=np.float32)

    if periodic and (unitcell_vectors is not None):
        box = ensure_type(
            unitcell_vectors,
            dtype=np.float32,
            ndim=3,
            name="unitcell_vectors",
            shape=(len(xyz), 3, 3),
            warn_on_cast=False,
        )

        # convert to angles
        unitcell_angles = []
        for fr_unitcell_vectors in unitcell_vectors:
            _, _, _, alpha, beta, gamma = box_vectors_to_lengths_and_angles(
                fr_unitcell_vectors[0],
                fr_unitcell_vectors[1],
                fr_unitcell_vectors[2],
            )
            unitcell_angles.append(np.array([alpha, beta, gamma]))

        orthogonal = np.allclose(np.array(unitcell_angles), 90)

        if opt:
            out = np.empty((xyz.shape[0], pairs.shape[0]), dtype=np.float32)
            _geometry._dist_mic(
                xyz,
                pairs,
                box.transpose(0, 2, 1).copy(),
                out,
                orthogonal,
            )
            return out

        else:
            return _distance_mic(xyz, pairs, box.transpose(0, 2, 1), orthogonal)

    # either there are no unitcell vectors or they dont want to use them
    if opt:
        out = np.empty((xyz.shape[0], pairs.shape[0]), dtype=np.float32)
        _geometry._dist(xyz, pairs, out)
        return out
    else:
        return _distance(xyz, pairs)


[docs] def compute_distances(traj, atom_pairs, periodic=True, opt=True): """Compute the distances between pairs of atoms in each frame. Parameters ---------- traj : Trajectory An mtraj trajectory. atom_pairs : np.ndarray, shape=(num_pairs, 2), dtype=int Each row gives the indices of two atoms involved in the interaction. periodic : bool, default=True If `periodic` is True and the trajectory contains unitcell information, we will compute distances under the minimum image convention. opt : bool, default=True Use an optimized native library to calculate distances. Our optimized SSE minimum image convention calculation implementation is over 1000x faster than the naive numpy implementation. Returns ------- distances : np.ndarray, shape=(n_frames, num_pairs), dtype=float The distance, in each frame, between each pair of atoms. """ return compute_distances_core( traj.xyz, atom_pairs, unitcell_vectors=traj.unitcell_vectors, periodic=periodic, opt=opt, )
def compute_distances_t(traj, atom_pairs, time_pairs, periodic=True, opt=True): """Compute the distances between pairs of atoms at pairs of times. Parameters ---------- traj : Trajectory An mtraj trajectory. atom_pairs : np.ndarray, shape=(num_atom_pairs, 2), dtype=int Each row gives the indices of two atoms involved in the interaction. time_pairs : nd.array, shape=(num_times, 2), dtype=int Each row gives the indices of two frames. periodic : bool, default=True If `periodic` is True and the trajectory contains unitcell information, we will compute distances under the minimum image convention. opt : bool, default=True Use an optimized native library to calculate distances. Our optimized SSE minimum image convention calculation implementation is over 1000x faster than the naive numpy implementation. Returns ------- distances : np.ndarray, shape=(num_times, num_atom_pairs), dtype=float The distance between each pair of atoms at t=t1 and t=t2. """ xyz = ensure_type( traj.xyz, dtype=np.float32, ndim=3, name="traj.xyz", shape=(None, None, 3), warn_on_cast=False, ) pairs = ensure_type( atom_pairs, dtype=np.int32, ndim=2, name="atom_pairs", shape=(None, 2), warn_on_cast=False, ) times = ensure_type( time_pairs, dtype=np.int32, ndim=2, name="time_pairs", shape=(None, 2), warn_on_cast=False, ) if not np.all(np.logical_and(pairs < traj.n_atoms, pairs >= 0)): raise ValueError("atom_pairs must be between 0 and %d" % traj.n_atoms) if not np.all(np.logical_and(times < traj.n_frames, times >= 0)): raise ValueError("time_pairs must be between 0 and %d" % traj.n_frames) if len(pairs) == 0: return np.zeros((len(xyz), 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), warn_on_cast=False, ) orthogonal = np.allclose(traj.unitcell_angles, 90) if opt: out = np.empty((times.shape[0], pairs.shape[0]), dtype=np.float32) _geometry._dist_mic_t( xyz, pairs, times, box.transpose(0, 2, 1).copy(), out, orthogonal, ) out = out.reshape((times.shape[0], pairs.shape[0])) return out else: return _distance_mic_t( xyz, pairs, times, box.transpose(0, 2, 1), orthogonal, ) # either there are no unitcell vectors or they dont want to use them if opt: out = np.empty((times.shape[0], pairs.shape[0]), dtype=np.float32) _geometry._dist_t(xyz, pairs, times, out) return out else: return _distance_t(xyz, pairs, times)
[docs] def compute_displacements(traj, atom_pairs, periodic=True, opt=True): """Compute the displacement vector between pairs of atoms in each frame of a trajectory. Parameters ---------- traj : Trajectory Trajectory to compute distances in atom_pairs : np.ndarray, shape[num_pairs, 2], dtype=int Each row gives the indices of two atoms. periodic : bool, default=True If `periodic` is True and the trajectory contains unitcell information, we will compute distances under the minimum image convention. opt : bool, default=True Use an optimized native library to calculate distances. Our optimized minimum image convention calculation implementation is over 1000x faster than the naive numpy implementation. Returns ------- displacements : np.ndarray, shape=[n_frames, n_pairs, 3], dtype=float32 The displacememt vector, in each frame, between each pair of atoms. """ xyz = ensure_type( traj.xyz, dtype=np.float32, ndim=3, name="traj.xyz", shape=(None, None, 3), warn_on_cast=False, ) pairs = ensure_type( np.asarray(atom_pairs), dtype=np.int32, ndim=2, name="atom_pairs", shape=(None, 2), warn_on_cast=False, ) if not np.all(np.logical_and(pairs < traj.n_atoms, pairs >= 0)): raise ValueError("atom_pairs must be between 0 and %d" % traj.n_atoms) if len(pairs) == 0: # If pairs is an empty slice of an array return np.zeros((len(xyz), 0, 3), 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), warn_on_cast=False, ) orthogonal = np.allclose(traj.unitcell_angles, 90) if opt: out = np.empty((xyz.shape[0], pairs.shape[0], 3), dtype=np.float32) _geometry._dist_mic_displacement( xyz, pairs, box.transpose(0, 2, 1).copy(), out, orthogonal, ) return out else: return _displacement_mic(xyz, pairs, box.transpose(0, 2, 1), orthogonal) # either there are no unitcell vectors or they dont want to use them if opt: out = np.empty((xyz.shape[0], pairs.shape[0], 3), dtype=np.float32) _geometry._dist_displacement(xyz, pairs, out) return out return _displacement(xyz, pairs)
[docs] def compute_center_of_mass(traj, select=None): """Compute the center of mass for each frame. Parameters ---------- traj : Trajectory Trajectory to compute center of mass for select : str, optional, default=all a mdtraj.Topology selection string that defines the set of atoms of which to calculate the center of mass, the default is all atoms Returns ------- com : np.ndarray, shape=(n_frames, 3) Coordinates of the center of mass for each frame """ com = np.empty((traj.n_frames, 3)) if select is None: masses = np.array([a.element.mass for a in traj.top.atoms]) masses /= masses.sum() xyz = traj.xyz else: atoms_of_interest = traj.topology.select(select) masses = np.array([traj.top.atom(i).element.mass for i in atoms_of_interest]) masses /= masses.sum() xyz = traj.xyz[:, atoms_of_interest] for i, x in enumerate(xyz): com[i, :] = x.astype("float64").T.dot(masses) return com
def compute_center_of_geometry(traj): """Compute the center of geometry for each frame. Parameters ---------- traj : Trajectory Trajectory to compute center of geometry for Returns ------- com : np.ndarray, shape=(n_frames, 3) Coordinates of the center of geometry for each frame """ centers = np.zeros((traj.n_frames, 3)) for i, x in enumerate(traj.xyz): centers[i, :] = x.astype("float64").T.mean(axis=1) return centers def find_closest_contact(traj, group1, group2, frame=0, periodic=True): """Find the closest contact between two groups of atoms. Given a frame of a Trajectory and two groups of atoms, identify the pair of atoms (one from each group) that form the closest contact between the two groups. Parameters ---------- traj : Trajectory An mtraj trajectory. group1 : np.ndarray, shape=(num_atoms), dtype=int The indices of atoms in the first group. group2 : np.ndarray, shape=(num_atoms), dtype=int The indices of atoms in the second group. frame : int, default=0 The frame of the Trajectory to take positions from periodic : bool, default=True If `periodic` is True and the trajectory contains unitcell information, we will compute distances under the minimum image convention. Returns ------- result : tuple (int, int, float) The indices of the two atoms forming the closest contact, and the distance between them. """ xyz = ensure_type( traj.xyz, dtype=np.float32, ndim=3, name="traj.xyz", shape=(None, None, 3), warn_on_cast=False, )[frame] atoms1 = ensure_type( group1, dtype=np.int32, ndim=1, name="group1", warn_on_cast=False, ) atoms2 = ensure_type( group2, dtype=np.int32, ndim=1, name="group2", warn_on_cast=False, ) if periodic and traj._have_unitcell: box = ensure_type( traj.unitcell_vectors, dtype=np.float32, ndim=3, name="unitcell_vectors", shape=(len(traj.xyz), 3, 3), warn_on_cast=False, )[frame] else: box = None return _geometry._find_closest_contact(xyz, atoms1, atoms2, box) def _distance(xyz, pairs): "Distance between pairs of points in each frame" delta = np.diff(xyz[:, pairs], axis=2)[:, :, 0] return (delta**2.0).sum(-1) ** 0.5 def _distance_t(xyz, pairs, times): "Distance between pairs of points in specified frames" frame1 = xyz[:, pairs[:, 0]][times[:, 0]] frame2 = xyz[:, pairs[:, 1]][times[:, 1]] out = np.linalg.norm(frame1 - frame2, axis=2) return out def _displacement(xyz, pairs): "Displacement vector between pairs of points in each frame" value = np.diff(xyz[:, pairs], axis=2)[:, :, 0] assert value.shape == ( xyz.shape[0], pairs.shape[0], 3, ), f"v.shape {str(value.shape)}, xyz.shape {str(xyz.shape)}, pairs.shape {str(pairs.shape)}" return value def _reduce_box_vectors(vectors): """Make sure box vectors are in reduced form.""" (bv1, bv2, bv3) = vectors bv3 -= bv2 * round(bv3[1] / bv2[1]) bv3 -= bv1 * round(bv3[0] / bv1[0]) bv2 -= bv1 * round(bv2[0] / bv1[0]) return (bv1, bv2, bv3) def _distance_mic(xyz, pairs, box_vectors, orthogonal): """Distance between pairs of points in each frame under the minimum image convention for periodic boundary conditions. The computation follows scheme B.9 in Tukerman, M. "Statistical Mechanics: Theory and Molecular Simulation", 2010. This is a slow pure python implementation, mostly for testing. """ out = np.empty((xyz.shape[0], pairs.shape[0]), dtype=np.float32) for i in range(len(xyz)): bv1, bv2, bv3 = _reduce_box_vectors(box_vectors[i].T) for j, (a, b) in enumerate(pairs): r12 = xyz[i, b, :] - xyz[i, a, :] r12 -= bv3 * round(r12[2] / bv3[2]) r12 -= bv2 * round(r12[1] / bv2[1]) r12 -= bv1 * round(r12[0] / bv1[0]) dist = np.linalg.norm(r12) if not orthogonal: for ii in range(-1, 2): v1 = bv1 * ii for jj in range(-1, 2): v12 = bv2 * jj + v1 for kk in range(-1, 2): new_r12 = r12 + v12 + bv3 * kk dist = min(dist, np.linalg.norm(new_r12)) out[i, j] = dist return out def _distance_mic_t(xyz, pairs, times, box_vectors, orthogonal): """Distance between pairs of points between specified frames under the minimum image convention for periodic boundary conditions. The computation is modified from scheme B.9 in Tukerman, M. "Statistical Mechanics: Theory and Molecular Simulation", 2010. This is a slow pure python implementation, mostly for testing. """ out = np.empty((times.shape[0], pairs.shape[0]), dtype=np.float32) for i, (a, b) in enumerate(times): bv1, bv2, bv3 = _reduce_box_vectors(box_vectors[a].T) for j, (c, d) in enumerate(pairs): r12 = xyz[a, c] - xyz[b, d] r12 -= bv3 * round(r12[2] / bv3[2]) r12 -= bv2 * round(r12[1] / bv2[1]) r12 -= bv1 * round(r12[0] / bv1[0]) dist = np.linalg.norm(r12) if not orthogonal: for ii in range(-1, 2): v1 = bv1 * ii for jj in range(-1, 2): v12 = bv2 * jj + v1 for kk in range(-1, 2): new_r12 = r12 + v12 + bv3 * kk dist = min(dist, np.linalg.norm(new_r12)) out[i, j] = dist return out def _displacement_mic(xyz, pairs, box_vectors, orthogonal): """Displacement vector between pairs of points in each frame under the minimum image convention for periodic boundary conditions. The computation follows scheme B.9 in Tukerman, M. "Statistical Mechanics: Theory and Molecular Simulation", 2010. This is a very slow pure python implementation, mostly for testing. """ out = np.empty((xyz.shape[0], pairs.shape[0], 3), dtype=np.float32) for i in range(len(xyz)): bv1, bv2, bv3 = _reduce_box_vectors(box_vectors[i].T) # hinv, not used _ = np.linalg.inv(np.array([bv1, bv2, bv3]).T) for j, (a, b) in enumerate(pairs): r12 = xyz[i, b, :] - xyz[i, a, :] r12 -= bv3 * round(r12[2] / bv3[2]) r12 -= bv2 * round(r12[1] / bv2[1]) r12 -= bv1 * round(r12[0] / bv1[0]) min_disp = r12 dist2 = (r12 * r12).sum() if not orthogonal: for ii in range(-1, 2): v1 = bv1 * ii for jj in range(-1, 2): v12 = bv2 * jj + v1 for kk in range(-1, 2): tmp = r12 + v12 + bv3 * kk new_dist2 = (tmp * tmp).sum() if new_dist2 < dist2: dist2 = new_dist2 min_disp = tmp out[i, j] = min_disp return out