Skip to content

Commit

Permalink
Fix distancearray result passing (#170)
Browse files Browse the repository at this point in the history 
* fix ver

* make distance_array result input 2D

* result passing

* fixify

* change doc example
  • Loading branch information
@hmacdope
hmacdope authored Oct 15, 2024
1 parent f77b4ba commit 50fd4a9
Show file tree
Hide file tree
Showing 3 changed files with 57 additions and 36 deletions.
63 changes: 31 additions & 32 deletions distopia/_distopia.pyx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -150,10 +150,10 @@ def _check_results(results, nvals):

def _check_results_darray(results, nvals0 , nvals1 ):
"""Check that results is the right shape"""
if results.ndim > 1:
raise ValueError("results must be a 1D array")
if results.shape[0] != nvals0 * nvals1:
raise ValueError(f"results must be a flattened 2D array of length MxN ({ nvals0, nvals1} -> {nvals0 *nvals1}), you provided {results.shape[0]}")
if results.ndim > 2:
raise ValueError("results must be a 2D array")
if results.shape[0] != nvals0 or results.shape[1] != nvals1:
raise ValueError(f"results must be a 2D array of length MxN ({ nvals0, nvals1}), you provided {results.shape}")


def _check_shapes(*args):
Expand All @@ -164,8 +164,6 @@ def _check_shapes(*args):
f"{', '.join(str(t.shape) for t in args)}")






@cython.boundscheck(False)
Expand Down Expand Up @@ -598,35 +596,36 @@ def calc_dihedrals_triclinic(
def calc_distance_array_no_box(
floating[:, ::1] coords0,
floating[:, ::1] coords1,
floating[::1] results=None):
floating[:, ::1] results=None):
"""Calculate pairwise distance matrix between coordinates with no periodic boundary conditions
Parameters
----------
coords0, coords1 : float32 or float64 array
must be same length and dtype
results: float32 or float64 array (optional)
array to store results in, must be a single dimension of length MxN where M is the length of coords0 and N is the length of coords1
array to store results in, must be of length MxN where M is the length of coords0 and N is the length of coords1
Returns
-------
distances : np.array
MxN array of distances
"""

cdef floating[::1] results_view
cdef floating[:, ::1] results_view
cdef size_t nvals0 = coords0.shape[0]
cdef size_t nvals1 = coords1.shape[0]
cdef cnp.npy_intp[1] dims
cdef cnp.npy_intp[2] dims

dims[0] = <ssize_t > nvals0 * nvals1
dims[0] = <ssize_t > nvals0
dims[1] = <ssize_t> nvals1


if results is None:
if floating is float:
results = cnp.PyArray_EMPTY(1, dims, cnp.NPY_FLOAT32, 0)
results = cnp.PyArray_EMPTY(2, dims, cnp.NPY_FLOAT32, 0)
else:
results = cnp.PyArray_EMPTY(1, dims, cnp.NPY_FLOAT64, 0)
results = cnp.PyArray_EMPTY(2, dims, cnp.NPY_FLOAT64, 0)

else:
_check_results_darray(results, nvals0, nvals1)
Expand All @@ -636,7 +635,7 @@ def calc_distance_array_no_box(

CalcDistanceArrayNoBox(&coords0[0][0], &coords1[0][0],
nvals0, nvals1,
&results_view[0])
&results_view[0][0])

return np.array(results).reshape(coords0.shape[0], coords1.shape[0])

Expand All @@ -646,7 +645,7 @@ def calc_distance_array_ortho(
floating[:, ::1] coords0,
floating[:, ::1] coords1,
floating[::1] box,
floating[::1] results=None):
floating[:, ::1] results=None):
"""Calculate pairwise distance matrix between coordinates under orthorhombic boundary conditions
Parameters
Expand All @@ -656,26 +655,26 @@ def calc_distance_array_ortho(
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 a single dimension of length MxN where M is the length of coords0 and N is the length of coords1
array to store results in, must be of length MxN where M is the length of coords0 and N is the length of coords1
Returns
-------
distances : np.array
MxN array of distances
"""
cdef floating[::1] results_view
cdef floating[:, ::1] results_view
cdef size_t nvals0 = coords0.shape[0]
cdef size_t nvals1 = coords1.shape[0]
cdef cnp.npy_intp[1] dims

dims[0] = <ssize_t > nvals0 * nvals1
cdef cnp.npy_intp[2] dims

dims[0] = <ssize_t > nvals0
dims[1] = <ssize_t> nvals1

if results is None:
if floating is float:
results = cnp.PyArray_EMPTY(1, dims, cnp.NPY_FLOAT32, 0)
results = cnp.PyArray_EMPTY(2, dims, cnp.NPY_FLOAT32, 0)
else:
results = cnp.PyArray_EMPTY(1, dims, cnp.NPY_FLOAT64, 0)
results = cnp.PyArray_EMPTY(2, dims, cnp.NPY_FLOAT64, 0)

else:
_check_results_darray(results, nvals0, nvals1)
Expand All @@ -685,7 +684,7 @@ def calc_distance_array_ortho(
CalcDistanceArrayOrtho(&coords0[0][0], &coords1[0][0],
nvals0, nvals1,
&box[0],
&results_view[0])
&results_view[0][0])

return np.array(results).reshape(coords0.shape[0], coords1.shape[0])

Expand All @@ -696,7 +695,7 @@ def calc_distance_array_triclinic(
floating[:, ::1] coords0,
floating[:, ::1] coords1,
floating[:, ::1] box,
floating[::1] results=None):
floating[:, ::1] results=None):
"""Calculate pairwise distance matrix between coordinates under triclinic boundary conditions
Parameters
Expand All @@ -706,28 +705,28 @@ def calc_distance_array_triclinic(
box : float32 or float64 array
periodic boundary dimensions, in 3x3 format
results: float32 or float64 array (optional)
array to store results in, must be a single dimension of length MxN where M is the length of coords0 and N is the length of coords1
array to store results in, must be of length MxN where M is the length of coords0 and N is the length of coords1
Returns
-------
distances : np.array
MxN array of distances
"""
cdef floating[::1] results_view
cdef floating[:, ::1] results_view
cdef size_t nvals0 = coords0.shape[0]
cdef size_t nvals1 = coords1.shape[0]
cdef cnp.npy_intp[2] dims

cdef cnp.npy_intp[1] dims

dims[0] = <ssize_t > nvals0 * nvals1
dims[0] = <ssize_t > nvals0
dims[1] = <ssize_t> nvals1



if results is None:
if floating is float:
results = cnp.PyArray_EMPTY(1, dims, cnp.NPY_FLOAT32, 0)
results = cnp.PyArray_EMPTY(2, dims, cnp.NPY_FLOAT32, 0)
else:
results = cnp.PyArray_EMPTY(1, dims, cnp.NPY_FLOAT64, 0)
results = cnp.PyArray_EMPTY(2, dims, cnp.NPY_FLOAT64, 0)

else:
_check_results_darray(results, nvals0, nvals1)
Expand All @@ -737,7 +736,7 @@ def calc_distance_array_triclinic(
CalcDistanceArrayTriclinic(&coords0[0][0], &coords1[0][0],
nvals0, nvals1,
&box[0][0],
&results_view[0])
&results_view[0][0])

return np.array(results).reshape(coords0.shape[0], coords1.shape[0])

Expand Down
28 changes: 25 additions & 3 deletions distopia/tests/test_distopia.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -195,24 +195,35 @@ def test_no_box_bad_result(self):
c0 = np.zeros(6, dtype=np.float32).reshape(2, 3)
c1 = np.zeros(6, dtype=np.float32).reshape(2, 3)
with pytest.raises(ValueError, match="results must be"):
distopia.calc_distance_array_no_box(c0, c1, results=np.empty(1, dtype=np.float32))
distopia.calc_distance_array_no_box(c0, c1, results=np.empty((1,1), dtype=np.float32))

def test_ortho_bad_result(self):
c0 = np.zeros(6, dtype=np.float32).reshape(2, 3)
c1 = np.zeros(6, dtype=np.float32).reshape(2, 3)
box = np.array([10, 10, 10], dtype=np.float32)
with pytest.raises(ValueError, match="results must be"):
distopia.calc_distance_array_ortho(c0, c1, box, results=np.empty(1, dtype=np.float32))
distopia.calc_distance_array_ortho(c0, c1, box, results=np.empty((1,1), dtype=np.float32))


def test_triclinic_bad_result(self):
c0 = np.zeros(6, dtype=np.float32).reshape(2, 3)
c1 = np.zeros(6, dtype=np.float32).reshape(2, 3)
box = np.array([[10, 0, 0], [0, 10, 0], [0, 0, 10]], dtype=np.float32)
with pytest.raises(ValueError, match="results must be"):
distopia.calc_distance_array_triclinic(c0, c1, box, results=np.empty(1, dtype=np.float32))
distopia.calc_distance_array_triclinic(c0, c1, box, results=np.empty((1,1), dtype=np.float32))


def test_distance_array_arange(self):
c0 = np.ones(9, dtype=np.float32).reshape(3, 3)
c1 = np.ones(9, dtype=np.float32).reshape(3, 3)
results = distopia.calc_distance_array_no_box(c0, c1)
assert_almost_equal(results, np.zeros((3,3), dtype=np.float32))

def test_distance_array_results(self):
c0 = np.ones(9, dtype=np.float32).reshape(3, 3)
c1 = np.ones(9, dtype=np.float32).reshape(3, 3)
results = distopia.calc_distance_array_no_box(c0, c1, results=np.empty((3,3), dtype=np.float32))
assert_almost_equal(results, np.zeros((3,3), dtype=np.float32))


class TestSelfDistanceArray:
Expand All @@ -234,6 +245,17 @@ def test_triclinic_bad_result(self):
with pytest.raises(ValueError, match="results must be"):
distopia.calc_self_distance_array_triclinic(c0, box, results=np.empty(1, dtype=np.float32))

def test_self_distance_ones(self):
c0 = np.ones(9, dtype=np.float32).reshape(3, 3)
results = distopia.calc_self_distance_array_no_box(c0)
assert_almost_equal(results, np.zeros(3, dtype=np.float32))


def test_self_distance_ones_result(self):
c0 = np.ones(9, dtype=np.float32).reshape(3, 3)
# n(n-1) //2
results = distopia.calc_self_distance_array_no_box(c0, results=np.empty(3, dtype=np.float32))
assert_almost_equal(results, np.zeros(3, dtype=np.float32))

class TestMDA:
"""
Expand Down
2 changes: 1 addition & 1 deletion docs/docs/examples.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -113,7 +113,7 @@ result = distopia.calc_distance_array_no_box(coordinates0, coordinates1)
result # -> will be NxM

# passing in a result buffer is also possible for distance arrays
buffer = np.empty(N*M, dtype=np.float32)
buffer = np.empty((N,M), dtype=np.float32)
result = distopia.calc_distance_array_no_box(coordinates0, coordinates1, results=buffer)
```

Expand Down

0 comments on commit 50fd4a9

Please sign in to comment.