[WIP] Idx python

distopia/__init__.py

@@ -16,6 +16,9 @@
     calc_self_distance_array_no_box,
     calc_self_distance_array_ortho,
     calc_self_distance_array_triclinic,
+    calc_bonds_no_box_idx,
+    calc_bonds_ortho_idx,
+    calc_bonds_triclinic_idx,
 )
 
 __all__ = [
@@ -34,5 +37,8 @@
     'calc_self_distance_array_no_box',
     'calc_self_distance_array_ortho',
     'calc_self_distance_array_triclinic',
+    'calc_bonds_no_box_idx',
+    'calc_bonds_ortho_idx',
+    'calc_bonds_triclinic_idx',
     '__version__',
 ]

distopia/_distopia.pyx

@@ -13,7 +13,7 @@ The python functions for distopia
 import numpy as np
 cimport cython
 cimport numpy as cnp
-from cython cimport floating
+from cython cimport floating, integral
 cnp.import_array()
 
 cdef extern from "distopia.h" namespace "distopia" nogil:
@@ -129,6 +129,95 @@ cdef extern from "distopia.h" namespace "distopia" nogil:
         const T *box,
         T *out,
     )
+    void CalcBondsNoBoxIdx[T](
+        const T *coords,
+        const int* a_idx,
+        const int* b_idx,
+        size_t n,
+        T *out,
+    )
+
+    void CalcBondsOrthoIdx[T](
+        const T *coords,
+        const int* a_idx,
+        const int* b_idx,
+        size_t n,
+        const T *box,
+        T *out,
+    )
+
+    void CalcBondsTriclinicIdx[T](
+        const T *coords,
+        const int* a_idx,
+        const int* b_idx,
+        size_t n,
+        const T *box,
+        T *out,
+    )
+
+    void CalcAnglesNoBoxIdx[T](
+        const T *coords,
+        const int* a_idx,
+        const int* b_idx,
+        const int* c_idx,
+        size_t n,
+        T *out,
+    )
+
+    void CalcAnglesOrthoIdx[T](
+        const T *coords,
+        const int* a_idx,
+        const int* b_idx,
+        const int* c_idx,
+        size_t n,
+        const T *box,
+        T *out,
+    )
+
+    void CalcAnglesTriclinicIdx[T](
+        const T *coords,
+        const int* a_idx,
+        const int* b_idx,
+        const int* c_idx,
+        size_t n,
+        const T *box,
+        T *out,
+    )
+
+    void CalcDihedralsNoBoxIdx[T](
+        const T *coords,
+        const int* a_idx,
+        const int* b_idx,
+        const int* c_idx,
+        const int* d_idx,
+        size_t n,
+        T *out,
+    )
+
+    void CalcDihedralsOrthoIdx[T](
+        const T *coords,
+        const int* a_idx,
+        const int* b_idx,
+        const int* c_idx,
+        const int* d_idx,
+        size_t n,
+        const T *box,
+        T *out,
+    )
+
+
+    void CalcDihedralsTriclinicIdx[T](
+        const T *coords,
+        const int* a_idx,
+        const int* b_idx,
+        const int* c_idx,
+        const int* d_idx,
+        size_t n,
+        const T *box,
+        T *out,
+    )
+
+
 
 def get_n_float_lanes():
     """The number of floats per register distopia will handle on this system"""
@@ -890,4 +979,152 @@ def calc_self_distance_array_triclinic(
                                &box[0][0],
                                &results_view[0])
 
-    return np.array(results).reshape(coords0.shape[0], coords0.shape[0])
+    return np.array(results).reshape(coords0.shape[0], coords0.shape[0])
+
+
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def calc_bonds_no_box_idx(
+     floating[:, ::1] coords,
+     int[::1] a_idx,
+     int[::1] b_idx,
+     floating[::1] results=None):
+    """Calculate pairwise distances between coords[a_idx] and coords[b_idx] with no periodic boundary conditions
+
+    Parameters
+    ----------
+    coords : float32 or float64 array
+      must be same length and dtype
+    a_idx, b_idx : int array
+      must be same length and dtype
+    results: float32 or float64 array (optional)
+      array to store results in, must be same size and dtype as a_idx/b_idx
+
+    Returns
+    -------
+    distances : float32 or float64 array
+      same length and dtype as a_idx/b_idx
+    """
+    cdef floating[::1] results_view
+    cdef size_t nvals = a_idx.shape[0]
+    cdef cnp.npy_intp[1] dims
+    dims[0] = <ssize_t > nvals  # FIXME truncation?
+
+    _check_shapes(a_idx, b_idx)
+
+    if results is None:
+        if floating is float:
+            results = cnp.PyArray_EMPTY(1, dims, cnp.NPY_FLOAT32, 0)
+        else:
+            results = cnp.PyArray_EMPTY(1, dims, cnp.NPY_FLOAT64, 0)
+
+    else:
+        _check_results(results, nvals)
+
+    results_view = results
+
+    CalcBondsNoBoxIdx(& coords[0][0], & a_idx[0], & b_idx[0], nvals, & results_view[0])
+
+    return np.array(results)
+
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def calc_bonds_ortho_idx(
+     floating[:, ::1] coords,
+     int[::1] a_idx,
+     int[::1] b_idx,
+     floating[::1] box,
+     floating[::1] results=None):
+    """Calculate pairwise distances between coords[a_idx] and coords[b_idx] under orthorhombic boundary conditions
+
+    Parameters
+    ----------
+    coords : float32 or float64 array
+      must be same length and dtype
+    a_idx, b_idx : int array
+      must be same length and dtype
+    box : float32 or float64 array
+      orthorhombic periodic boundary dimensions in [L, L, L] format
+    results: float32 or float64 array (optional)
+      array to store results in, must be same size and dtype as a_idx/b_idx
+
+    Returns
+    -------
+    distances : float32 or float64 array
+      same length and dtype as a_idx/b_idx
+    """
+    cdef floating[::1] results_view
+    cdef size_t nvals = a_idx.shape[0]
+    cdef cnp.npy_intp[1] dims
+    dims[0] = <ssize_t > nvals  # FIXME truncation?
+
+    _check_shapes(a_idx, b_idx)
+
+
+    if results is None:
+        if floating is float:
+            results = cnp.PyArray_EMPTY(1, dims, cnp.NPY_FLOAT32, 0)
+        else:
+            results = cnp.PyArray_EMPTY(1, dims, cnp.NPY_FLOAT64, 0)
+
+    else:
+        _check_results(results, nvals)
+
+    results_view = results
+
+    CalcBondsOrthoIdx(& coords[0][0], & a_idx[0], & b_idx[0], nvals, & box[0], & results_view[0])
+
+    return np.array(results)
+
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def calc_bonds_triclinic_idx(
+     floating[:, ::1] coords,
+     int[::1] a_idx,
+     int[::1] b_idx,
+     floating[:, ::1] box,
+     floating[::1] results=None):
+    """Calculate pairwise distances between coords[a_idx] and coords[b_idx] under triclinic boundary conditions
+
+    Parameters
+    ----------
+    coords : float32 or float64 array
+      must be same length and dtype
+    a_idx, b_idx : int array
+      must be same length and dtype
+    box : float32 or float64 array
+      periodic boundary dimensions, in 3x3 format
+    results: float32 or float64 array (optional)
+      array to store results in, must be same size and dtype as a_idx/b_idx
+
+    Returns
+    -------
+    distances : float32 or float64 array
+      same length and dtype as a_idx/b_idx
+    """
+    cdef floating[::1] results_view
+    cdef size_t nvals = a_idx.shape[0]
+    cdef cnp.npy_intp[1] dims
+    dims[0] = <ssize_t > nvals  # FIXME truncation?
+
+    _check_shapes(a_idx, b_idx)
+
+
+    if results is None:
+        if floating is float:
+            results = cnp.PyArray_EMPTY(1, dims, cnp.NPY_FLOAT32, 0)
+        else:
+            results = cnp.PyArray_EMPTY(1, dims, cnp.NPY_FLOAT64, 0)
+
+    else:
+        _check_results(results, nvals)
+
+    results_view = results
+
+    CalcBondsTriclinicIdx(& coords[0][0], & a_idx[0], & b_idx[0], nvals, & box[0][0], & results_view[0])
+
+    return np.array(results)
+

distopia/tests/test_distopia.py

@@ -240,6 +240,36 @@ def test_triclinic_bad_result(self):
             distopia.calc_self_distance_array_triclinic(c0, box, results=np.empty(1, dtype=np.float32))
 
 
+
+class TestDistancesIdx:
+    def arange_input(self, N, dtype):
+        return np.arange(3 * N, dtype=dtype).reshape(N, 3)
+
+    def result_shim(self, make_arr, N, dtype):
+        if not make_arr:
+            return None
+        else:
+            return np.empty(N, dtype=dtype)
+
+
+    @pytest.mark.parametrize("dtype", (np.float32, np.float64))
+    @pytest.mark.parametrize("N", (3, 6, 10, 1000, 10000))
+    @pytest.mark.parametrize("use_result_buffer", (True, False))
+    def test_calc_bonds_no_box_idx(self, N, use_result_buffer, dtype):
+        c = self.arange_input(N, dtype)
+
+        a_idx = np.asarray(np.arange(N -3)).astype(np.int32)
+        b_idx = np.asarray(np.arange(N -3)).astype(np.int32)
+        result_buffer = self.result_shim(use_result_buffer, len(a_idx), dtype)
+        result = distopia.calc_bonds_no_box_idx(
+            c, a_idx, b_idx, results=result_buffer
+        )
+        assert_allclose(result, np.zeros(N))
+        # check dtype of result
+        assert result.dtype == dtype
+
+
+
 class TestMDA:
     """
     Copy of some of the tests from MDAnalysisTests repo
@@ -381,4 +411,8 @@ def test_periodic_angles(self, box_bonds, positions_angles, dtype):
         test3 = distopia.calc_angles_ortho(a, b, c2, box=box)
 
         for val in [test1, test2, test3]:
-            assert_almost_equal(ref, val, self.prec, err_msg="Min image in angle calculation failed")
+            assert_almost_equal(ref, val, self.prec, err_msg="Min image in angle calculation failed")
+
+
+
+

libdistopia/CMakeLists.txt

@@ -41,10 +41,14 @@ Include(GoogleTest)
 add_subdirectory("googletest")
 enable_testing()
 include_directories(${gtest_SOURCE_DIR}/include ${gtest_SOURCE_DIR})
+include_directories(${gmock_SOURCE_DIR}/include ${gmock_SOURCE_DIR})
+
 
 add_executable(test)
 target_sources(test PRIVATE "test/test.cpp")
 target_link_libraries(test PUBLIC gtest gtest_main)
+target_link_libraries(test PUBLIC gmock gmock_main)
+
 target_link_libraries(test PUBLIC libdistopia)
 target_include_directories(test PUBLIC  ${CMAKE_SOURCE_DIR})
 gtest_discover_tests(test WORKING_DIRECTORY ${CMAKE_RUNTIME_OUTPUT_DIRECTORY})
@@ -58,6 +62,7 @@ target_include_directories(targets PUBLIC  ${CMAKE_SOURCE_DIR})
 add_executable(test_mda_match)
 target_sources(test_mda_match PRIVATE "test/test_mda_match.cpp")
 target_link_libraries(test_mda_match PUBLIC gtest gtest_main)
+target_link_libraries(test PUBLIC gmock gmock_main)
 target_link_libraries(test_mda_match PUBLIC libdistopia)
 target_include_directories(test_mda_match PUBLIC  ${CMAKE_SOURCE_DIR})
 gtest_discover_tests(test_mda_match WORKING_DIRECTORY ${CMAKE_RUNTIME_OUTPUT_DIRECTORY})