Add support for MDA minimize_vectors

.github/workflows/make_and_test.yml

@@ -96,7 +96,7 @@ jobs:
 
     - name: python_test
       # run python API tests
-      run: pytest distopia/tests
+      run: pytest -v -p no:logging distopia/tests
 
 
   external_hwy_no_test:

distopia/__init__.py

@@ -15,6 +15,8 @@
     self_distance_array_no_box,
     self_distance_array_ortho,
     self_distance_array_triclinic,
+    minimize_vectors_ortho,
+    minimize_vectors_triclinic,
 )
 
 __all__ = [
@@ -33,6 +35,8 @@
     'self_distance_array_no_box',
     'self_distance_array_ortho',
     'self_distance_array_triclinic',
+    'minimize_vectors_ortho',
+    'minimize_vectors_triclinic',
 ]
 
 from importlib.metadata import version

distopia/pydistopia.pyx

@@ -129,6 +129,19 @@ cdef extern from "distopia.h" namespace "distopia" nogil:
         const T *box,
         T *out,
     )
+    void MinimizeVectorsOrtho[T](
+        const T *vectors,
+        size_t n,
+        const T *box,
+        T *output,
+    )
+    void MinimizeVectorsTriclinic[T](
+        const T *vectors,
+        size_t n,
+        const T *box,
+        T *output,
+    )
+
 
 def get_n_float_lanes():
     """The number of floats per register distopia will handle on this system"""
@@ -902,4 +915,118 @@ def self_distance_array_triclinic(
                                &box[0][0],
                                &results_view[0])
 
-    return np.array(results)
+    return np.array(results)
+
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def minimize_vectors_ortho(
+    floating[:, ::1] vectors,
+    floating[::1] box,
+    floating[:, ::1] output=None):
+    """Apply minimum image convention to an array of vectors under orthorhombic boundary conditions
+
+    Parameters
+    ----------
+    vectors : float32 or float64 array
+        Vector array of shape ``(n, 3)``, either float32 or float64.  These
+        represent many vectors (such as between two particles).
+    box : float32 or float64 array
+        periodic boundary dimensions, in 3x3 format
+    output: float32 or float64 array
+        Same shape and dtype as input.  The vectors from the input, but
+        minimized according to the size of the box.
+
+    Returns
+    -------
+    minimized_vectors : np.array
+        Same shape and dtype as input.  The vectors from the input, but
+        minimized according to the size of the box.
+    """
+    cdef floating[:, ::1] output_view
+    cdef size_t nvals0 = vectors.shape[0]
+
+    cdef cnp.npy_intp[1] dims
+
+    cdef ssize_t final_size = nvals0 * (nvals0 -1) // 2
+
+    dims[0] = <ssize_t > final_size
+
+    # return early, will seg
+    if nvals0 == 0:
+        return np.array([])
+
+    if output is None:
+        # cnp.PyArray_NewLikeArray is not directly accessible
+        output = np.empty_like(vectors)
+    else:
+        if output.ndim != vectors.ndim:
+            raise ValueError("output dimensions not same as vectors")
+        if output.shape[0] != nvals0:
+            raise ValueError(f"output must be the same length as vectors ({nvals0}), you provided {output.shape[0]}")
+
+    output_view = output
+
+    MinimizeVectorsOrtho(&vectors[0][0],
+                            nvals0,
+                            &box[0],
+                            &output_view[0][0])
+
+    return np.array(output)
+
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def minimize_vectors_triclinic(
+    floating[:, ::1] vectors,
+    floating[:, ::1] box,
+    floating[:, ::1] output=None):
+    """Apply minimum image convention to an array of vectors under triclinic boundary conditions
+
+    Parameters
+    ----------
+    vectors : float32 or float64 array
+        Vector array of shape ``(n, 3)``, either float32 or float64.  These
+        represent many vectors (such as between two particles).
+    box : float32 or float64 array
+        periodic boundary dimensions, in 3x3 format
+    output: float32 or float64 array
+        Same shape and dtype as input.  The vectors from the input, but
+        minimized according to the size of the box.
+
+    Returns
+    -------
+    minimized_vectors : np.array
+        Same shape and dtype as input.  The vectors from the input, but
+        minimized according to the size of the box.
+    """
+    cdef floating[:, ::1] output_view
+    cdef size_t nvals0 = vectors.shape[0]
+
+    cdef cnp.npy_intp[1] dims
+
+    cdef ssize_t final_size = nvals0 * (nvals0 -1) // 2
+
+    dims[0] = <ssize_t > final_size
+
+    # return early, will seg
+    if nvals0 == 0:
+        return np.array([])
+
+    if output is None:
+        # cnp.PyArray_NewLikeArray is not directly accessible
+        output = np.empty_like(vectors)
+    else:
+        if output.ndim != vectors.ndim:
+            raise ValueError("output dimensions not same as vectors")
+        if output.shape[0] != nvals0:
+            raise ValueError(f"output must be the same length as vectors ({nvals0}), you provided {output.shape[0]}")
+
+    output_view = output
+
+    MinimizeVectorsTriclinic(&vectors[0][0],
+                                nvals0,
+                                &box[0][0],
+                                &output_view[0][0])
+
+    return np.array(output)

distopia/tests/test_distopia.py

@@ -1,3 +1,4 @@
+import itertools
 import pytest
 import numpy as np
 import distopia
@@ -479,3 +480,75 @@ def test_periodic_angles(self, box_bonds, positions_angles, dtype):
 
         for val in [test1, test2, test3]:
             assert_almost_equal(ref, val, self.prec, err_msg="Min image in angle calculation failed")
+
+    # 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_ortho(self, 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 = np.asarray([[10, 0, 0], [0, 10, 0], [0, 0, 10]], dtype=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 = np.asarray([10, 10, 10], dtype=dtype)
+        # empty input
+        res = distopia.minimize_vectors_ortho(np.array([[]], dtype=dtype), box2)
+        assert res.size == 0
+        # pass invalid output buffer size
+        output = np.array([[]]).astype(dtype)
+        with pytest.raises(ValueError, match="must be the same length"):
+            res = distopia.minimize_vectors_ortho(shifted_vec, box2, output)
+        # pass output buffer
+        output = np.empty_like(shifted_vec)
+        res = distopia.minimize_vectors_ortho(shifted_vec, box2, output)
+        assert_equal(res, output)
+        assert_allclose(res, vec, atol=0.00001)
+        assert res.dtype == dtype
+        # don't pass output buffer
+        res = distopia.minimize_vectors_ortho(shifted_vec, box2)
+        assert_allclose(res, vec, atol=0.00001)
+        assert res.dtype == dtype
+
+    # 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_triclinic(self, 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 = np.asarray([[10, 0, 0], [2, 10, 0], [2, 2, 10]], dtype=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 = np.asarray([[10, 0, 0, 2, 10, 0, 2, 2, 10]], dtype=dtype)
+        # empty input
+        res = distopia.minimize_vectors_triclinic(np.array([[]], dtype=dtype), box2)
+        assert res.size == 0
+        # pass invalid output buffer size
+        output = np.array([[]]).astype(dtype)
+        with pytest.raises(ValueError, match="must be the same length"):
+            res = distopia.minimize_vectors_triclinic(shifted_vec, box2, output)
+        # pass output buffer
+        output = np.empty_like(shifted_vec)
+        res = distopia.minimize_vectors_triclinic(shifted_vec, box2, output)
+        assert_equal(res, output)
+        assert_allclose(res, vec, atol=0.00001)
+        assert res.dtype == dtype
+        # don't pass output buffer
+        res = distopia.minimize_vectors_triclinic(shifted_vec, box2)
+        assert_allclose(res, vec, atol=0.00001)
+        assert res.dtype == dtype

docs/docs/api.md

@@ -46,4 +46,10 @@
     :docstring:
 
 ::: distopia.self_distance_array_triclinic
-    :docstring:
+    :docstring:
+
+::: distopia.minimize_vectors_ortho
+    :docstring:
+
+::: distopia.minimize_vectors_triclinic
+    :docstring:

docs/docs/examples.md

@@ -144,6 +144,41 @@ for i in range(N):
 ```
 
 
+## Minimizing Vectors
+
+### Orthorhombic periodic boundary conditions
+
+For orthorhombic boxes you need to provide the three box vector lengths `[l, l, l]`
+
+```python
+import numpy as np
+import distopia
+
+# make N x 3 coordinate arrays
+N = 10000
+coordinates0 = np.random.rand(3 * N).reshape(N, 3).astype(np.float32)
+coordinates1 = np.random.rand(3 * N).reshape(N, 3).astype(np.float32)
+box = np.asarray([10, 10, 10]).astype(np.float32)
+vectors = coordinates1 - coordinates0
+output = distopia.minimize_vectors_ortho(vectors, box)
+```
+
+### Triclinic periodic boundary conditions
+
+For triclinic boxes the box matrix must be provided in 3x3 matrix form.
+
+```python
+import numpy as np
+import distopia
+
+# make N x 3 coordinate arrays
+N = 10000
+coordinates0 = np.random.rand(3 * N).reshape(N, 3).astype(np.float32)
+coordinates1 = np.random.rand(3 * N).reshape(N, 3).astype(np.float32)
+box = np.asarray([[10, 0, 0], [0, 10, 0], [0, 0, 10]]).astype(np.float32)
+vectors = coordinates1 - coordinates0
+output = distopia.minimize_vectors_triclinic(vectors, box)
+```
 
 ## Questions
 

libdistopia/include/distopia.h

@@ -48,6 +48,8 @@ namespace distopia {
     template <typename T> void SelfDistanceArrayNoBoxIdx(const T *coords, const int *idx, int n, T *out);
     template <typename T> void SelfDistanceArrayOrthoIdx(const T *coords, const int *idx, int n, const T *box, T *out);
     template <typename T> void SelfDistanceArrayTriclinicIdx(const T *coords, const int *idx, int n, const T *box, T *out);
+    template <typename T> void MinimizeVectorsOrtho(const T *vectors, int n, const T *box, T *output);
+    template <typename T> void MinimizeVectorsTriclinic(const T *vectors, int n, const T *box, T *output);
     int GetNFloatLanes();
     int GetNDoubleLanes();
     std::vector<std::string> DistopiaSupportedAndGeneratedTargets();