Implements compact wrapping and nojump, and updates associated transforms.

package/CHANGELOG

@@ -15,7 +15,7 @@ The rules for this file:
 ------------------------------------------------------------------------------
 ??/??/?? tylerjereddy, richardjgowers, IAlibay, hmacdope, orbeckst, cbouy,
          lilyminium, daveminh, jbarnoud, yuxuanzhuang, VOD555, ianmkenney,
-         calcraven，xiki-tempula, mieczyslaw
+         calcraven，xiki-tempula, mieczyslaw, manuel.nuno.melo
 
   * 2.0.0
 
@@ -60,6 +60,9 @@ Fixes
     (Issue #2934)
 
 Enhancements
+  * Can now do wrapping to compact unit cell representations; sped up
+    `AtomGroup.wrap()` treatment of compounds (PR #2982)
+  * Added a `transformations.nojump` (PR #2982)
   * Refactored analysis.helanal into analysis.helix_analysis
     (Issue #2452, PR #2622)
   * Improved MSD code to accept `AtomGroup` and reject `UpdatingAtomGroup`

package/MDAnalysis/core/groups.py

@@ -1345,7 +1345,8 @@ def pack_into_box(self, box=None, inplace=True):
         """
         return self.wrap(box=box, inplace=inplace)
 
-    def wrap(self, compound="atoms", center="com", box=None, inplace=True):
+    def wrap(self, compound="atoms", center="com", compact=False, box=None,
+             inplace=True):
         r"""Shift the contents of this group back into the primary unit cell
         according to periodic boundary conditions.
 
@@ -1354,6 +1355,28 @@ def wrap(self, compound="atoms", center="com", box=None, inplace=True):
         ``'atoms'``, each compound as a whole will be shifted so that its
         `center` lies within the primary unit cell.
 
+        Wrapping with `compound`='atoms':
+           +-----------+              +-----------+
+           |           |              |           |
+           |           | 3 6          | 3 6       |
+           |           | ! !          | ! !       |
+           |         1-|-2-5-8  ->    |-2-5-8   1-|
+           |           | ! !          | ! !       |
+           |           | 4 7          | 4 7       |
+           |           |              |           |
+           +-----------+              +-----------+
+
+        Wrapping with `compound`='fragments':
+           +-----------+              +-----------+
+           |           |              |           |
+           |           | 3 6          | 3 6       |
+           |           | ! !          | ! !       |
+           |         1-|-2-5-8  ->  1-|-2-5-8     |
+           |           | ! !          | ! !       |
+           |           | 4 7          | 4 7       |
+           |           |              |           |
+           +-----------+              +-----------+
+
         Parameters
         ----------
         compound : {'atoms', 'group', 'segments', 'residues', 'molecules', \
@@ -1364,6 +1387,10 @@ def wrap(self, compound="atoms", center="com", box=None, inplace=True):
         center : {'com', 'cog'}
             How to define the center of a given group of atoms. If `compound` is
             ``'atoms'``, this parameter is meaningless and therefore ignored.
+        compact : bool, optional
+            If ``True``, wrapping will be done to the space-filling volume
+            closest to the primary unit cell's center (only useful for
+            non-orthogonal periodic boxes).
         box : array_like, optional
             The unitcell dimensions of the system, which can be orthogonal or
             triclinic and must be provided in the same format as returned by
@@ -1407,6 +1434,12 @@ def wrap(self, compound="atoms", center="com", box=None, inplace=True):
 
            x_i' = x_i - \left\lfloor\frac{x_i}{L_i}\right\rfloor\,.
 
+        `compact` changes the shifting behavior: centers will be shifted by
+        the integer combination of box vectors that minimizes
+        :math:`\left\|\mathbf{r_{compact}}-\mathbf{r_{box\_center}}\right\|`,
+        the distance of a center's final position,
+        :math:`\mathbf{r_{compact}}`, to the primary unit cell's center.
+
         When specifying a `compound`, the translation is calculated based on
         each compound. The same translation is applied to all atoms
         within this compound, meaning it will not be broken by the shift.
@@ -1417,8 +1450,9 @@ def wrap(self, compound="atoms", center="com", box=None, inplace=True):
         taken into account!
 
         `center` allows to define how the center of each group is computed.
-        This can be either ``'com'`` for center of mass, or ``'cog'`` for center
-        of geometry.
+        This can be either ``'com'`` for center of mass, or ``'cog'`` for
+        center of geometry.
+
 
         `box` allows a unit cell to be given for the transformation. If not
         specified, the :attr:`~MDAnalysis.coordinates.base.Timestep.dimensions`
@@ -1451,6 +1485,7 @@ def wrap(self, compound="atoms", center="com", box=None, inplace=True):
             raise ValueError("Invalid box: Box has invalid shape or not all "
                              "box dimensions are positive. You can specify a "
                              "valid box using the 'box' argument.")
+        packer = distances.apply_compact_PBC if compact else distances.apply_PBC
 
         # no matter what kind of group we have, we need to work on its (unique)
         # atoms:
@@ -1472,7 +1507,7 @@ def wrap(self, compound="atoms", center="com", box=None, inplace=True):
             return np.zeros((0, 3), dtype=np.float32)
 
         if comp == "atoms" or len(atoms) == 1:
-            positions = distances.apply_PBC(atoms.positions, box)
+            positions = packer(atoms.positions, box)
         else:
             ctr = center.lower()
             if ctr == 'com':
@@ -1496,7 +1531,7 @@ def wrap(self, compound="atoms", center="com", box=None, inplace=True):
                 else:  # ctr == 'cog'
                     ctrpos = atoms.center_of_geometry(pbc=False, compound=comp)
                 ctrpos = ctrpos.astype(np.float32, copy=False)
-                target = distances.apply_PBC(ctrpos, box)
+                target = packer(ctrpos, box)
                 positions += target - ctrpos
             else:
                 if comp == 'segments':
@@ -1510,6 +1545,7 @@ def wrap(self, compound="atoms", center="com", box=None, inplace=True):
                         errmsg = ("Cannot use compound='molecules', this "
                                   "requires molnums.")
                         raise NoDataError(errmsg) from None
+
                 else:  # comp == 'fragments'
                     try:
                         compound_indices = atoms.fragindices
@@ -1529,16 +1565,16 @@ def wrap(self, compound="atoms", center="com", box=None, inplace=True):
                 else:  # ctr == 'cog'
                     ctrpos = atoms.center_of_geometry(pbc=False, compound=comp)
                 ctrpos = ctrpos.astype(np.float32, copy=False)
-                target = distances.apply_PBC(ctrpos, box)
+                target = packer(ctrpos, box)
                 shifts = target - ctrpos
 
+                # There may be gaps in the used indices. We must translate
+                # them to a (0, 1, 2, ..., n_indices-1) sequence
+                unique_indices = unique_int_1d(compound_indices)
+                transl_indices = np.arange(unique_indices.max()+1)
+                transl_indices[unique_indices] = np.arange(len(unique_indices))
                 # apply the shifts:
-                unique_compound_indices = unique_int_1d(compound_indices)
-                shift_idx = 0
-                for i in unique_compound_indices:
-                    mask = np.where(compound_indices == i)
-                    positions[mask] += shifts[shift_idx]
-                    shift_idx += 1
+                positions += shifts[transl_indices[compound_indices]]
 
         if inplace:
             atoms.positions = positions
@@ -1559,7 +1595,7 @@ def unwrap(self, compound='fragments', reference='com', inplace=True):
            |           |       |           |
            | 6       3 |       |         3 | 6
            | !       ! |       |         ! | !
-           |-5-8   1-2-|  ==>  |       1-2-|-5-8
+           |-5-8   1-2-|  ->   |       1-2-|-5-8
            | !       ! |       |         ! | !
            | 7       4 |       |         4 | 7
            |           |       |           |

package/MDAnalysis/lib/distances.py

@@ -61,13 +61,15 @@
 .. autofunction:: calc_angles
 .. autofunction:: calc_dihedrals
 .. autofunction:: apply_PBC
+.. autofunction:: apply_compact_PBC
 .. autofunction:: transform_RtoS
 .. autofunction:: transform_StoR
 .. autofunction:: augment_coordinates(coordinates, box, r)
 .. autofunction:: undo_augment(results, translation, nreal)
 """
 import numpy as np
 from numpy.lib.utils import deprecate
+import itertools
 
 from .util import check_coords, check_box
 from .mdamath import triclinic_vectors
@@ -1133,7 +1135,7 @@ def transform_RtoS(coords, box, backend="serial"):
     Returns
     -------
     newcoords : numpy.ndarray (``dtype=numpy.float32``, ``shape=coords.shape``)
-        An array containing fractional coordiantes.
+        An array containing fractional coordinates.
 
 
     .. versionchanged:: 0.13.0
@@ -1482,8 +1484,8 @@ def calc_dihedrals(coords1, coords2, coords3, coords4, box=None, result=None,
 
 
 @check_coords('coords')
-def apply_PBC(coords, box, backend="serial"):
-    """Moves coordinates into the primary unit cell.
+def apply_PBC(coords, box, center=None, backend="serial"):
+    """Moves coordinates into an arbitrarily centered unit cell.
 
     Parameters
     ----------
@@ -1495,12 +1497,15 @@ def apply_PBC(coords, box, backend="serial"):
         triclinic and must be provided in the same format as returned by
         :attr:`MDAnalysis.coordinates.base.Timestep.dimensions`:
         ``[lx, ly, lz, alpha, beta, gamma]``.
+    center : numpy.ndarray
+        Coordinate array of shape ``(3,)`` of the center of the output unit
+        cell representation. Defaults to the primary unit cell's center.
     backend : {'serial', 'OpenMP'}, optional
         Keyword selecting the type of acceleration.
 
     Returns
     -------
-    newcoords : numpy.ndarray  (``dtype=numpy.float32``, ``shape=coords.shape``)
+    newcoords : numpy.ndarray (``dtype=numpy.float32``, ``shape=coords.shape``)
         Array containing coordinates that all lie within the primary unit cell
         as defined by `box`.
 
@@ -1511,13 +1516,101 @@ def apply_PBC(coords, box, backend="serial"):
     .. versionchanged:: 0.19.0
        Internal dtype conversion of input coordinates to ``numpy.float32``.
        Now also accepts (and, likewise, returns) single coordinates.
+    .. versionchanged:: 2.0.0
+       Added *center* keyword.
     """
     if len(coords) == 0:
         return coords
     boxtype, box = check_box(box)
+
+    if center is not None:
+        box_center = box.sum(axis=0) * 0.5
+        center_displacement = box_center - center
+        coords += center_displacement
+
     if boxtype == 'ortho':
         _run("ortho_pbc", args=(coords, box), backend=backend)
     else:
         _run("triclinic_pbc", args=(coords, box), backend=backend)
 
+    if center is not None:
+        coords -= center_displacement
+
+    return coords
+
+
+@check_coords('coords')
+def apply_compact_PBC(coords, box, center=None, backend="serial"):
+    """Moves coordinates into an arbitrarily centered, space-filling compact
+       volume.
+
+    Parameters
+    ----------
+    coords : numpy.ndarray
+        Coordinate array of shape ``(3,)`` or ``(n, 3)`` (dtype is arbitrary,
+        will be converted to ``numpy.float32`` internally).
+    box : numpy.ndarray
+        The unitcell dimensions of the system, which can be orthogonal or
+        triclinic and must be provided in the same format as returned by
+        :attr:`MDAnalysis.coordinates.base.Timestep.dimensions`:
+        ``[lx, ly, lz, alpha, beta, gamma]``.
+    center : numpy.ndarray
+        Coordinate array of shape ``(3,)`` of the center of the output compact
+        representation. Defaults to the primary unit cell's center.
+    backend : {'serial', 'OpenMP'}, optional
+        Keyword selecting the type of acceleration.
+
+    Returns
+    -------
+    newcoords : numpy.ndarray (``dtype=numpy.float32``, ``shape=coords.shape``)
+        Array containing coordinates that all lie within a compact
+        space-filling volume centered on the point given with `center` (compact
+        in that all that volume's points are closest to the center than to any
+        of the center's periodic images).
+
+
+    .. versionadded:: 2.0.0
+    """
+    boxtype, tric_vecs = check_box(box)
+    # shortcut for orthogonal boxes
+    if boxtype == 'ortho':
+        return apply_PBC(coords, box, center=center, backend=backend)
+
+    pbc_img_coeffs = np.array(list(itertools.product([0, -1, 1], repeat=3)))
+    pbc_img_vecs = np.matmul(pbc_img_coeffs, tric_vecs)
+
+    # Stuff could get sped up because this can work with squared distance;
+    # no accelerated function exists for it yet, though.
+    min_compact_norm = np.linalg.norm(pbc_img_vecs[1:], axis=1).min() * 0.5
+
+    if center is None:
+        center = tric_vecs.sum(axis=0) * 0.5
+    else:
+        center = np.asarray(center, dtype=np.float32)
+    pbc_imgs = pbc_img_vecs + center
+
+    # Also improvable as distance^2
+    center_dists = distance_array(center, coords)[0]
+    outside_ats = np.where(center_dists > min_compact_norm)[0]
+    # Let's make sure we're dealing with everyone in the primary cell
+    outside_pos = coords[outside_ats] = apply_PBC(coords[outside_ats],
+                                                  box, backend=backend)
+    # Potentially save cycles by re-filtering after PBC. It makes sense for
+    # GROMACS trajectories, that are saved as the rectangular cell.
+    center_dists = distance_array(center, outside_pos)[0]
+    outside_ats = outside_ats[center_dists > min_compact_norm]
+    outside_pos = coords[outside_ats]
+
+    # This works for the general case by checking the distance to all 27
+    # unit cell centers (primary + 1st neighbors). Will break for extremely
+    # skewed cases where the closest center is in a 2nd neighbor unit cell.
+    # The distance check to all 26 1st neighbors can also probably be
+    # optimized, both in general and with boxtype-specific shortcuts.
+    # Also improvable as distance^2
+    closest_img_ndx = distance_array(outside_pos, pbc_imgs).argmin(axis=1)
+    to_move = closest_img_ndx.nonzero()[0]
+    closest_img_ndx = closest_img_ndx[to_move]
+    outside_ats = outside_ats[to_move]
+    outside_pos = coords[outside_ats] - pbc_img_vecs[closest_img_ndx]
+    coords[outside_ats] = outside_pos
     return coords

package/MDAnalysis/lib/mdamath.py

@@ -58,7 +58,7 @@ def norm(v):
 
     .. math:: v = \sqrt{\mathbf{v}\cdot\mathbf{v}}
 
-    This version is faster then numpy.linalg.norm because it only works for a
+    This version is faster than numpy.linalg.norm because it only works for a
     single vector and therefore can skip a lot of the additional fuss
     linalg.norm does.
 

package/MDAnalysis/transformations/__init__.py

@@ -129,4 +129,4 @@ def wrapped(ts):
 from .rotate import rotateby
 from .positionaveraging import PositionAverager
 from .fit import fit_translation, fit_rot_trans
-from .wrap import wrap, unwrap
+from .wrap import wrap, unwrap, nojump