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

package/CHANGELOG

@@ -13,11 +13,13 @@ The rules for this file:
   * release numbers follow "Semantic Versioning" http://semver.org
 
 ------------------------------------------------------------------------------
-??/??/?? IAlibay, ianmkenney, PicoCentauri, pgbarletta
+??/??/?? IAlibay, ianmkenney, PicoCentauri, pgbarletta, p-j-smith
 
  * 2.7.0
 
 Fixes
+  * Fix `NoJump` unwrapping for jumps on the second frame in a
+    trajectory (PR #4258, Issue #4257)
   * Deprecated np.float_ and np.NaN aliases have been replaced with
     their original type value (np.float64 and np.nan) (PR #4272)
 
@@ -45,6 +47,7 @@ Fixes
   * Fix Atom type guessing error (PR #4168, Issue #4167)
 
 Enhancements
+  * Allow `set_dimensions` transformation to handle NPT trajectories (PR #4258)
 
 Changes
   * Reverts PR #4108, builds are now again made using the oldest

package/MDAnalysis/transformations/boxdimensions.py

@@ -25,44 +25,64 @@
 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 = 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 = [2, 2, 2, 90, 90, 90]
+        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,
@@ -76,19 +96,34 @@ def __init__(self,
         try:
             self.dimensions = np.asarray(self.dimensions, np.float32)
         except ValueError:
-            errmsg = (f'{self.dimensions} cannot be converted into '
-                       'np.float32 numpy.ndarray')
+            errmsg = (
+                f'{self.dimensions} cannot be converted into '
+                '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: '
-                       '[a, b, c, alpha, beta, gamma]')
+            errmsg = (
+                f'{self.dimensions} array does not have valid box '
+                'dimension shape.\nSimulation box dimensions are '
+                '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,9 +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 self.prev is None:
+        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 = ts.frame
+            self.old_frame = 0
+            self.older_frame = -1
             return ts
         if (
             self.check_c

testsuite/MDAnalysisTests/datafiles.py

@@ -633,6 +633,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)