Feature minimise vectors

package/CHANGELOG

@@ -34,6 +34,7 @@ Enhancements
   * Add option for custom compiler flags for C/C++ on build and remove march
     option from setup.cfg (Issue #3428, PR #3429).
   * Added R/S Chirality support to RDKit parsed molecules
+  * Added MDAnalysis.lib.distances.minimize_vectors
 
 Changes
   * TRZReader now defaults to a `dt` of 1.0 (instead of 0.0 previously) when

package/MDAnalysis/lib/c_distances.pyx

@@ -34,6 +34,9 @@ import numpy
 cimport numpy
 numpy.import_array()
 
+from libc.math cimport fabs, round as cround
+from libc.float cimport FLT_MAX, DBL_MAX
+
 cdef extern from "string.h":
     void* memcpy(void* dst, void* src, int len)
 
@@ -59,6 +62,7 @@ cdef extern from "calc_distances.h":
     void _ortho_pbc(coordinate* coords, int numcoords, float* box)
     void _triclinic_pbc(coordinate* coords, int numcoords, float* box)
     void minimum_image(double* x, float* box, float* inverse_box)
+    void minimum_image_triclinic(float* x, float* box, float* inverse_box)
 
 OPENMP_ENABLED = True if USED_OPENMP else False
 
@@ -268,3 +272,147 @@ def contact_matrix_pbc(coord, sparse_contacts, box, cutoff):
             if dist < cutoff2:
                 sparse_contacts[i, j] = True
                 sparse_contacts[j, i] = True
+
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+cdef inline void _minimum_image_orthogonal(cython.floating[:] dx,
+                                           cython.floating[:] box,
+                                           cython.floating[:] inverse_box) nogil:
+    """Minimize dx to be the shortest vector
+
+    Parameters
+    ----------
+    dx : numpy.array, shape (3,)
+      vector to minimize
+    box : numpy.array, shape (3,)
+      box length in each dimension
+    inverse_box : numpy.array, shape (3,)
+      inverse of box
+
+    Operates in-place on dx!
+    """
+    cdef int i, j
+    cdef cython.floating s
+
+    for i in range(3):
+        if box[i] > 0:
+            s = inverse_box[i] * dx[i]
+            dx[i] = box[i] * (s - cround(s))
+
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+cdef inline void _minimum_image_triclinic(cython.floating[:] dx,
+                                          cython.floating[:] box,
+                                          cython.floating[:] inverse_box) nogil:
+    """Minimise dx to be the shortest vector
+
+    Parameters
+    ----------
+    dx : numpy.array, shape (3,)
+      the vector to apply minimum image convention to
+    box : numpy.array, shape (9,)
+      flattened 3x3 representation of the unit cell
+    inverse_box : numpy.array, shape (3,)
+      inverse of the **diagonal** of the 3x3 representation
+
+    This version is near identical to the version in calc_distances.h, with the
+    difference being that this version does not assume that coordinates are at
+    most a single box length apart.
+
+    Operates in-place on dx!
+    """
+    cdef cython.floating dx_min[3]
+    cdef cython.floating s, dsq, dsq_min, rx
+    cdef cython.floating ry[2]
+    cdef cython.floating rz[3]
+    cdef int ix, iy, iz
+
+    # first make shift only 1 cell in any direction
+    s = cround(inverse_box[2] * dx[2])
+    dx[0] -= s * box[6]
+    dx[1] -= s * box[7]
+    dx[2] -= s * box[8]
+    s = cround(inverse_box[1] * dx[1])
+    dx[0] -= s * box[3]
+    dx[1] -= s * box[4]
+    s = cround(inverse_box[0] * dx[0])
+    dx[0] -= s * box[0]
+
+    if cython.floating is float:
+        dsq_min = FLT_MAX
+    else:
+        dsq_min = DBL_MAX
+    dx_min[0] = 0.0
+    dx_min[1] = 0.0
+    dx_min[2] = 0.0
+
+    # then check all images to see which combination of 1 cell shifts gives the best shift
+    for ix in range(-1, 2):
+        rx = dx[0] + box[0] * ix
+        for iy in range(-1, 2):
+            ry[0] = rx + box[3] * iy
+            ry[1] = dx[1] + box[4] * iy
+            for iz in range(-1, 2):
+                rz[0] = ry[0] + box[6] * iz
+                rz[1] = ry[1] + box[7] * iz
+                rz[2] = dx[2] + box[8] * iz
+                dsq = rz[0] * rz[0] + rz[1] * rz[1] + rz[2] * rz[2]
+                if (dsq < dsq_min):
+                    dsq_min = dsq
+                    dx_min[0] = rz[0]
+                    dx_min[1] = rz[1]
+                    dx_min[2] = rz[2]
+
+    dx[0] = dx_min[0]
+    dx[1] = dx_min[1]
+    dx[2] = dx_min[2]
+
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+@cython.cdivision(True)
+def _minimize_vectors_ortho(cython.floating[:, :] vectors not None, cython.floating[:] box not None,
+                            cython.floating[:, :] output not None):
+    cdef int i, n
+    cdef cython.floating box_inverse[3]
+    cdef cython.floating[:] box_inverse_view
+
+    box_inverse[0] = 1.0 / box[0]
+    box_inverse[1] = 1.0 / box[1]
+    box_inverse[2] = 1.0 / box[2]
+
+    box_inverse_view = box_inverse
+
+    n = len(vectors)
+    with nogil:
+        for i in range(n):
+            output[i, 0] = vectors[i, 0]
+            output[i, 1] = vectors[i, 1]
+            output[i, 2] = vectors[i, 2]
+            _minimum_image_orthogonal(output[i, :], box, box_inverse_view)
+
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def _minimize_vectors_triclinic(cython.floating[:, :] vectors not None, cython.floating[:] box not None,
+                                cython.floating[:, :] output not None):
+    cdef int i, n
+    cdef cython.floating box_inverse[3]
+    cdef cython.floating[:] box_inverse_view
+
+    # this is only inverse of diagonal, used for initial shift to ensure vector is within a single shift away
+    box_inverse[0] = 1.0 / box[0]
+    box_inverse[1] = 1.0 / box[4]
+    box_inverse[2] = 1.0 / box[8]
+
+    box_inverse_view = box_inverse
+
+    n = len(vectors)
+    with nogil:
+        for i in range(n):
+            output[i, 0] = vectors[i, 0]
+            output[i, 1] = vectors[i, 1]
+            output[i, 2] = vectors[i, 2]
+            _minimum_image_triclinic(output[i, :], box, box_inverse_view)

package/MDAnalysis/lib/distances.py

@@ -73,6 +73,8 @@
 from .mdamath import triclinic_vectors
 from ._augment import augment_coordinates, undo_augment
 from .nsgrid import FastNS
+from .c_distances import _minimize_vectors_ortho, _minimize_vectors_triclinic
+
 
 # hack to select backend with backend=<backend> kwarg. Note that
 # the cython parallel code (prange) in parallel.distances is
@@ -1549,3 +1551,45 @@ def apply_PBC(coords, box, backend="serial"):
         _run("triclinic_pbc", args=(coords, box), backend=backend)
 
     return coords
+
+
+@check_coords('vectors', enforce_copy=False, enforce_dtype=False)
+def minimize_vectors(vectors, box):
+    """Apply minimum image convention to an array of vectors
+
+    This function is required for calculating the correct vectors between two
+    points.  A naive approach of ``ag1.positions - ag2.positions`` will not
+    provide the minimum vectors between particles, even if all particles are
+    within the primary unit cell (box).
+
+    Parameters
+    ----------
+    vectors : numpy.ndarray
+        Vector array of shape ``(n, 3)``, either float32 or float64.  These
+        represent many vectors (such as between two particles).
+    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]``.
+
+    Returns
+    -------
+    minimized_vectors : numpy.ndarray
+        Same shape and dtype as input.  The vectors from the input, but
+        minimized according to the size of the box.
+
+    .. versionadded:: 2.1.0
+    """
+    boxtype, box = check_box(box)
+    output = np.empty_like(vectors)
+
+    # use box which is same precision as input vectors
+    box = box.astype(vectors.dtype)
+
+    if boxtype == 'ortho':
+        _minimize_vectors_ortho(vectors, box, output)
+    else:
+        _minimize_vectors_triclinic(vectors, box.ravel(), output)
+
+    return output

package/MDAnalysis/lib/util.py

@@ -1961,6 +1961,8 @@ def check_coords(*coord_names, **options):
         * **enforce_copy** (:class:`bool`, optional) -- Enforce working on a
           copy of the coordinate arrays. This is useful to ensure that the input
           arrays are left unchanged. Default: ``True``
+        * **enforce_dtype** (:class:`bool`, optional) -- Enforce a conversion
+          to float32.  Default: ``True``
         * **allow_single** (:class:`bool`, optional) -- Allow the input
           coordinate array to be a single coordinate with shape ``(3,)``.
         * **convert_single** (:class:`bool`, optional) -- If ``True``, single
@@ -2016,6 +2018,7 @@ def check_coords(*coord_names, **options):
     .. versionadded:: 0.19.0
     """
     enforce_copy = options.get('enforce_copy', True)
+    enforce_dtype = options.get('enforce_dtype', True)
     allow_single = options.get('allow_single', True)
     convert_single = options.get('convert_single', True)
     reduce_result_if_single = options.get('reduce_result_if_single', True)
@@ -2061,13 +2064,14 @@ def _check_coords(coords, argname):
                 if (coords.ndim != 2) or (coords.shape[1] != 3):
                     raise ValueError("{}(): {}.shape must be (n, 3), got {}."
                                      "".format(fname, argname, coords.shape))
-            try:
-                coords = coords.astype(
-                    np.float32, order='C', copy=enforce_copy)
-            except ValueError:
-                errmsg = (f"{fname}(): {argname}.dtype must be convertible to "
-                          f"float32, got {coords.dtype}.")
-                raise TypeError(errmsg) from None
+            if enforce_dtype:
+                try:
+                    coords = coords.astype(
+                        np.float32, order='C', copy=enforce_copy)
+                except ValueError:
+                    errmsg = (f"{fname}(): {argname}.dtype must be convertible to "
+                              f"float32, got {coords.dtype}.")
+                    raise TypeError(errmsg) from None
             return coords, is_single
 
         @wraps(func)

testsuite/MDAnalysisTests/lib/test_distances.py

@@ -23,6 +23,7 @@
 import pytest
 import numpy as np
 from numpy.testing import assert_equal, assert_almost_equal
+import itertools
 from itertools import combinations_with_replacement as comb
 
 import MDAnalysis
@@ -1376,3 +1377,27 @@ def test_wront_backend(self, backend_selection_pos):
 
 def test_used_openmpflag():
     assert isinstance(distances.USED_OPENMP, bool)
+
+
+# test both orthognal and triclinic boxes
+@pytest.mark.parametrize('box', (np.eye(3) * 10, np.array([[10, 0, 0], [2, 10, 0], [2, 2, 10]])))
+# try shifts of -2 to +2 in each dimension, and all combinations of shifts
+@pytest.mark.parametrize('shift', itertools.product(range(-2, 3), range(-2, 3), range(-2, 3)))
+@pytest.mark.parametrize('dtype', (np.float32, np.float64))
+def test_minimize_vectors(box, shift, dtype):
+    # test vectors pointing in all directions
+    # these currently all obey minimum convention as they're much smaller than the box
+    vec = np.array(list(itertools.product(range(-1, 2), range(-1, 2), range(-1, 2))), dtype=dtype)
+    box = box.astype(dtype)
+
+    # box is 3x3 representation
+    # multiply by shift, then sum xyz components then add these to the vector
+    # this technically doesn't alter the vector because of periodic boundaries
+    shifted_vec = (vec + (box.T * shift).sum(axis=1)).astype(dtype)
+
+    box2 = mdamath.triclinic_box(*box).astype(dtype)
+
+    res = distances.minimize_vectors(shifted_vec, box2)
+
+    assert_almost_equal(res, vec, decimal=5)
+    assert res.dtype == dtype