diff --git a/distopia/_distopia.pyx b/distopia/_distopia.pyx index cbf37c12..f796f8bd 100644 --- a/distopia/_distopia.pyx +++ b/distopia/_distopia.pyx @@ -39,17 +39,20 @@ cdef extern from "distopia.h" nogil: @cython.boundscheck(False) @cython.wraparound(False) -cpdef cnp.ndarray[cnp.float32_t, ndim=1] calc_bonds_ortho_float(float[:, ::1] coords0, - float[:, ::1] coords1, - float[::1] box): +cpdef cnp.ndarray[cnp.float32_t, ndim = 1] calc_bonds_ortho_float(float[:, ::1] coords0, + float[:, ::1] coords1, + float[::1] box, + cnp.ndarray results=None): """Calculate pairwise distances between coords0 and coords1 Parameters ---------- coords0, coords1 : float32 array must be same length - box : float32 array + box : float32 array periodic boundary dimensions + results: float32 array (optional) + array to store results in, must be same size as coords0/coords1 Returns ------- @@ -60,7 +63,8 @@ cpdef cnp.ndarray[cnp.float32_t, ndim=1] calc_bonds_ortho_float(float[:, ::1] co cdef size_t nvals = coords0.shape[0] cdef cnp.npy_intp[1] dims dims[0] = nvals # FIXME truncation? - cdef cnp.ndarray[cnp.float32_t, ndim=1] results = cnp.PyArray_EMPTY(1, dims, cnp.NPY_FLOAT32, 0) + if results is None: + results = cnp.PyArray_EMPTY(1, dims, cnp.NPY_FLOAT32, 0) results_view = results @@ -71,9 +75,10 @@ cpdef cnp.ndarray[cnp.float32_t, ndim=1] calc_bonds_ortho_float(float[:, ::1] co @cython.boundscheck(False) @cython.wraparound(False) -cpdef cnp.ndarray[cnp.float64_t, ndim=1] calc_bonds_ortho_double(double[:, ::1] coords0, +cpdef cnp.ndarray[cnp.float64_t, ndim = 1] calc_bonds_ortho_double(double[:, ::1] coords0, double[:, ::1] coords1, - double[::1] box): + double[::1] box, + cnp.ndarray results=None): """Calculate pairwise distances between coords0 and coords1 Parameters @@ -82,6 +87,8 @@ cpdef cnp.ndarray[cnp.float64_t, ndim=1] calc_bonds_ortho_double(double[:, ::1] must be same length box : float64 array periodic boundary dimensions + results: float64 array (optional) + array to store results in, must be same size as coords0/coords1 Returns ------- @@ -92,7 +99,8 @@ cpdef cnp.ndarray[cnp.float64_t, ndim=1] calc_bonds_ortho_double(double[:, ::1] cdef size_t nvals = coords0.shape[0] cdef cnp.npy_intp[1] dims dims[0] = nvals # FIXME truncation? - cdef cnp.ndarray[cnp.float64_t, ndim=1] results = cnp.PyArray_EMPTY(1, dims, cnp.NPY_FLOAT64, 0) + if results is None: + results = cnp.PyArray_EMPTY(1, dims, cnp.NPY_FLOAT64, 0) results_view = results @@ -103,14 +111,17 @@ cpdef cnp.ndarray[cnp.float64_t, ndim=1] calc_bonds_ortho_double(double[:, ::1] @cython.boundscheck(False) @cython.wraparound(False) -cpdef cnp.ndarray[cnp.float32_t, ndim=1] calc_bonds_no_box_float(float[:, ::1] coords0, - float[:, ::1] coords1): +cpdef cnp.ndarray[cnp.float32_t, ndim = 1] calc_bonds_no_box_float(float[:, ::1] coords0, + float[:, ::1] coords1, + cnp.ndarray results=None): """Calculate pairwise distances between coords0 and coords1 Parameters ---------- coords0, coords1 : float32 array must be same length + results: float32 array (optional) + array to store results in, must be same size as coords0/coords1 Returns ------- @@ -121,7 +132,8 @@ cpdef cnp.ndarray[cnp.float32_t, ndim=1] calc_bonds_no_box_float(float[:, ::1] c cdef size_t nvals = coords0.shape[0] cdef cnp.npy_intp[1] dims dims[0] = nvals # FIXME truncation? - cdef cnp.ndarray[cnp.float32_t, ndim=1] results = cnp.PyArray_EMPTY(1, dims, cnp.NPY_FLOAT32, 0) + if results is None: + results = cnp.PyArray_EMPTY(1, dims, cnp.NPY_FLOAT32, 0) results_view = results @@ -132,8 +144,9 @@ cpdef cnp.ndarray[cnp.float32_t, ndim=1] calc_bonds_no_box_float(float[:, ::1] c @cython.boundscheck(False) @cython.wraparound(False) -cpdef cnp.ndarray[cnp.float64_t, ndim=1] calc_bonds_no_box_double(double[:, ::1] coords0, - double[:, ::1] coords1): +cpdef cnp.ndarray[cnp.float64_t, ndim = 1] calc_bonds_no_box_double(double[:, ::1] coords0, + double[:, ::1] coords1, + cnp.ndarray results=None): """Calculate pairwise distances between coords0 and coords1 Parameters @@ -142,6 +155,8 @@ cpdef cnp.ndarray[cnp.float64_t, ndim=1] calc_bonds_no_box_double(double[:, ::1] must be same length box : float64 array periodic boundary dimensions + results: float64 array (optional) + array to store results in, must be same size as coords0/coords1 Returns ------- @@ -152,7 +167,8 @@ cpdef cnp.ndarray[cnp.float64_t, ndim=1] calc_bonds_no_box_double(double[:, ::1] cdef size_t nvals = coords0.shape[0] cdef cnp.npy_intp[1] dims dims[0] = nvals # FIXME truncation? - cdef cnp.ndarray[cnp.float64_t, ndim=1] results = cnp.PyArray_EMPTY(1, dims, cnp.NPY_FLOAT64, 0) + if results is None: + results = cnp.PyArray_EMPTY(1, dims, cnp.NPY_FLOAT64, 0) results_view = results @@ -163,9 +179,10 @@ cpdef cnp.ndarray[cnp.float64_t, ndim=1] calc_bonds_no_box_double(double[:, ::1] @cython.boundscheck(False) @cython.wraparound(False) -cpdef cnp.ndarray[cnp.float32_t, ndim=1] calc_bonds_idx_ortho_float(float[:, ::1] coords, - size_t[::1] idx, - float[::1] box): +cpdef cnp.ndarray[cnp.float32_t, ndim = 1] calc_bonds_idx_ortho_float(float[:, ::1] coords, + size_t[::1] idx, + float[::1] box, + cnp.ndarray results=None): """Calculate distances between pairs of coordinates by index Parameters @@ -176,6 +193,8 @@ cpdef cnp.ndarray[cnp.float32_t, ndim=1] calc_bonds_idx_ortho_float(float[:, ::1 array of integers to calculate distances for box : float32 array periodic boundary dimensions + results: float32 array (optional) + array to store results in, must be half the size of idx Returns ------- @@ -186,7 +205,8 @@ cpdef cnp.ndarray[cnp.float32_t, ndim=1] calc_bonds_idx_ortho_float(float[:, ::1 cdef size_t nvals = idx.shape[0] // 2 # SAFE? cdef cnp.npy_intp[1] dims dims[0] = nvals # FIXME truncation? - cdef cnp.ndarray[cnp.float32_t, ndim=1] results = cnp.PyArray_EMPTY(1, dims, cnp.NPY_FLOAT32, 0) + if results is None: + results = cnp.PyArray_EMPTY(1, dims, cnp.NPY_FLOAT32, 0) results_view = results @@ -197,9 +217,10 @@ cpdef cnp.ndarray[cnp.float32_t, ndim=1] calc_bonds_idx_ortho_float(float[:, ::1 @cython.boundscheck(False) @cython.wraparound(False) -cpdef cnp.ndarray[cnp.float64_t, ndim=1] calc_bonds_idx_ortho_double(double[:, ::1] coords, - size_t[::1] idx, - double[::1] box): +cpdef cnp.ndarray[cnp.float64_t, ndim = 1] calc_bonds_idx_ortho_double(double[:, ::1] coords, + size_t[::1] idx, + double[::1] box, + cnp.ndarray results=None): """Calculate distances between pairs of coordinates by index Parameters @@ -210,6 +231,8 @@ cpdef cnp.ndarray[cnp.float64_t, ndim=1] calc_bonds_idx_ortho_double(double[:, : array of integers to calculate distances for box : float64 array periodic boundary dimensions + results: float64 array (optional) + array to store results in, must be half the size of idx Returns ------- @@ -220,7 +243,8 @@ cpdef cnp.ndarray[cnp.float64_t, ndim=1] calc_bonds_idx_ortho_double(double[:, : cdef size_t nvals = idx.shape[0] // 2 # SAFE? cdef cnp.npy_intp[1] dims dims[0] = nvals # FIXME truncation? - cdef cnp.ndarray[cnp.float64_t, ndim=1] results = cnp.PyArray_EMPTY(1, dims, cnp.NPY_FLOAT64, 0) + if results is None: + results = cnp.PyArray_EMPTY(1, dims, cnp.NPY_FLOAT64, 0) results_view = results @@ -231,8 +255,9 @@ cpdef cnp.ndarray[cnp.float64_t, ndim=1] calc_bonds_idx_ortho_double(double[:, : @cython.boundscheck(False) @cython.wraparound(False) -cpdef cnp.ndarray[cnp.float32_t, ndim=1] calc_bonds_idx_no_box_float(float[:, ::1] coords, - size_t[::1] idx): +cpdef cnp.ndarray[cnp.float32_t, ndim = 1] calc_bonds_idx_no_box_float(float[:, ::1] coords, + size_t[::1] idx, + cnp.ndarray results=None): """Calculate distances between pairs of coordinates by index Parameters @@ -241,7 +266,9 @@ cpdef cnp.ndarray[cnp.float32_t, ndim=1] calc_bonds_idx_no_box_float(float[:, :: array of coordinates idx: int array array of integers to calculate distances for - + results: float32 array (optional) + array to store results in, must be half the size of idx + Returns ------- distances : float32 array @@ -251,7 +278,8 @@ cpdef cnp.ndarray[cnp.float32_t, ndim=1] calc_bonds_idx_no_box_float(float[:, :: cdef size_t nvals = idx.shape[0] // 2 # SAFE? cdef cnp.npy_intp[1] dims dims[0] = nvals # FIXME truncation? - cdef cnp.ndarray[cnp.float32_t, ndim=1] results = cnp.PyArray_EMPTY(1, dims, cnp.NPY_FLOAT32, 0) + if results is None: + results = cnp.PyArray_EMPTY(1, dims, cnp.NPY_FLOAT32, 0) results_view = results @@ -262,8 +290,9 @@ cpdef cnp.ndarray[cnp.float32_t, ndim=1] calc_bonds_idx_no_box_float(float[:, :: @cython.boundscheck(False) @cython.wraparound(False) -cpdef cnp.ndarray[cnp.float64_t, ndim=1] calc_bonds_idx_no_box_double(double[:, ::1] coords, - size_t[::1] idx): +cpdef cnp.ndarray[cnp.float64_t, ndim = 1] calc_bonds_idx_no_box_double(double[:, ::1] coords, + size_t[::1] idx, + cnp.ndarray results=None): """Calculate distances between pairs of coordinates by index Parameters @@ -272,6 +301,8 @@ cpdef cnp.ndarray[cnp.float64_t, ndim=1] calc_bonds_idx_no_box_double(double[:, array of coordinates idx: int array array of integers to calculate distances for + results: float64 array (optional) + array to store results in, must be half the size of idx Returns ------- @@ -282,7 +313,8 @@ cpdef cnp.ndarray[cnp.float64_t, ndim=1] calc_bonds_idx_no_box_double(double[:, cdef size_t nvals = idx.shape[0] // 2 # SAFE? cdef cnp.npy_intp[1] dims dims[0] = nvals # FIXME truncation? - cdef cnp.ndarray[cnp.float64_t, ndim=1] results = cnp.PyArray_EMPTY(1, dims, cnp.NPY_FLOAT64, 0) + if results is None: + results = cnp.PyArray_EMPTY(1, dims, cnp.NPY_FLOAT64, 0) results_view = results diff --git a/distopia/tests/test_distopia.py b/distopia/tests/test_distopia.py index 08bb39da..80786227 100644 --- a/distopia/tests/test_distopia.py +++ b/distopia/tests/test_distopia.py @@ -11,84 +11,110 @@ class TestDistances: - def arange_input(self, N, dtype): - return np.arange(3*N, dtype=dtype).reshape(N,3) + return np.arange(3 * N, dtype=dtype).reshape(N, 3) def idx_input(self, N): idxs = np.zeros((N, 2), dtype=np.ulonglong) for i in range(N): - idxs[i,0]=i - idxs[i,1]=i+N + idxs[i, 0] = i + idxs[i, 1] = i + N return np.ravel(idxs) - @pytest.mark.parametrize('box', ([10, 10, 10], [100, 20, 10])) - @pytest.mark.parametrize('N', (0, 10, 1000, 10000)) - def test_calc_bonds_ortho_float_all_zero(self, N, box): + def result_shim(self, make_arr, N, dtype): + if not make_arr: + return None + else: + return np.empty(N, dtype=dtype) + + @pytest.mark.parametrize("box", ([10, 10, 10], [100, 20, 10])) + @pytest.mark.parametrize("N", (0, 10, 1000, 10000)) + @pytest.mark.parametrize("use_result_buffer", (True, False)) + def test_calc_bonds_ortho_float_all_zero(self, N, box, use_result_buffer): c0 = self.arange_input(N, np.float32) c1 = self.arange_input(N, np.float32) + result_buffer = self.result_shim(use_result_buffer, N, np.float32) result = distopia.calc_bonds_ortho_float( - c0, c1, np.asarray(box, dtype=np.float32)) + c0, c1, np.asarray(box, dtype=np.float32), results=result_buffer + ) assert_allclose(result, np.zeros(N)) - @pytest.mark.parametrize('box', ([10, 10, 10], [100, 20, 10])) - @pytest.mark.parametrize('N', (0, 10, 1000, 10000)) - def test_calc_bonds_ortho_double_all_zero(self, N, box): + @pytest.mark.parametrize("box", ([10, 10, 10], [100, 20, 10])) + @pytest.mark.parametrize("N", (0, 10, 1000, 10000)) + @pytest.mark.parametrize("use_result_buffer", (True, False)) + def test_calc_bonds_ortho_double_all_zero(self, N, box, use_result_buffer): c0 = self.arange_input(N, np.float64) c1 = self.arange_input(N, np.float64) + result_buffer = self.result_shim(use_result_buffer, N, np.float64) result = distopia.calc_bonds_ortho_double( - c0, c1, np.asarray(box, dtype=np.float64)) + c0, c1, np.asarray(box, dtype=np.float64), results=result_buffer + ) assert_allclose(result, np.zeros(N)) - @pytest.mark.parametrize('N', (0, 10, 1000, 10000)) - def test_calc_bonds_nobox_float_all_zero(self, N): + @pytest.mark.parametrize("N", (0, 10, 1000, 10000)) + @pytest.mark.parametrize("use_result_buffer", (True, False)) + def test_calc_bonds_nobox_float_all_zero(self, N, use_result_buffer): c0 = self.arange_input(N, np.float32) c1 = self.arange_input(N, np.float32) - result = distopia.calc_bonds_no_box_float(c0, c1) + result_buffer = self.result_shim(use_result_buffer, N, np.float32) + result = distopia.calc_bonds_no_box_float(c0, c1, results=result_buffer) assert_allclose(result, np.zeros(N)) - - @pytest.mark.parametrize('N', (0, 10, 1000, 10000)) - def test_calc_bonds_nobox_double_all_zero(self, N): + + @pytest.mark.parametrize("N", (0, 10, 1000, 10000)) + @pytest.mark.parametrize("use_result_buffer", (True, False)) + def test_calc_bonds_nobox_double_all_zero(self, N, use_result_buffer): c0 = self.arange_input(N, np.float64) c1 = self.arange_input(N, np.float64) - result = distopia.calc_bonds_no_box_double(c0, c1) + result_buffer = self.result_shim(use_result_buffer, N, np.float64) + result = distopia.calc_bonds_no_box_double(c0, c1, results=result_buffer) assert_allclose(result, np.zeros(N)) - @pytest.mark.parametrize('box', ([10, 10, 10], [100, 20, 10])) - @pytest.mark.parametrize('N', (0, 10, 1000, 10000)) - def test_calc_bonds_idx_ortho_float_all_zero(self, N, box): + @pytest.mark.parametrize("box", ([10, 10, 10], [100, 20, 10])) + @pytest.mark.parametrize("N", (0, 10, 1000, 10000)) + @pytest.mark.parametrize("use_result_buffer", (True, False)) + def test_calc_bonds_idx_ortho_float_all_zero(self, N, box, use_result_buffer): c0 = self.arange_input(N, np.float32) c1 = self.arange_input(N, np.float32) idx = self.idx_input(N) + result_buffer = self.result_shim(use_result_buffer, N, np.float32) coords = np.vstack((c0, c1)) - result = distopia.calc_bonds_idx_ortho_float(coords, idx, np.asarray(box).astype(np.float32)) + result = distopia.calc_bonds_idx_ortho_float( + coords, idx, np.asarray(box).astype(np.float32), results=result_buffer + ) assert_allclose(result, np.zeros(N)) - @pytest.mark.parametrize('box', ([10, 10, 10], [100, 20, 10])) - @pytest.mark.parametrize('N', (0, 10, 1000, 10000)) - def test_calc_bonds_idx_ortho_double_all_zero(self, N, box): + @pytest.mark.parametrize("box", ([10, 10, 10], [100, 20, 10])) + @pytest.mark.parametrize("N", (0, 10, 1000, 10000)) + @pytest.mark.parametrize("use_result_buffer", (True, False)) + def test_calc_bonds_idx_ortho_double_all_zero(self, N, box, use_result_buffer): c0 = self.arange_input(N, np.float64) c1 = self.arange_input(N, np.float64) idx = self.idx_input(N) + result_buffer = self.result_shim(use_result_buffer, N, np.float64) coords = np.vstack((c0, c1)) - result = distopia.calc_bonds_idx_ortho_double(coords, idx, np.asarray(box).astype(np.float64)) + result = distopia.calc_bonds_idx_ortho_double( + coords, idx, np.asarray(box).astype(np.float64), results=result_buffer + ) assert_allclose(result, np.zeros(N)) - - @pytest.mark.parametrize('N', (0, 10, 1000, 10000)) - def test_calc_bonds_idx_no_box_float_all_zero(self, N): + @pytest.mark.parametrize("N", (0, 10, 1000, 10000)) + @pytest.mark.parametrize("use_result_buffer", (True, False)) + def test_calc_bonds_idx_no_box_float_all_zero(self, N, use_result_buffer): c0 = self.arange_input(N, np.float32) c1 = self.arange_input(N, np.float32) idx = self.idx_input(N) + result_buffer = self.result_shim(use_result_buffer, N, np.float32) coords = np.vstack((c0, c1)) - result = distopia.calc_bonds_idx_no_box_float(coords, idx) + result = distopia.calc_bonds_idx_no_box_float(coords, idx, results=result_buffer) assert_allclose(result, np.zeros(N)) - @pytest.mark.parametrize('N', (0, 10, 1000, 10000)) - def test_calc_bonds_idx_no_box_double_all_zero(self, N): + @pytest.mark.parametrize("N", (0, 10, 1000, 10000)) + @pytest.mark.parametrize("use_result_buffer", (True, False)) + def test_calc_bonds_idx_no_box_double_all_zero(self, N, use_result_buffer): c0 = self.arange_input(N, np.float64) c1 = self.arange_input(N, np.float64) idx = self.idx_input(N) + result_buffer = self.result_shim(use_result_buffer, N, np.float64) coords = np.vstack((c0, c1)) - result = distopia.calc_bonds_idx_no_box_double(coords, idx) + result = distopia.calc_bonds_idx_no_box_double(coords, idx, results=result_buffer) assert_allclose(result, np.zeros(N))