Nojump

package/AUTHORS

@@ -206,6 +206,7 @@ Chronological list of authors
   - Christian Pfaendner
   - Pratham Chauhan
   - Meet Brijwani
+  - Josh Vermaas
 
 
 External code

package/CHANGELOG

@@ -26,6 +26,8 @@ Fixes
     (Issue #3336)
 
 Enhancements
+  * Add a nojump transformation, which unwraps trajectories so that particle
+    paths are continuous. (Issue #3703, PR #4031)
   * Add distopia distance calculation library bindings as a selectable backend
     for `calc_bonds` in `MDA.lib.distances`. (Issue #3783, PR #3914)
   * AuxReaders are now pickle-able and copy-able (Issue #1785, PR #3887)

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,127 @@
+# -*- 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 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
+
+
+class NoJump(TransformationBase):
+    """
+    Returns transformed coordinates for the given timestep so that an atom
+    does not move more than half the periodic box size, and will not jump
+    across the periodic boundary. The algorithm used is based on :cite:p:`Kulke2022`, 
+    equation B6 for non-orthogonal systems.
+
+    Example
+    -------
+
+    To calculate diffusion, one of the first steps is to compute a trajectory
+    where the atoms do not cross the periodic boundary.
+
+    .. 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.
+    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, "
+               "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.check_c = check_continuity
+
+    def _transform(self, ts):
+        if self.prev is None:
+            self.prev = ts.positions @ np.linalg.inv(ts.triclinic_dimensions)
+            self.old_frame = ts.frame
+            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.",
+                Warning,
+            )
+        # Remember that @ is a shorthand for matrix multiplication.
+        # np.matmul(a, b) is equivalent to a @ b
+        L = ts.triclinic_dimensions
+        Linverse = np.linalg.inv(L)
+        # 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.old_frame = ts.frame
+
+        return ts

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,57 @@
+###
+
+import numpy as np
+import pytest
+from numpy.testing import assert_array_almost_equal
+
+import MDAnalysis as mda
+from MDAnalysis.transformations import NoJump
+from MDAnalysisTests import datafiles as data
+
+
+@pytest.fixture()
+def nojump_universes():
+    """
+    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_universes_nocheck():
+    """
+    Create the universe objects for the tests.
+    Do not check for continunity
+    """
+    u = mda.Universe(data.PSF_TRICLINIC, data.DCD_TRICLINIC)
+    transformation = NoJump(check_continuity=False)
+    u.trajectory.add_transformations(transformation)
+    return u
+
+
+def test_nojump_fwd(nojump_universes):
+    """
+    Test if the nojump transform is returning the correct
+    values when iterating forwards over the sample trajectory.
+    """
+    # These were determined based on determining what the TRICLINIC system
+    # would return on a validated version that worked on triclinic and
+    # orthogonal systems for a small water trajectory.
+    # These are specific to atoms 166 and 362, which moved the most in the
+    # trajectory.
+    ref_matrix_fwd1 = np.asarray([-3.42261, 11.28495, -1.37211])
+    ref_matrix_fwd2 = np.asarray([3.674243, -8.725193, -0.07884017])
+    size = (
+        nojump_universes.trajectory.ts.positions.shape[0],
+        nojump_universes.trajectory.ts.positions.shape[1],
+        len(nojump_universes.trajectory),
+    )
+    parr = np.empty(size)
+    for ts in nojump_universes.trajectory:
+        parr[..., ts.frame] = ts.positions.copy()
+    # Atoms 166 and 362 happen to move alot.
+    assert_array_almost_equal(ref_matrix_fwd1, parr[166, :, -1], decimal=5)
+    assert_array_almost_equal(ref_matrix_fwd2, parr[362, :, -1], decimal=5)