Ensure NoJump handles jumps that occur on the second frame in NPT trajectories

package/CHANGELOG

@@ -13,14 +13,17 @@ The rules for this file:
   * release numbers follow "Semantic Versioning" http://semver.org
 
 ------------------------------------------------------------------------------
-??/??/?? IAlibay, hmacdope, pillose, jaclark5, tylerjereddy
+??/??/?? IAlibay, hmacdope, pillose, jaclark5, tylerjereddy, p-j-smith
 
  * 2.7.0
 
 Fixes
+  * Fix `NoJump`` unwrapping for jumps on the second frame in a
+    trajectory (PR #4258, Issue #4257)
   * Fix Atom type guessing error (PR #4168, Issue #4167)
 
 Enhancements
+  * Allow `set_dimensions` transformation to handle NPT trajectories (PR #4258)
 
 Changes
   * NumPy `in1d` replaced with `isin` in preparation for NumPy

package/MDAnalysis/transformations/boxdimensions.py

@@ -25,44 +25,63 @@
 Set box dimensions --- :mod:`MDAnalysis.transformations.boxdimensions`
 =======================================================================
 
-Set dimensions of the simulation box to a constant vector across all timesteps.
-
+Set dimensions of the simulation box, either to a constant vector across
+all timesteps or to a specified vector at each frame.
 
 .. autoclass:: set_dimensions
 """
 import numpy as np
 
 from .base import TransformationBase
 
-
 class set_dimensions(TransformationBase):
     """
     Set simulation box dimensions.
 
     Timestep dimensions are modified in place.
 
-    Example
-    -------
+    Examples
+    --------
 
     e.g. set simulation box dimensions to a vector containing unit cell
     dimensions [*a*, *b*, *c*, *alpha*, *beta*, *gamma*], lengths *a*,
     *b*, *c* are in the MDAnalysis length unit (Å), and angles are in degrees.
+    The same dimensions will be used for every frame in the trajectory.
 
     .. code-block:: python
 
-        dim = [2, 2, 2, 90, 90, 90]
+        dim = np.array([2, 2, 2, 90, 90, 90])
+        transform = mda.transformations.boxdimensions.set_dimensions(dim)
+        u.trajectory.add_transformations(transform)
+
+    Or e.g. set simulation box dimensions to a vector containing unit cell
+    dimensions [*a*, *b*, *c*, *alpha*, *beta*, *gamma*] at the first frame,
+    and [*2a*, *2b*, *2c*, *alpha*, *beta*, *gamma*] at the second frame.
+
+    .. code-block:: python
+
+        dim = np.array([
+            [2, 2, 2, 90, 90, 90],
+            [4, 4, 4, 90, 90, 90],
+        ])
         transform = mda.transformations.boxdimensions.set_dimensions(dim)
         u.trajectory.add_transformations(transform)
 
     Parameters
     ----------
-    dimensions: iterable of floats
+    dimensions: iterable of floats or two-dimensional np.typing.NDArrayLike
         vector that contains unit cell lengths and angles.
-        Expected shapes are (6, 0) or (1, 6)
+        Expected shapes are (6, 0) or (1, 6) or (N, 6), where
+        N is the number of frames in the trajectory. If shape is (6, 0) or
+        (1, 6), the same dimensions will be used at every frame in the
+        trajectory.
 
     Returns
     -------
     :class:`~MDAnalysis.coordinates.timestep.Timestep` object
+
+    .. versionchanged:: 2.7.0
+       Added the option to set varying box dimensions (i.e. an NPT trajectory).
     """
 
     def __init__(self,
@@ -80,15 +99,20 @@ def __init__(self,
                        'np.float32 numpy.ndarray')
             raise ValueError(errmsg)
         try:
-            self.dimensions = self.dimensions.reshape(6, )
+            self.dimensions = self.dimensions.reshape(-1, 6)
         except ValueError:
             errmsg = (f'{self.dimensions} array does not have valid box '
                        'dimension shape.\nSimulation box dimensions are '
-                       'given by an float array of shape (6, ), '
-                       ' containing 3 lengths and 3 angles: '
+                       'given by an float array of shape (6, 0), (1, 6), '
+                       'or (N, 6) where N is the number of frames in the '
+                       'trajectory and the dimension vector(s) containing '
+                       '3 lengths and 3 angles: '
                        '[a, b, c, alpha, beta, gamma]')
             raise ValueError(errmsg)
 
     def _transform(self, ts):
-        ts.dimensions = self.dimensions
+        try:
+            ts.dimensions = self.dimensions[0] if self.dimensions.shape[0] == 1 else self.dimensions[ts.frame]
+        except IndexError as e:
+            raise ValueError(f"Dimensions array has no data for frame {ts.frame}") from e
         return ts

package/MDAnalysis/transformations/nojump.py

@@ -123,6 +123,17 @@ def _transform(self, ts):
         except np.linalg.LinAlgError:
             msg = f"Periodic box dimensions are not invertible at step {ts.frame}"
             raise NoDataError(msg)
+        if ts.frame == 0:
+            # We don't need to apply the transformation here. However, we need to ensure we
+            # have the 0th frame coordinates in reduced form. We also need to set an
+            # an appropriate value for self.older frame. This is so that on the following frame
+            # we don't return early when we check `self.older_frame != "A"`. If we return early,
+            # then the transformation is not applied, and any jumps across boundaries that occur
+            # at that frame will not be accounted for.
+            self.prev = ts.positions @ Linverse
+            self.old_frame = 0
+            self.older_frame = -1
+            return ts
         if self.prev is None:
             self.prev = ts.positions @ Linverse
             self.old_frame = ts.frame

testsuite/MDAnalysisTests/datafiles.py

@@ -623,6 +623,5 @@
 SURFACE_PDB = (_data_ref / 'surface.pdb.bz2').as_posix()
 SURFACE_TRR = (_data_ref / 'surface.trr').as_posix()
 
-
 # This should be the last line: clean up namespace
 del resources

testsuite/MDAnalysisTests/transformations/test_boxdimensions.py

@@ -37,6 +37,17 @@ def boxdimensions_universe():
     return new_u
 
 
+@pytest.fixture()
+def variable_boxdimensions_universe():
+    # create Universe objects for tests
+    n_atoms = 1
+    n_frames = 3
+    u = mdanalysis.Universe.empty(n_atoms, trajectory=True)
+    coordinates = np.zeros((n_frames, u.atoms.n_atoms, 3))
+    u.load_new(coordinates, order="fac")
+    return u
+
+
 def test_boxdimensions_dims(boxdimensions_universe):
     new_dims = np.float32([2, 2, 2, 90, 90, 90])
     set_dimensions(new_dims)(boxdimensions_universe.trajectory.ts)
@@ -72,7 +83,6 @@ def test_dimensions_vector_asarray(boxdimensions_universe,
     with pytest.raises(ValueError, match='cannot be converted'):
         set_dimensions(dim_vector_forms_dtypes)(ts)
 
-
 def test_dimensions_transformations_api(boxdimensions_universe):
     # test if transformation works with on-the-fly transformations API
     new_dims = np.float32([2, 2, 2, 90, 90, 90])
@@ -81,3 +91,39 @@ def test_dimensions_transformations_api(boxdimensions_universe):
     for ts in boxdimensions_universe.trajectory:
         assert_array_almost_equal(boxdimensions_universe.dimensions,
                                   new_dims, decimal=6)
+
+
+def test_varying_dimensions_transformations_api(
+    variable_boxdimensions_universe,
+):
+    """
+    Test if transformation works with on-the-fly transformations API
+    when we have varying dimensions.
+    """
+    new_dims = np.float32([
+        [2, 2, 2, 90, 90, 90],
+        [4, 4, 4, 90, 90, 90],
+        [8, 8, 8, 90, 90, 90],
+    ])
+    transform = set_dimensions(new_dims)
+    variable_boxdimensions_universe.trajectory.add_transformations(transform)
+    for ts in variable_boxdimensions_universe.trajectory:
+        assert_array_almost_equal(variable_boxdimensions_universe.dimensions,
+                                  new_dims[ts.frame], decimal=6)
+
+
+def test_varying_dimensions_no_data(
+    variable_boxdimensions_universe,
+):
+    """
+    Check error is raised when dimensions are missing for a frame
+    in a trajectory.
+    """
+    # trjactory has three frames
+    new_dims = np.float32([
+        [2, 2, 2, 90, 90, 90],
+        [4, 4, 4, 90, 90, 90],
+    ])
+    transform = set_dimensions(new_dims)
+    with pytest.raises(ValueError, match="Dimensions array has no data for frame 2"):
+        variable_boxdimensions_universe.trajectory.add_transformations(transform)

testsuite/MDAnalysisTests/transformations/test_nojump.py

@@ -48,6 +48,90 @@ def nojump_constantvel_universe():
     return reference
 
 
+@pytest.fixture
+def nojump_universe_npt_2nd_frame():
+    """
+    Create a Universe in which a single atom jumps across the periodic boundary in the
+    x-dimensions at the second frame.
+
+    Unwrapped coordinates should all be 97.5.
+    """
+    n_atoms = 1
+    n_frames = 4
+    u = mda.Universe.empty(n_atoms, trajectory=True)
+    coordinates = np.empty((n_frames, u.atoms.n_atoms, 3))
+    coordinates[0] = [97.5, 50.0, 50.0]
+    coordinates[1] = [2.5, 50.0, 50.0]
+    coordinates[2] = [2.5, 50.0, 50.0]
+    coordinates[3] = [2.5, 50.0, 50.0]
+    u.load_new(coordinates, order="fac")
+    return u
+
+
+@pytest.fixture
+def nojump_universe_npt_3rd_frame():
+    """
+    Create a Universe in which a single atom jumps across the periodic boundary in the
+    x-dimensions at the third frame.
+
+    Unwrapped coordinates should all be 97.5.
+    """
+    n_atoms = 1
+    n_frames = 4
+    u = mda.Universe.empty(n_atoms, trajectory=True)
+    coordinates = np.empty((n_frames, u.atoms.n_atoms, 3))
+    coordinates[0] = [97.5, 50.0, 50.0]
+    coordinates[1] = [97.5, 50.0, 50.0]
+    coordinates[2] = [2.5, 50.0, 50.0]
+    coordinates[3] = [2.5, 50.0, 50.0]
+    u.load_new(coordinates, order="fac")
+    return u
+
+
+@pytest.fixture(scope="module")
+def nojump_universe_npt_2nd_frame_from_file(tmp_path_factory):
+    """
+    Write the `nojump_universe_npt_2nd_frame` fixture to file, read it in and
+    return the Universe.
+
+    Used for testing that coordinates can be unwrapped correctly when iterating
+    over the trajectory multiple times.
+
+    We can't use an in-memory trajectory to test this because the
+    transformation would only be applied once.
+
+    Note, we use `tmp_path_factory` because this fixture requies `module` scope
+    so we can read the file after the fixture has been created, and
+    `tmp_path` has function-level scope.
+    """
+    n_atoms = 1
+    n_frames = 4
+    u = mda.Universe.empty(n_atoms, trajectory=True)
+    coordinates = np.empty((n_frames, u.atoms.n_atoms, 3))
+    coordinates[0] = [97.5, 50.0, 50.0]
+    coordinates[1] = [2.5, 50.0, 50.0]
+    coordinates[2] = [2.5, 50.0, 50.0]
+    coordinates[3] = [2.5, 50.0, 50.0]
+    u.load_new(coordinates, order="fac")
+    dim = np.asarray([
+        [100, 100, 100, 90, 90, 90],
+        [95, 100, 100, 90, 90, 90], # Box shrinks by 5 in the x-dimension at the second frame
+        [95, 100, 100, 90, 90, 90],
+        [95, 100, 100, 90, 90, 90],
+    ])
+    workflow = [
+        mda.transformations.boxdimensions.set_dimensions(dim),
+    ]
+    u.trajectory.add_transformations(*workflow)
+    tmp_pdb = (tmp_path_factory.getbasetemp() / "nojump_npt_2nd_frame.pdb").as_posix()
+    tmp_xtc = (tmp_path_factory.getbasetemp() / "nojump_npt_2nd_frame.xtc").as_posix()
+    u.atoms.write(tmp_pdb)
+    with mda.Writer(tmp_xtc) as f:
+        for ts in u.trajectory:
+            f.write(u.atoms)
+    return mda.Universe(tmp_pdb, tmp_xtc)
+
+
 def test_nojump_orthogonal_fwd(nojump_universe):
     """
     Test if the nojump transform is returning the correct
@@ -122,6 +206,86 @@ def test_nojump_constantvel(nojump_constantvel_universe):
     )
 
 
+def test_nojump_2nd_frame(nojump_universe_npt_2nd_frame):
+    """
+    Test if the nojump transform returns the correct values
+    at all frames when iterating over an npt trajectory
+    and an atom crosses the x-boundary at the second frame
+
+    Wrapped coordinates are:
+    coordinates = [
+        [97.5, 50.0, 50.0],
+        [2.5, 50.0, 50.0],
+        [2.5, 50.0, 50.0],
+        [2.5, 50.0, 50.0],
+    ]
+    where each row is a different timestep.
+
+    Unwrapped coordinates are the same at each frame:
+    unwrapped = [97.5, 50.0, 50.0]
+    """
+    u = nojump_universe_npt_2nd_frame
+    dim = np.asarray([
+        [100, 100, 100, 90, 90, 90],
+        [95, 100, 100, 90, 90, 90], # Box shrinks by 5 in the x-dimension at the second frame
+        [95, 100, 100, 90, 90, 90],
+        [95, 100, 100, 90, 90, 90],
+    ])
+    workflow = [
+        mda.transformations.boxdimensions.set_dimensions(dim),
+        NoJump(),
+    ]
+    u.trajectory.add_transformations(*workflow)
+    x_position = 97.5
+    np.testing.assert_allclose(u.trajectory.timeseries()[:, 0, 0], x_position)
+
+
+def test_nojump_3rd_frame(nojump_universe_npt_3rd_frame):
+    """
+    Test if the nojump transform returns the correct values
+    at all frames when iterating over an npt trajectory
+    and an atom crosses the x-boundary at the third frame.
+
+    Wrapped coordinates are:
+    coordinates = [
+        [97.5, 50.0, 50.0],
+        [97.5, 50.0, 50.0],
+        [2.5, 50.0, 50.0],
+        [2.5, 50.0, 50.0],
+    ]
+    where each row is a different timestep.
+
+    Unwarpped coordinates are the same at each frame:
+    unwrapped = [97.5, 50.0, 50.0]
+    """
+    u = nojump_universe_npt_3rd_frame
+    dim = np.asarray([
+        [100, 100, 100, 90, 90, 90],
+        [100, 100, 100, 90, 90, 90],
+        [95, 100, 100, 90, 90, 90],  # Box shrinks by 5 in the x-dimension at the third frame
+        [95, 100, 100, 90, 90, 90],
+    ])
+    workflow = [
+        mda.transformations.boxdimensions.set_dimensions(dim),
+        NoJump(),
+    ]
+    u.trajectory.add_transformations(*workflow)
+    x_position = 97.5
+    np.testing.assert_allclose(u.trajectory.timeseries()[:, 0, 0], x_position)
+
+
+def test_nojump_iterate_twice(nojump_universe_npt_2nd_frame_from_file):
+    """
+    Test if the nojump transform always returns the correct values
+    at all frames when iterating over multiple times.
+    """
+    u = nojump_universe_npt_2nd_frame_from_file
+    u.trajectory.add_transformations(NoJump())
+    timeseries_first_iteration = u.trajectory.timeseries()
+    timeseries_second_iteration = u.trajectory.timeseries()
+    np.testing.assert_allclose(timeseries_first_iteration, timeseries_second_iteration)
+
+
 def test_nojump_constantvel_skip(nojump_universes_fromfile):
     """
     Test if the nojump transform warning is emitted.