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

…jectories (#4258)

* Ensure atoms are unwrapped correctly when they jump across boundaries on the second frame

Previosuly, jumps wouldn't be unwrapped until the third frame due to the way the class was initialised.
To fix the issue, the NoJump transformation now requires an 'ag' argument, which is used to determine the box size at the first frame

* Add testfiles for nojump transformation

* Add tests for nojump transformation where the jump occurs at the second and third frames

Also update other tests to pass the 'ag' argument to the transformation

* Remove unncessary second assignment of L in NoJump initialisation

* Temporarily put test datafiles in the test directory

* Add transformation to set variable box dimensions

For creating NPT trajectories

* Use set_variable_dimensions transformation in test for NoJump

* Remove the test datafiles for the npt NoJump

* Import set_variable_dimensions in the transformation sub-package

* Fix typos in docsttring of set_variable_dimensions transformation

* remove unncessary comments in docstring of npt tests for nojump

* Add tests for set_variable_dimensions transformation

* Add test to check that set_variable_dimensions raises an error when dimensions are missing for a frame

* Tidy formatting of set_variable_dimensions test functions

* Add docs for set_variable_dimensions

* Add versionchanged note to NoJump about adding 'ag' argument

* Add test to check if the nojump transform always returns the correct values at all frames when iterating over multiple times

* ensure coordinates can be unwrapped correctly when iterating
    over the trajectory multiple times

* Use tmp_path_factory to write tmp trajectory for testing NoJump repeated unwrapping

Also don't pass 'ag' argument to NoJump in tests

* Fix typo in boxdimensions module docstring

Co-authored-by: Hugo MacDermott-Opeskin <[email protected]>

* Fix typo in docstring of NoJump test

Co-authored-by: Hugo MacDermott-Opeskin <[email protected]>

* remove unnecessary second calculation of Linverse from NoJump

* Add a versionadded note to the set_variable_dimensions docstring

* Update CHANGELOG

* Remove added set_variable_dimensions class

Incorportate the functionality into set_dimensions

* Update tests to reflect removal of set_variable_dimensions transformation

* Update CHANGELOG

* Update NoJump tests to reflect removal of set_variable_dimensions transformation

* Remove unnecessary whitespace from NoJump tests

Co-authored-by: Irfan Alibay <[email protected]>

* Use numpy array in docstring example of set_dimensions transformation

* Update changelog

* Use ValueError rather than NoDataError in set_dimensions

* Set NoJump.older_frame to -1 when ts.frame == 0.

It can be set to any number, just not 'A'.

* Add comment to explain how we ensure jumps on the second frame are handled correctly by NoJump

* Make flake8 happy

* Remove unnecessary if-block 'if self.prev is None:' from NoJump

This will never be true

* Add test to check warning is emitted by NoJump when iterating trajectory with stride > 1

* Fix formatting of set_dimensions docstring

Co-authored-by: Irfan Alibay <[email protected]>

---------

Co-authored-by: Hugo MacDermott-Opeskin <[email protected]>
Co-authored-by: Irfan Alibay <[email protected]>

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)