Nojump

package/AUTHORS

@@ -212,6 +212,7 @@ Chronological list of authors
   - Domenico Marson
   - Ahmed Salah Ghoneim
   - Alexander Schlaich
+  - Josh Vermaas
 
 External code
 -------------

package/CHANGELOG

@@ -37,6 +37,8 @@ Fixes
     (Issue #3336)
 
 Enhancements
+  * Add a nojump transformation, which unwraps trajectories so that particle
+    paths are continuous. (Issue #3703, PR #4031)
   * Added AtomGroup TopologyAttr to calculate gyration moments (Issue #3904,
     PR #3905)
   * Add support for TPR files produced by Gromacs 2023 (Issue #4047)

package/MDAnalysis/transformations/__init__.py

@@ -129,6 +129,7 @@ def wrapped(ts):
 from .translate import translate, center_in_box
 from .rotate import rotateby
 from .positionaveraging import PositionAverager
+from .nojump import NoJump
 from .fit import fit_translation, fit_rot_trans
 from .wrap import wrap, unwrap
 from .boxdimensions import set_dimensions

package/MDAnalysis/transformations/nojump.py

@@ -0,0 +1,168 @@
+# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding:utf-8 -*-
+# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4
+#
+# MDAnalysis --- https://www.mdanalysis.org
+# Copyright (c) 2006-2017 The MDAnalysis Development Team and contributors
+# (see the file AUTHORS for the full list of names)
+#
+# Released under the GNU Public Licence, v2 or any higher version
+#
+# Please cite your use of MDAnalysis in published work:
+#
+# R. J. Gowers, M. Linke, J. Barnoud, T. J. E. Reddy, M. N. Melo, S. L. Seyler,
+# D. L. Dotson, J. Domanski, S. Buchoux, I. M. Kenney, and O. Beckstein.
+# MDAnalysis: A Python package for the rapid analysis of molecular dynamics
+# simulations. In S. Benthall and S. Rostrup editors, Proceedings of the 15th
+# Python in Science Conference, pages 102-109, Austin, TX, 2016. SciPy.
+# doi: 10.25080/majora-629e541a-00e
+#
+# N. Michaud-Agrawal, E. J. Denning, T. B. Woolf, and O. Beckstein.
+# MDAnalysis: A Toolkit for the Analysis of Molecular Dynamics Simulations.
+# J. Comput. Chem. 32 (2011), 2319--2327, doi:10.1002/jcc.21787
+#
+
+"""\
+No Jump Trajectory Unwrapping --- :mod:`MDAnalysis.transformations.nojump`
+=======================================================================================
+
+Unwraps the trajectory such that atoms never move more than half a periodic box length.
+The consequence of this is that particles will diffuse across periodic boundaries when
+needed. This unwrapping method is suitable as a preprocessing step to calculate
+molecular diffusion, or more commonly to keep multi-domain proteins whole during trajectory
+analysis.
+The algorithm used is based on :cite:p:`Kulke2022`.
+
+.. autoclass:: NoJump
+
+"""
+import numpy as np
+import warnings
+
+from .base import TransformationBase
+from ..due import due, Doi
+from ..exceptions import NoDataError
+
+
+class NoJump(TransformationBase):
+    """
+    Returns transformed coordinates for the given timestep so that an atom
+    does not move more than half the periodic box size between two timesteps, and will move
+    across periodic boundary edges. The algorithm used is based on :cite:p:`Kulke2022`,
+    equation B6 for non-orthogonal systems, so it is general to most applications where
+    molecule trajectories should not "jump" from one side of a periodic box to another.
+
+    Note that this transformation depends on a periodic box dimension being set for every
+    frame in the trajectory, and that this box dimension can be transformed to an orthonormal
+    unit cell. If not, an error is emitted. Since it is typical to transform all frames
+    sequentially when unwrapping trajectories, warnings are emitted if non-sequential timesteps
+    are transformed using NoJump.
+
+    Example
+    -------
+
+    Suppose you had a molecular trajectory with a protein dimer that moved across a periodic
+    boundary. This will normally appear to make the dimer split apart. This transformation
+    uses the molecular motions between frames in a wrapped trajectory to create an unwrapped
+    trajectory where molecules never move more than half a periodic box dimension between subsequent
+    frames. Thus, the normal use case is to apply the transformation to every frame of a trajectory,
+    and to do so sequentially.
+
+    .. code-block:: python
+
+        transformation = NoJump()
+        u.trajectory.add_transformations(transformation)
+        for ts in u.trajectory:
+            print(ts.positions)
+
+    In this case, ``ts.positions`` will return the NoJump unwrapped trajectory, which would keep
+    protein dimers together if they are together in the inital timestep. The unwrapped trajectory may
+    also be useful to compute diffusion coefficients via the Einstein relation.
+    To reverse the process, wrap the coordinates again.
+
+    Returns
+    -------
+    MDAnalysis.coordinates.timestep.Timestep
+
+    References
+    ----------
+    .. bibliography::
+        :filter: False
+        :style: MDA
+
+        Kulke2022
+
+    """
+
+    @due.dcite(
+        Doi("10.1021/acs.jctc.2c00327"),
+        description="Works through the orthogonal case for unwrapping, "
+        "and proposes the non-orthogonal approach.",
+        path=__name__,
+    )
+    def __init__(
+        self,
+        check_continuity=True,
+        max_threads=None,
+        # NoJump transforms are inherently unparallelizable, since
+        # it depends on the previous frame's unwrapped coordinates
+        parallelizable=False,
+    ):
+        super().__init__(max_threads=max_threads, parallelizable=parallelizable)
+        self.prev = None
+        self.old_frame = 0
+        self.older_frame = "A"
+        self.check_c = check_continuity
+
+    def _transform(self, ts):
+        L = ts.triclinic_dimensions
+        if L is None:
+            msg = f"Periodic box dimensions not provided at step {ts.frame}"
+            raise NoDataError(msg)
+        try:
+            Linverse = np.linalg.inv(L)
+        except np.linalg.LinAlgError:
+            msg = f"Periodic box dimensions are not invertible at step {ts.frame}"
+            raise NoDataError(msg)
+        if self.prev is None:
+            self.prev = ts.positions @ Linverse
+            self.old_frame = ts.frame
+            return ts
+        if (
+            self.check_c
+            and self.older_frame != "A"
+            and (self.old_frame - self.older_frame) != (ts.frame - self.old_frame)
+        ):
+            warnings.warn(
+                "NoJump detected that the interval between frames is unequal."
+                "We are resetting the reference frame to the current timestep.",
+                UserWarning,
+            )
+            self.prev = ts.positions @ Linverse
+            self.old_frame = ts.frame
+            self.older_frame = "A"
+            return ts
+        if self.check_c and np.abs(self.old_frame - ts.frame) != 1:
+            warnings.warn(
+                "NoJump transform is only accurate when positions"
+                "do not move by more than half a box length."
+                "Currently jumping between frames with a step of more than 1."
+                "This might be fine, but depending on the trajectory stride,"
+                "this might be inaccurate.",
+                UserWarning,
+            )
+        # Convert into reduced coordinate space
+        fcurrent = ts.positions @ Linverse
+        fprev = self.prev  # Previous unwrapped coordinates in reduced box coordinates.
+        # Calculate the new positions in reduced coordinate space (Equation B6 from
+        # 10.1021/acs.jctc.2c00327). As it turns out, the displacement term can
+        # be moved inside the round function in this coordinate space, as the
+        # difference between wrapped and unwrapped coordinates is an integer.
+        newpositions = fcurrent - np.round(fcurrent - fprev)
+        # Convert back into real space
+        ts.positions = newpositions @ L
+        # Set things we need to save for the next frame.
+        self.prev = newpositions  # Note that this is in reduced coordinate space.
+        self.older_frame = self.old_frame
+        self.old_frame = ts.frame
+
+        return ts

package/doc/sphinx/source/documentation_pages/trajectory_transformations.rst

@@ -270,5 +270,6 @@ Currently implemented transformations
    ./transformations/positionaveraging
    ./transformations/fit
    ./transformations/wrap
+   ./transformations/nojump
    ./transformations/boxdimensions
 

package/doc/sphinx/source/documentation_pages/transformations/nojump.rst

@@ -0,0 +1 @@
+.. automodule:: MDAnalysis.transformations.nojump

package/doc/sphinx/source/references.bib

@@ -759,3 +759,14 @@ @incollection{Waveren2005
     booktitle = {},
     id = {transformations}
 }
+
+@article{Kulke2022,
+  title = {Reversible Unwrapping Algorithm for Constant-Pressure Molecular Dynamics Simulations},
+  author = {Kulke, Martin and Vermaas, Josh V.},
+  year = {2022},
+  journal = {Journal of Chemical Theory and Computation}, 
+  volume = {18}, 
+  number = {10},
+  pages = {6161--6171},
+  doi = {10.1021/acs.jctc.2c00327}
+}

testsuite/MDAnalysisTests/transformations/test_nojump.py

@@ -0,0 +1,181 @@
+import numpy as np
+import pytest
+from numpy.testing import assert_allclose
+
+import MDAnalysis as mda
+from MDAnalysis.transformations import NoJump, wrap
+from MDAnalysisTests import datafiles as data
+
+
+@pytest.fixture()
+def nojump_universes_fromfile():
+    '''
+    Create the universe objects for the tests.
+    '''
+    u = mda.Universe(data.PSF_TRICLINIC, data.DCD_TRICLINIC)
+    transformation = NoJump()
+    u.trajectory.add_transformations(transformation)
+    return u
+
+
+@pytest.fixture()
+def nojump_universe():
+    """
+    Create the universe objects for the tests.
+    """
+    u = mda.Universe.empty(1, trajectory=True)
+    coordinates = np.empty((100, u.atoms.n_atoms, 3))  # number of frames
+    coordinates[::3, 0] = 0 * np.ones(3) / 3
+    coordinates[1::3, 0] = 1 * np.ones(3) / 3
+    coordinates[2::3, 0] = 2 * np.ones(3) / 3
+    u.load_new(coordinates, order="fac")
+    return u
+
+
+@pytest.fixture()
+def nojump_constantvel_universe():
+    """
+    Create the universe objects for the tests.
+    """
+    Natom = 1
+    Nframe = 100
+    coordinates = np.empty((Nframe, Natom, 3))  # number of frames
+    coordinates[:, 0, 0] = np.linspace(0, 45, Nframe)
+    coordinates[:, 0, 1] = np.linspace(0, 15, Nframe)
+    coordinates[:, 0, 2] = np.linspace(0, 10, Nframe)
+    reference = mda.Universe.empty(Natom, trajectory=True)
+    reference.load_new(coordinates, order="fac")
+    return reference
+
+
+def test_nojump_orthogonal_fwd(nojump_universe):
+    """
+    Test if the nojump transform is returning the correct
+    values when iterating forwards over the sample trajectory.
+    """
+    u = nojump_universe
+    dim = np.asarray([1, 1, 1, 90, 90, 90], np.float32)
+    workflow = [mda.transformations.boxdimensions.set_dimensions(dim), NoJump()]
+    u.trajectory.add_transformations(*workflow)
+    transformed_coordinates = u.trajectory.timeseries()[0]
+    # Step is 1 unit every 3 steps. After 99 steps from the origin,
+    # we'll end up at 33.
+    assert_allclose(
+        transformed_coordinates, np.outer(np.linspace(0, 33, 100), np.ones(3))
+    )
+
+
+def test_nojump_nonorthogonal_fwd(nojump_universe):
+    """
+    Test if the nojump transform is returning the correct
+    values when iterating forwards over the sample trajectory.
+    """
+    u = nojump_universe
+    # Set a non-orthogonal box dimension. The box below works out to be this cell:
+    # [[1.        0.        0.       ]
+    # [0.        1.        0.       ]
+    # [0.5       0.        0.8660254]]
+    dim = np.asarray([1, 1, 1, 90, 60, 90], np.float32)
+    workflow = [mda.transformations.boxdimensions.set_dimensions(dim), NoJump()]
+    u.trajectory.add_transformations(*workflow)
+    transformed_coordinates = u.trajectory.timeseries()[0]
+    # After the transformation, you should end up in a repeating pattern, since you are
+    # working in a hexagonal unit cell system. Since you jump every third timestep across
+    # a periodic boundary, the shift in each axis is saved. As a consequence, the correct
+    # jump every third step is just the original position + the size of the periodic cells.
+    assert_allclose(
+        transformed_coordinates[::3],
+        np.outer(np.arange(33.5), np.array([0.5, 1, np.sqrt(3) / 2])),
+    )
+    assert_allclose(
+        transformed_coordinates[1::3],
+        np.outer(np.arange(32.5), np.array([0.5, 1, np.sqrt(3) / 2])) + 1 * np.ones(3) / 3,
+        rtol=1.2e-7
+    )
+    assert_allclose(
+        transformed_coordinates[2::3],
+        np.outer(np.arange(32.5), np.array([0.5, 1, np.sqrt(3) / 2])) + 2 * np.ones(3) / 3,
+        rtol=1.2e-7
+    )
+
+
+def test_nojump_constantvel(nojump_constantvel_universe):
+    """
+    Test if the nojump transform is returning the correct
+    values when iterating forwards over the sample trajectory.
+    """
+    ref = nojump_constantvel_universe
+    towrap = ref.copy() # This copy of the universe will be wrapped, then unwrapped,
+    # and should be equal to ref.
+    dim = np.asarray([5, 5, 5, 54, 60, 90], np.float32)
+    workflow = [
+        mda.transformations.boxdimensions.set_dimensions(dim),
+        wrap(towrap.atoms),
+        NoJump(),
+    ]
+    towrap.trajectory.add_transformations(*workflow)
+    assert_allclose(
+        towrap.trajectory.timeseries(),
+        ref.trajectory.timeseries(),
+        rtol=5e-07,
+        atol=5e-06,
+    )
+
+
+def test_nojump_constantvel_skip(nojump_universes_fromfile):
+    """
+    Test if the nojump transform warning is emitted.
+    """
+    with pytest.warns(UserWarning):
+        u = nojump_universes_fromfile
+        u.trajectory[0]
+        u.trajectory[9] #Exercises the warning.
+
+
+def test_nojump_constantvel_jumparound(nojump_universes_fromfile):
+    """
+    Test if the nojump transform is emitting a warning.
+    """
+    with pytest.warns(UserWarning):
+        u = nojump_universes_fromfile
+        u.trajectory[0]
+        u.trajectory[1]
+        u.trajectory[9]
+
+
+def test_missing_dimensions_init(nojump_universe):
+    """
+    Test if the nojump transform raises a NoDataError if there is no
+    initial dimension for the periodic unit cell.
+    """
+    with pytest.raises(mda.exceptions.NoDataError):
+        u = nojump_universe
+        workflow = [NoJump()]
+        u.trajectory.add_transformations(*workflow)
+        transformed_coordinates = u.trajectory.timeseries()[0]
+
+
+def test_missing_dimensions(nojump_universe):
+    """
+    Test if the nojump transform raises a NoDataError if there is no
+    dimension for the periodic unit cell in a subsequent timestep.
+    """
+    with pytest.raises(mda.exceptions.NoDataError):
+        u = nojump_universe
+        u.dimensions = [73, 73, 73, 90, 90, 90]
+        workflow = [NoJump()]
+        u.trajectory.add_transformations(*workflow)
+        transformed_coordinates = u.trajectory.timeseries()[0]
+
+
+def test_notinvertible(nojump_universe):
+    """
+    Test if the nojump transform raises a NoDataError if the dimensions
+    are invalid for the periodic unit cell.
+    """
+    with pytest.raises(mda.exceptions.NoDataError):
+        u = nojump_universe
+        dim = [1, 0, 0, 90, 90, 90]
+        workflow = [mda.transformations.boxdimensions.set_dimensions(dim),NoJump()]
+        u.trajectory.add_transformations(*workflow)
+        transformed_coordinates = u.trajectory.timeseries()[0]