Add Transformation to set box dimensions (#3175)

Fixes #2691 
Co-authored-by: Lily Wang <[email protected]>
Co-authored-by: Irfan Alibay <[email protected]>

package/AUTHORS

@@ -159,6 +159,7 @@ Chronological list of authors
   - Jan Stevens
   - Orion Cohen
   - Dimitrios Papageorgiou
+  - Hannah Pollak
 
 External code
 -------------

package/CHANGELOG

@@ -17,7 +17,7 @@ The rules for this file:
          lilyminium, daveminh, jbarnoud, yuxuanzhuang, VOD555, ianmkenney,
          calcraven，xiki-tempula, mieczyslaw, manuel.nuno.melo, PicoCentauri,
          hanatok, rmeli, aditya-kamath, tirkarthi, LeonardoBarneschi, hejamu,
-         biogen98, orioncohen, z3y50n
+         biogen98, orioncohen, z3y50n, hp115
 
   * 2.0.0
 
@@ -89,6 +89,8 @@ Fixes
   * Fix syntax warning over comparison of literals using is (Issue #3066)
 
 Enhancements
+  * Add `set_dimensions` transformation class for setting constant 
+    box dimensions for all timesteps in trajectory (Issue #2691)
   * Added a ValueError raised when not given a gridcenter while
     providing grid dimensions to DensityAnalysis, also added
     check for NaN in the input (Issue #3148, PR #3154)

package/MDAnalysis/transformations/__init__.py

@@ -130,3 +130,4 @@ def wrapped(ts):
 from .positionaveraging import PositionAverager
 from .fit import fit_translation, fit_rot_trans
 from .wrap import wrap, unwrap
+from .boxdimensions import set_dimensions

package/MDAnalysis/transformations/boxdimensions.py

@@ -0,0 +1,87 @@
+# -*- 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
+#
+
+"""\
+Set box dimensions --- :mod:`MDAnalysis.transformations.boxdimensions`
+=======================================================================
+
+Set dimensions of the simulation box to a constant vector across all timesteps.
+
+
+.. autoclass:: set_dimensions
+"""
+import numpy as np
+
+
+class set_dimensions:
+    """
+    Set simulation box dimensions.
+
+    Timestep dimensions are modified in place.
+
+    Example
+    -------
+
+    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.
+
+    .. code-block:: python
+
+        dim = [2, 2, 2, 90, 90, 90]
+        transform = mda.transformations.boxdimensions.set_dimensions(dim)
+        u.trajectory.add_transformations(transform)
+
+    Parameters
+    ----------
+    dimensions: iterable of floats
+        vector that contains unit cell lengths and angles.
+        Expected shapes are (6, 0) or (1, 6)
+
+    Returns
+    -------
+    :class:`~MDAnalysis.coordinates.base.Timestep` object
+    """
+
+    def __init__(self, dimensions):
+        self.dimensions = dimensions
+
+        try:
+            self.dimensions = np.asarray(self.dimensions, np.float32)
+        except ValueError:
+            errmsg = (f'{self.dimensions} cannot be converted into '
+                       'np.float32 numpy.ndarray')
+            raise ValueError(errmsg)
+        try:
+            self.dimensions = self.dimensions.reshape(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]')
+            raise ValueError(errmsg)
+
+    def __call__(self, ts):
+        ts.dimensions = self.dimensions
+        return ts

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

@@ -233,4 +233,5 @@ Currently implemented transformations
    ./transformations/positionaveraging
    ./transformations/fit
    ./transformations/wrap
+   ./transformations/boxdimensions
 

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

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

testsuite/MDAnalysisTests/transformations/test_boxdimensions.py

@@ -0,0 +1,83 @@
+# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding:utf-8 -*-
+# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4 fileencoding=utf-8
+#
+# 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
+#
+
+import numpy as np
+import pytest
+from numpy.testing import assert_array_almost_equal
+
+import MDAnalysis as mdanalysis
+from MDAnalysis.transformations import set_dimensions
+from MDAnalysisTests import make_Universe
+
+
+@pytest.fixture()
+def boxdimensions_universe():
+    # create Universe objects for tests
+    new_u = make_Universe(trajectory=True)
+    return new_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)
+    assert_array_almost_equal(boxdimensions_universe.dimensions,
+                              new_dims, decimal=6)
+
+
+@pytest.mark.parametrize('dim_vector_shapes', (
+    [1, 1, 1, 90, 90],
+    [1, 1, 1, 1, 90, 90, 90],
+    np.array([[1], [1], [90], [90], [90]]),
+    np.array([1, 1, 1, 90, 90]),
+    np.array([1, 1, 1, 1, 90, 90, 90]),
+    [1, 1, 1, 90, 90],
+    111909090)
+    )
+def test_dimensions_vector(boxdimensions_universe, dim_vector_shapes):
+    # wrong box dimension vector shape
+    ts = boxdimensions_universe.trajectory.ts
+    with pytest.raises(ValueError, match='valid box dimension shape'):
+        set_dimensions(dim_vector_shapes)(ts)
+
+
+@pytest.mark.parametrize('dim_vector_forms_dtypes', (
+    ['a', 'b', 'c', 'd', 'e', 'f'],
+    np.array(['a', 'b', 'c', 'd', 'e', 'f']),
+    'abcd')
+    )
+def test_dimensions_vector_asarray(boxdimensions_universe,
+                                   dim_vector_forms_dtypes):
+    # box dimension input type not convertible into array
+    ts = boxdimensions_universe.trajectory.ts
+    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])
+    transform = set_dimensions(new_dims)
+    boxdimensions_universe.trajectory.add_transformations(transform)
+    for ts in boxdimensions_universe.trajectory:
+        assert_array_almost_equal(boxdimensions_universe.dimensions,
+                                  new_dims, decimal=6)