From e58a4f0a4ec6a13e036aa7cd8a747f9e05bf00f2 Mon Sep 17 00:00:00 2001 From: Brady Johnston Date: Fri, 9 Oct 2026 12:01:43 +0100 Subject: [PATCH 1/3] fix triclinic pkdtree distances --- package/CHANGELOG | 6 ++ .../MDAnalysis/lib/include/calc_distances.h | 4 +- package/MDAnalysis/lib/nsgrid.pyx | 37 +++++------ .../MDAnalysisTests/lib/test_distances.py | 62 +++++++++++++++++++ testsuite/MDAnalysisTests/lib/test_nsgrid.py | 33 ++++++++++ 5 files changed, 118 insertions(+), 24 deletions(-) diff --git a/package/CHANGELOG b/package/CHANGELOG index dede22e74ab..3bb64be96f8 100644 --- a/package/CHANGELOG +++ b/package/CHANGELOG @@ -22,6 +22,12 @@ The rules for this file: * 2.11.0 Fixes + * Fixed `capped_distance` and `self_capped_distance` (and hence bond + guessing, selections and other users of the `nsgrid` and `pkdtree` + methods) missing pairs within the cutoff in triclinic boxes where both + b_x and c_y are non-zero, e.g. truncated octahedra; `apply_PBC` also + wrapped some coordinates in such boxes to just outside the primary + image (Issue #4906) * `MDAnalysis.analysis.nucleicacids.WatsonCrickDist`, `MinorPairDist`, and `MajorPairDist` now match residue names against the full resname instead of only the first character, fixing incorrect behaviour with diff --git a/package/MDAnalysis/lib/include/calc_distances.h b/package/MDAnalysis/lib/include/calc_distances.h index ff130d2c049..106996740ac 100644 --- a/package/MDAnalysis/lib/include/calc_distances.h +++ b/package/MDAnalysis/lib/include/calc_distances.h @@ -294,7 +294,7 @@ static void _triclinic_pbc(coordinate* coords, uint64_t numcoords, float* box) } if (!msr) { // translate remainder of crd to central cell along a-axis - lbound = crd[1] * a_ax_yfactor + crd[2] * a_ax_zfactor; + lbound = (crd[1] - crd[2] * b_ax_zfactor) * a_ax_yfactor + crd[2] * a_ax_zfactor; ubound = lbound + box[0]; if (crd[0] < lbound) { crd[0] += box[0]; @@ -356,7 +356,7 @@ static void _triclinic_pbc(coordinate* coords, uint64_t numcoords, float* box) crd[1] -= box[4]; } // translate remainder of crd to central cell along a-axis - lbound = crd[1] * a_ax_yfactor + crd[2] * a_ax_zfactor; + lbound = (crd[1] - crd[2] * b_ax_zfactor) * a_ax_yfactor + crd[2] * a_ax_zfactor; ubound = lbound + box[0]; if (crd[0] < lbound) { crd[0] += box[0]; diff --git a/package/MDAnalysis/lib/nsgrid.pyx b/package/MDAnalysis/lib/nsgrid.pyx index a4e245e884e..46c259bb384 100644 --- a/package/MDAnalysis/lib/nsgrid.pyx +++ b/package/MDAnalysis/lib/nsgrid.pyx @@ -105,13 +105,6 @@ cdef extern from "calc_distances.h" nogil: void _triclinic_pbc(coordinate* coords, int numcoords, float* box) -cdef inline float fmax(float a, float b): - if a > b: - return a - else: - return b - - cdef inline float fmin(float a, float b): if a < b: return a @@ -305,7 +298,7 @@ cdef class FastNS(object): max_cutoff : float the maximum allowable cutoff given the box shape and size """ - cdef float cutoff, min_cellsize, max_cutoff, new_cellsize + cdef float cutoff, min_cellsize, max_cutoff cdef int i from MDAnalysis.lib.mdamath import triclinic_vectors @@ -337,22 +330,20 @@ cdef class FastNS(object): max_cutoff = fmin(max_cutoff, self.triclinic_dimensions[YY] * degsin(self.dimensions[5])) max_cutoff /= 2 - # for triclinic cells, we need to worry about the shortest path across the cells - min_cellsize = cutoff - if self.triclinic: - for i in range(3, 6): - # cutoff/sin(theta) to elongate the XX/YY/ZZ dimension to make smallest diagonal large enough - new_cellsize = cutoff / degsin(self.dimensions[i]) - min_cellsize = fmax(new_cellsize, min_cellsize) - + # Cells must be at least `cutoff` thick perpendicular to each pair of + # box faces, so size the grid by the box's perpendicular widths. + tv = np.asarray(self.triclinic_dimensions, dtype=np.float64).reshape(3, 3) + volume = abs(np.linalg.det(tv)) + widths = [volume / np.linalg.norm(np.cross(tv[1], tv[2])), + volume / np.linalg.norm(np.cross(tv[2], tv[0])), + volume / np.linalg.norm(np.cross(tv[0], tv[1]))] # add 0.001 here to avoid floating point errors # will make cells slightly too large as a result, ah well - min_cellsize += 0.001 + min_cellsize = cutoff + 0.001 # If the cell size is too small, indexing overflow will occur. Limit the number # of cells in any dimension to the cube root of the maximum of 32 bit integer values. - self.ncells[0] = min(math.floor(self.triclinic_dimensions[XX] / min_cellsize), MAX_GRID_DIM) - self.ncells[1] = min(math.floor(self.triclinic_dimensions[YY] / min_cellsize), MAX_GRID_DIM) - self.ncells[2] = min(math.floor(self.triclinic_dimensions[ZZ] / min_cellsize), MAX_GRID_DIM) + for i in range(3): + self.ncells[i] = min(math.floor(widths[i] / min_cellsize), MAX_GRID_DIM) self.pbc = pbc # If there aren't enough cells in a given dimension it's equivalent to one @@ -436,8 +427,10 @@ cdef class FastNS(object): """Calculate cell coordinate for coord""" # This assumes coordinate is inside the primary unit cell xyz[2] = (coord[2] / self.cellsize[ZZ]) - xyz[1] = ((coord[1] - coord[2] * self.cellsize[YZ]) / self.cellsize[YY]) - xyz[0] = ((coord[0] - coord[1] * self.cellsize[XY] + # y with the c-vector contribution removed, needed for both y and x + cdef double y = coord[1] - coord[2] * self.cellsize[YZ] + xyz[1] = (y / self.cellsize[YY]) + xyz[0] = ((coord[0] - y * self.cellsize[XY] - coord[2] * self.cellsize[XZ]) / self.cellsize[XX]) # Make sure cell coordinate indices are within the primary unit cell # (better safe than sorry): diff --git a/testsuite/MDAnalysisTests/lib/test_distances.py b/testsuite/MDAnalysisTests/lib/test_distances.py index 06abe0a99c7..853f8674afb 100644 --- a/testsuite/MDAnalysisTests/lib/test_distances.py +++ b/testsuite/MDAnalysisTests/lib/test_distances.py @@ -84,6 +84,26 @@ def test_transform_StoR_pass(coord_dtype): assert_allclose(original_r, test_r) +# Triclinic boxes with both b_x and c_y non-zero, see Issue #4906 +SKEWED_TRICLINIC_BOXES = ( + # truncated octahedron + np.array([97.16, 97.16, 97.16, 109.47, 109.47, 109.47], dtype=np.float32), + # LAMMPS cell from Issue #4906 + np.array( + [15.173525, 14.923327, 14.475218, 109.888, 109.3786, 115.684875], + dtype=np.float32, + ), + np.array([26.24, 27.86, 34.66, 106.56, 77.20, 79.13], dtype=np.float32), +) + + +def _skewed_box_coords(box, n, seed): + """Random coordinates spread over the box and its neighbouring images""" + rng = np.random.default_rng(seed) + frac = rng.uniform(-1.5, 2.5, size=(n, 3)) + return (frac @ mdamath.triclinic_vectors(box)).astype(np.float32) + + class TestCappedDistances(object): npoints_1 = (1, 100) @@ -122,6 +142,38 @@ def query_2_atomgroup(self, query_2): min_cutoff_1 = (None, 0.1) + @pytest.mark.parametrize("box", SKEWED_TRICLINIC_BOXES) + @pytest.mark.parametrize("method", ("pkdtree", "nsgrid")) + def test_capped_distance_skewed_triclinic(self, box, method): + # Issue #4906: pkdtree and nsgrid missed pairs in skewed boxes + width = mdamath.triclinic_vectors(box).diagonal().min() + cutoff = 0.1 * width + a = _skewed_box_coords(box, 1000, seed=1) + b = _skewed_box_coords(box, 1000, seed=2) + + def pairs(m): + p = distances.capped_distance( + a, b, cutoff, box=box, method=m, return_distances=False + ) + return {tuple(x) for x in p} + + assert pairs(method) == pairs("bruteforce") + + @pytest.mark.parametrize("box", SKEWED_TRICLINIC_BOXES) + @pytest.mark.parametrize("method", ("pkdtree", "nsgrid")) + def test_self_capped_distance_skewed_triclinic(self, box, method): + width = mdamath.triclinic_vectors(box).diagonal().min() + cutoff = 0.1 * width + a = _skewed_box_coords(box, 1000, seed=3) + + def pairs(m): + p = distances.self_capped_distance( + a, cutoff, box=box, method=m, return_distances=False + ) + return {tuple(sorted(x)) for x in p} + + assert pairs(method) == pairs("bruteforce") + def test_capped_distance_noresults(self): point1 = np.array([0.1, 0.1, 0.1], dtype=np.float32) point2 = np.array([0.95, 0.1, 0.1], dtype=np.float32) @@ -1787,6 +1839,16 @@ def test_coords_in_central_image_tric(self, backend): assert np.all(relres >= 0.0) assert np.all(relres < 1.0) + @pytest.mark.parametrize("box", SKEWED_TRICLINIC_BOXES) + def test_coords_in_central_image_skewed_tric(self, backend, box): + # Issue #4906: boxes where both b_x and c_y are non-zero were wrapped + # to just outside the primary image along a. + coords = _skewed_box_coords(box, 1000, seed=4906) + res = distances.apply_PBC(coords, box, backend=backend) + relres = distances.transform_RtoS(res, box) + assert np.all(relres >= -1e-6) + assert np.all(relres < 1.0 + 1e-6) + @pytest.mark.parametrize("backend", ["serial", "openmp"]) class TestPeriodicAngles(object): diff --git a/testsuite/MDAnalysisTests/lib/test_nsgrid.py b/testsuite/MDAnalysisTests/lib/test_nsgrid.py index 582e780172b..de3536bdd3d 100644 --- a/testsuite/MDAnalysisTests/lib/test_nsgrid.py +++ b/testsuite/MDAnalysisTests/lib/test_nsgrid.py @@ -505,6 +505,39 @@ def test_issue_2670(): assert len(ag2.select_atoms("around 0.0 resid 3")) == 1 +@pytest.mark.parametrize( + "box", + [ + # truncated octahedron + [97.16, 97.16, 97.16, 109.47, 109.47, 109.47], + # LAMMPS cell from Issue #4906 + [15.173525, 14.923327, 14.475218, 109.888, 109.3786, 115.684875], + ], +) +def test_issue_4906(box): + # In boxes with both b_x and c_y non-zero, atoms were binned into the + # wrong cells and cells could be thinner than the cutoff, so neighbours + # within the cutoff were missed. + box = np.array(box, dtype=np.float32) + rng = np.random.default_rng(4906) + frac = rng.uniform(0, 1, size=(1000, 3)) + coords = (frac @ mda.lib.mdamath.triclinic_vectors(box)).astype(np.float32) + cutoff = 0.1 * box[:3].min() + + dists = mda.lib.distances.distance_array(coords, coords, box=box) + + searcher = nsgrid.FastNS(cutoff, coords, box) + pairs = searcher.self_search().get_pairs() + i, j = np.triu_indices(len(coords), k=1) + within = dists[i, j] <= cutoff + expected = set(zip(i[within], j[within])) + assert {tuple(sorted(p)) for p in pairs} == expected + + pairs = searcher.search(coords).get_pairs() + i, j = np.nonzero(dists <= cutoff) + assert {tuple(p) for p in pairs} == set(zip(i, j)) + + def high_mem_tests_enabled(): """Returns true if ENABLE_HIGH_MEM_UNIT_TESTS is set to true.""" env = os.getenv("ENABLE_HIGH_MEM_UNIT_TESTS", default="false").lower() From a3e37a85dcf5b0af782be36b94ff02b693d9bcd5 Mon Sep 17 00:00:00 2001 From: Brady Johnston Date: Fri, 9 Oct 2026 12:23:01 +0100 Subject: [PATCH 2/3] update incorrect expectations --- testsuite/MDAnalysisTests/analysis/test_dielectric.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/testsuite/MDAnalysisTests/analysis/test_dielectric.py b/testsuite/MDAnalysisTests/analysis/test_dielectric.py index 45d8e722cdb..3d22d364184 100644 --- a/testsuite/MDAnalysisTests/analysis/test_dielectric.py +++ b/testsuite/MDAnalysisTests/analysis/test_dielectric.py @@ -41,7 +41,8 @@ def test_broken_molecules(self, ag): ag.wrap() eps = DielectricConstant(ag, make_whole=False).run() - assert_allclose(eps.results["eps_mean"], 721.711, rtol=1e-03) + # reference values changed after the apply_PBC fix for Issue #4906 + assert_allclose(eps.results["eps_mean"], 630.326, rtol=1e-03) def test_broken_repaired_molecules(self, ag): # cut molecules apart @@ -50,7 +51,8 @@ def test_broken_repaired_molecules(self, ag): ag.wrap() eps = DielectricConstant(ag, make_whole=True).run() - assert_allclose(eps.results["eps_mean"], 5.088, rtol=1e-03) + # reference values changed after the apply_PBC fix for Issue #4906 + assert_allclose(eps.results["eps_mean"], 3.874, rtol=1e-03) def test_temperature(self, ag): eps = DielectricConstant(ag, temperature=100).run() From 4940789fedd9ffbba600a1c07fbc5830388e973b Mon Sep 17 00:00:00 2001 From: Brady Johnston Date: Thu, 15 Oct 2026 11:01:52 +0100 Subject: [PATCH 3/3] fix box widths slowdown; cleanup tests and add versionchanged to docs --- package/MDAnalysis/lib/distances.py | 4 +++ .../MDAnalysis/lib/include/calc_distances.h | 6 ++--- package/MDAnalysis/lib/nsgrid.pyx | 27 +++++++++++++++---- .../analysis/test_dielectric.py | 6 ++++- .../MDAnalysisTests/lib/test_distances.py | 25 ++++++++--------- testsuite/MDAnalysisTests/lib/test_nsgrid.py | 17 +++++++----- 6 files changed, 57 insertions(+), 28 deletions(-) diff --git a/package/MDAnalysis/lib/distances.py b/package/MDAnalysis/lib/distances.py index b3f862e005b..5e7ce4aec04 100644 --- a/package/MDAnalysis/lib/distances.py +++ b/package/MDAnalysis/lib/distances.py @@ -2086,6 +2086,10 @@ def apply_PBC( .. versionchanged:: 2.3.0 Can now accept an :class:`~MDAnalysis.core.groups.AtomGroup` as an argument in any position and checks inputs using type hinting. + .. versionchanged:: 2.11.0 + Fixed wrapping in triclinic boxes where both the *x* component of the + second box vector and the *y* component of the third are non-zero + (Issue #4906). """ # coords is an array, the check_coords decorator made sure of that. # Mypy, however, is not aware of that so we have to tell it explicitly. diff --git a/package/MDAnalysis/lib/include/calc_distances.h b/package/MDAnalysis/lib/include/calc_distances.h index 106996740ac..1b9a6b339e6 100644 --- a/package/MDAnalysis/lib/include/calc_distances.h +++ b/package/MDAnalysis/lib/include/calc_distances.h @@ -225,8 +225,6 @@ static void _triclinic_pbc(coordinate* coords, uint64_t numcoords, float* box) const double bi6 = (-bi3 * box[7] - box[6] * bi0) * bi8; const double bi7 = -box[7] * bi4 * bi8; // variables and constants for single box shifts: - const double a_ax_yfactor = (double) box[3] * bi4;; - const double a_ax_zfactor = (double) box[6] * bi8; const double b_ax_zfactor = (double) box[7] * bi8; @@ -294,7 +292,7 @@ static void _triclinic_pbc(coordinate* coords, uint64_t numcoords, float* box) } if (!msr) { // translate remainder of crd to central cell along a-axis - lbound = (crd[1] - crd[2] * b_ax_zfactor) * a_ax_yfactor + crd[2] * a_ax_zfactor; + lbound = -(crd[1] * bi3 + crd[2] * bi6) * box[0]; ubound = lbound + box[0]; if (crd[0] < lbound) { crd[0] += box[0]; @@ -356,7 +354,7 @@ static void _triclinic_pbc(coordinate* coords, uint64_t numcoords, float* box) crd[1] -= box[4]; } // translate remainder of crd to central cell along a-axis - lbound = (crd[1] - crd[2] * b_ax_zfactor) * a_ax_yfactor + crd[2] * a_ax_zfactor; + lbound = -(crd[1] * bi3 + crd[2] * bi6) * box[0]; ubound = lbound + box[0]; if (crd[0] < lbound) { crd[0] += box[0]; diff --git a/package/MDAnalysis/lib/nsgrid.pyx b/package/MDAnalysis/lib/nsgrid.pyx index 46c259bb384..afd3c0e06f0 100644 --- a/package/MDAnalysis/lib/nsgrid.pyx +++ b/package/MDAnalysis/lib/nsgrid.pyx @@ -193,6 +193,10 @@ cdef class FastNS(object): .. versionchanged:: 1.0.2 Rewrote to fix bugs with triclinic boxes + .. versionchanged:: 2.11.0 + Fixed missed neighbours in triclinic boxes where both the *x* + component of the second box vector and the *y* component of the third + are non-zero (Issue #4906). """ cdef readonly double cutoff cdef float[:, ::1] coords_bbox @@ -299,6 +303,8 @@ cdef class FastNS(object): the maximum allowable cutoff given the box shape and size """ cdef float cutoff, min_cellsize, max_cutoff + cdef double ax, bx, by, cx, cy, cz, volume + cdef double widths[3] cdef int i from MDAnalysis.lib.mdamath import triclinic_vectors @@ -332,11 +338,22 @@ cdef class FastNS(object): # Cells must be at least `cutoff` thick perpendicular to each pair of # box faces, so size the grid by the box's perpendicular widths. - tv = np.asarray(self.triclinic_dimensions, dtype=np.float64).reshape(3, 3) - volume = abs(np.linalg.det(tv)) - widths = [volume / np.linalg.norm(np.cross(tv[1], tv[2])), - volume / np.linalg.norm(np.cross(tv[2], tv[0])), - volume / np.linalg.norm(np.cross(tv[0], tv[1]))] + # The box matrix is lower triangular, so width = volume / face area + # has a closed form. + ax = self.triclinic_dimensions[XX] + bx = self.triclinic_dimensions[XY] + by = self.triclinic_dimensions[YY] + cx = self.triclinic_dimensions[XZ] + cy = self.triclinic_dimensions[YZ] + cz = self.triclinic_dimensions[ZZ] + volume = ax * by * cz + widths[0] = widths[1] = widths[2] = 0.0 + # invalid boxes give an all-zero matrix, leave widths at zero + if volume > 0: + widths[0] = volume / math.sqrt((by * cz) ** 2 + (bx * cz) ** 2 + + (bx * cy - by * cx) ** 2) + widths[1] = by * cz / math.sqrt(cy * cy + cz * cz) + widths[2] = cz # add 0.001 here to avoid floating point errors # will make cells slightly too large as a result, ah well min_cellsize = cutoff + 0.001 diff --git a/testsuite/MDAnalysisTests/analysis/test_dielectric.py b/testsuite/MDAnalysisTests/analysis/test_dielectric.py index 3d22d364184..ab7735a9159 100644 --- a/testsuite/MDAnalysisTests/analysis/test_dielectric.py +++ b/testsuite/MDAnalysisTests/analysis/test_dielectric.py @@ -45,13 +45,17 @@ def test_broken_molecules(self, ag): assert_allclose(eps.results["eps_mean"], 630.326, rtol=1e-03) def test_broken_repaired_molecules(self, ag): + intact = DielectricConstant(ag, make_whole=True).run() # cut molecules apart ag.universe.transfer_to_memory() for ts in ag.universe.trajectory: ag.wrap() + # repairing the molecules should recover the intact result eps = DielectricConstant(ag, make_whole=True).run() - # reference values changed after the apply_PBC fix for Issue #4906 + assert_allclose( + eps.results["eps_mean"], intact.results["eps_mean"], rtol=1e-03 + ) assert_allclose(eps.results["eps_mean"], 3.874, rtol=1e-03) def test_temperature(self, ag): diff --git a/testsuite/MDAnalysisTests/lib/test_distances.py b/testsuite/MDAnalysisTests/lib/test_distances.py index 853f8674afb..bf95e1587f9 100644 --- a/testsuite/MDAnalysisTests/lib/test_distances.py +++ b/testsuite/MDAnalysisTests/lib/test_distances.py @@ -104,6 +104,11 @@ def _skewed_box_coords(box, n, seed): return (frac @ mdamath.triclinic_vectors(box)).astype(np.float32) +def _away_from_cutoff(pairs, dists, cutoff, tol=1e-4): + """Drop pairs so close to the cutoff that rounding can decide them""" + return pairs[np.abs(dists - cutoff) > tol] + + class TestCappedDistances(object): npoints_1 = (1, 100) @@ -146,31 +151,27 @@ def query_2_atomgroup(self, query_2): @pytest.mark.parametrize("method", ("pkdtree", "nsgrid")) def test_capped_distance_skewed_triclinic(self, box, method): # Issue #4906: pkdtree and nsgrid missed pairs in skewed boxes - width = mdamath.triclinic_vectors(box).diagonal().min() - cutoff = 0.1 * width + min_diag = mdamath.triclinic_vectors(box).diagonal().min() + cutoff = 0.1 * min_diag a = _skewed_box_coords(box, 1000, seed=1) b = _skewed_box_coords(box, 1000, seed=2) def pairs(m): - p = distances.capped_distance( - a, b, cutoff, box=box, method=m, return_distances=False - ) - return {tuple(x) for x in p} + p, d = distances.capped_distance(a, b, cutoff, box=box, method=m) + return {tuple(x) for x in _away_from_cutoff(p, d, cutoff)} assert pairs(method) == pairs("bruteforce") @pytest.mark.parametrize("box", SKEWED_TRICLINIC_BOXES) @pytest.mark.parametrize("method", ("pkdtree", "nsgrid")) def test_self_capped_distance_skewed_triclinic(self, box, method): - width = mdamath.triclinic_vectors(box).diagonal().min() - cutoff = 0.1 * width + min_diag = mdamath.triclinic_vectors(box).diagonal().min() + cutoff = 0.1 * min_diag a = _skewed_box_coords(box, 1000, seed=3) def pairs(m): - p = distances.self_capped_distance( - a, cutoff, box=box, method=m, return_distances=False - ) - return {tuple(sorted(x)) for x in p} + p, d = distances.self_capped_distance(a, cutoff, box=box, method=m) + return {tuple(sorted(x)) for x in _away_from_cutoff(p, d, cutoff)} assert pairs(method) == pairs("bruteforce") diff --git a/testsuite/MDAnalysisTests/lib/test_nsgrid.py b/testsuite/MDAnalysisTests/lib/test_nsgrid.py index de3536bdd3d..7c46029dd59 100644 --- a/testsuite/MDAnalysisTests/lib/test_nsgrid.py +++ b/testsuite/MDAnalysisTests/lib/test_nsgrid.py @@ -525,17 +525,22 @@ def test_issue_4906(box): cutoff = 0.1 * box[:3].min() dists = mda.lib.distances.distance_array(coords, coords, box=box) + # skip pairs so close to the cutoff that rounding can decide them + ambiguous = np.abs(dists - cutoff) <= 1e-4 + + def found(results): + pairs = results.get_pairs() + return {tuple(p) for p in pairs if not ambiguous[p[0], p[1]]} searcher = nsgrid.FastNS(cutoff, coords, box) - pairs = searcher.self_search().get_pairs() i, j = np.triu_indices(len(coords), k=1) - within = dists[i, j] <= cutoff + within = (dists[i, j] <= cutoff) & ~ambiguous[i, j] expected = set(zip(i[within], j[within])) - assert {tuple(sorted(p)) for p in pairs} == expected + pairs = {tuple(sorted(p)) for p in found(searcher.self_search())} + assert pairs == expected - pairs = searcher.search(coords).get_pairs() - i, j = np.nonzero(dists <= cutoff) - assert {tuple(p) for p in pairs} == set(zip(i, j)) + i, j = np.nonzero((dists <= cutoff) & ~ambiguous) + assert found(searcher.search(coords)) == set(zip(i, j)) def high_mem_tests_enabled():