From 57e0b51bd384620502e823b02b200c38aa3be1ea Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 15 Oct 2026 15:47:28 +0100 Subject: [PATCH] perf: numba Sibson kernel for the NumPy DelaunayNN path Replace the per-query Python loop in the NumPy branch of sibson_mappings_weights_from_tables with a parallel numba kernel (module-level in _sibson_numba.py, lazily imported, cross-process cache) that ports _sibson_single_from_tables exactly: cavity walk and overflow, Watson contributions, outside-mesh/on-vertex/on-edge rows, degenerate and neighbour-overflow NaN rows and all six outputs. The reference loop stays as the fallback when numba is unavailable or disabled. JAX path untouched. Aris SMA interferometer, 1405 vertices: DelaunayNN 18.1 s -> 0.57 s per likelihood (Delaunay 0.41 s); Sibson step 8.5 s -> 13 ms. Phase 1 of #623. Co-Authored-By: Claude Opus 5.5 --- .../mesh/interpolator/_sibson_numba.py | 362 ++++++++++++++++++ .../inversion/mesh/interpolator/sibson.py | 181 ++++++++- .../interpolator/test_sibson_numba.py | 185 +++++++++ 3 files changed, 720 insertions(+), 8 deletions(-) create mode 100644 autoarray/inversion/mesh/interpolator/_sibson_numba.py create mode 100644 test_autoarray/inversion/pixelization/interpolator/test_sibson_numba.py diff --git a/autoarray/inversion/mesh/interpolator/_sibson_numba.py b/autoarray/inversion/mesh/interpolator/_sibson_numba.py new file mode 100644 index 000000000..c1f962b0d --- /dev/null +++ b/autoarray/inversion/mesh/interpolator/_sibson_numba.py @@ -0,0 +1,362 @@ +"""numba kernel for the NumPy Sibson path; see ``sibson.sibson_numba_kernel``. + +Kept in its own module, imported only on first use, so ``import numba`` stays +off the library import path while every jitted function is a plain +module-level function: numba's on-disk cache (``cache=True``) keys functions +by module file and qualified name, which closures built inside a factory do +not have stably, so they recompiled in every new process. +""" + +import numba +import numpy as np +from numba import prange + +# float64 machine constants as module globals, which numba freezes into the +# compiled code (and the on-disk cache) as literals. +EPS = 2.220446049250313e-16 +TINY = 2.2250738585072014e-308 + + +@numba.njit(cache=True, nogil=True) +def _np_maximum(a, b): + # np.maximum semantics: NaN propagates. + if a != a or b != b: + return np.nan + return a if a >= b else b + + +@numba.njit(cache=True, nogil=True) +def _block_sum(values, start, n): + # NumPy ``pairwise_sum`` leaf (n <= 128): plain loop below 8, otherwise + # eight interleaved accumulators combined as a fixed tree. + if n < 8: + result = -0.0 + for i in range(n): + result += values[start + i] + return result + r0 = values[start + 0] + r1 = values[start + 1] + r2 = values[start + 2] + r3 = values[start + 3] + r4 = values[start + 4] + r5 = values[start + 5] + r6 = values[start + 6] + r7 = values[start + 7] + i = 8 + stop = n - (n % 8) + while i < stop: + r0 += values[start + i + 0] + r1 += values[start + i + 1] + r2 += values[start + i + 2] + r3 += values[start + i + 3] + r4 += values[start + i + 4] + r5 += values[start + i + 5] + r6 += values[start + i + 6] + r7 += values[start + i + 7] + i += 8 + result = ((r0 + r1) + (r2 + r3)) + ((r4 + r5) + (r6 + r7)) + while i < n: + result += values[start + i] + i += 1 + return result + + +@numba.njit(cache=True, nogil=True) +def _pairwise_sum(values): + # Explicit-stack form of NumPy's recursive ``pairwise_sum``: blocks + # above 128 split at ``n2 = n // 2 - (n // 2) % 8`` and the two halves + # are added left + right. + node_start = np.empty(64, dtype=np.int64) + node_n = np.empty(64, dtype=np.int64) + node_state = np.empty(64, dtype=np.int64) + partials = np.empty(64, dtype=np.float64) + top = 0 + partial_top = 0 + node_start[0] = 0 + node_n[0] = values.shape[0] + node_state[0] = 0 + top = 1 + while top > 0: + top -= 1 + start = node_start[top] + n = node_n[top] + state = node_state[top] + if n <= 128: + partials[partial_top] = _block_sum(values, start, n) + partial_top += 1 + elif state == 0: + n2 = n // 2 + n2 -= n2 % 8 + node_start[top] = start + node_n[top] = n + node_state[top] = 1 + node_start[top + 1] = start + n2 + node_n[top + 1] = n - n2 + node_state[top + 1] = 0 + node_start[top + 2] = start + node_n[top + 2] = n2 + node_state[top + 2] = 0 + top += 3 + else: + right = partials[partial_top - 1] + left = partials[partial_top - 2] + partial_top -= 2 + partials[partial_top] = left + right + partial_top += 1 + return partials[0] + + +@numba.njit(cache=True, parallel=True, nogil=True) +def _sibson_from_tables_parallel( + query_points, + seeds, + fallbacks, + points, + simplices_padded, + simplex_neighbors, + circumcentres, + circumradii_squared, + circumcircle_valid, + max_cavity_triangles, + max_neighbors, +): + query_count = query_points.shape[0] + cavity_cap = max_cavity_triangles + compact_length = 3 * cavity_cap + + mappings = np.full((query_count, max_neighbors), -1, dtype=np.int32) + sizes = np.zeros(query_count, dtype=np.int32) + weights = np.zeros((query_count, max_neighbors), dtype=np.float64) + cavity_sizes = np.zeros(query_count, dtype=np.int32) + overflow_out = np.zeros(query_count, dtype=np.bool_) + degenerate_out = np.zeros(query_count, dtype=np.bool_) + + for q in prange(query_count): + qx = query_points[q, 0] + qy = query_points[q, 1] + seed = seeds[q] + + # Outside the mesh: nearest-vertex fallback row, no diagnostics. + if seed < 0: + mappings[q, 0] = fallbacks[q] + weights[q, 0] = 1.0 + sizes[q] = 1 + continue + + # Cavity walk (``_cavity_triangle_indexes_numpy``). + cavity = np.full(cavity_cap, -1, dtype=np.int64) + cavity[0] = seed + count = 1 + cavity_overflow = False + for position in range(cavity_cap): + if position >= count: + break + triangle_index = cavity[position] + for edge in range(3): + candidate = simplex_neighbors[triangle_index, edge] + if candidate < 0: + continue + seen = False + for k in range(count): + if cavity[k] == candidate: + seen = True + break + if seen: + continue + dx = circumcentres[candidate, 0] - qx + dy = circumcentres[candidate, 1] - qy + distance_squared = dx * dx + dy * dy + radius_squared = circumradii_squared[candidate] + tolerance = 64.0 * EPS * _np_maximum(radius_squared, 1.0) + if not ( + circumcircle_valid[candidate] + and distance_squared <= radius_squared + tolerance + ): + continue + if count == cavity_cap: + cavity_overflow = True + else: + cavity[count] = candidate + count += 1 + + # Watson contributions, grouped by vertex id into a sorted list. + unique_ids = np.full(compact_length, -1, dtype=np.int64) + compact_weights = np.zeros(compact_length, dtype=np.float64) + unique_count = 0 + inserted_invalid = False + vx = np.empty(3, dtype=np.float64) + vy = np.empty(3, dtype=np.float64) + vid = np.empty(3, dtype=np.int64) + rx = np.empty(3, dtype=np.float64) + ry = np.empty(3, dtype=np.float64) + ivalid = np.empty(3, dtype=np.bool_) + contribution = np.empty(3, dtype=np.float64) + for slot in range(count): + triangle_index = cavity[slot] + for j in range(3): + vertex = simplices_padded[triangle_index, j] + if vertex < 0: + vertex = 0 + vid[j] = vertex + vx[j] = points[vertex, 0] + vy[j] = points[vertex, 1] + ox = circumcentres[triangle_index, 0] + oy = circumcentres[triangle_index, 1] + # Circle j passes through the query and the edge opposite j. + for j in range(3): + ax = vx[(j + 1) % 3] + ay = vy[(j + 1) % 3] + bax = vx[(j + 2) % 3] - ax + bay = vy[(j + 2) % 3] - ay + cax = qx - ax + cay = qy - ay + denominator = 2.0 * (bax * cay - bay * cax) + ba2 = bax * bax + bay * bay + ca2 = cax * cax + cay * cay + scale = _np_maximum(_np_maximum(ba2, ca2), 1.0) + valid = abs(denominator) > 32.0 * EPS * scale + safe_denominator = denominator if valid else 1.0 + offset_x = (ba2 * cay - ca2 * bay) / safe_denominator + offset_y = (bax * ca2 - cax * ba2) / safe_denominator + if valid: + centre_x = ax + offset_x + centre_y = ay + offset_y + else: + centre_x = 0.0 + centre_y = 0.0 + inserted_invalid = True + ivalid[j] = valid + rx[j] = centre_x - ox + ry[j] = centre_y - oy + for j in range(3): + j1 = (j + 1) % 3 + j2 = (j + 2) % 3 + if ivalid[j1] and ivalid[j2]: + contribution[j] = rx[j1] * ry[j2] - ry[j1] * rx[j2] + else: + contribution[j] = 0.0 + for j in range(3): + vertex = vid[j] + position = 0 + while position < unique_count and unique_ids[position] < vertex: + position += 1 + if position < unique_count and unique_ids[position] == vertex: + compact_weights[position] += contribution[j] + else: + for k in range(unique_count, position, -1): + unique_ids[k] = unique_ids[k - 1] + compact_weights[k] = compact_weights[k - 1] + unique_ids[position] = vertex + compact_weights[position] = 0.0 + contribution[j] + unique_count += 1 + + # On-vertex / on-edge tests against the seed triangle. + s0 = simplices_padded[seed, 0] + s1 = simplices_padded[seed, 1] + s2 = simplices_padded[seed, 2] + if s0 < 0: + s0 = 0 + if s1 < 0: + s1 = 0 + if s2 < 0: + s2 = 0 + p0x = points[s0, 0] + p0y = points[s0, 1] + p1x = points[s1, 0] + p1y = points[s1, 1] + p2x = points[s2, 0] + p2y = points[s2, 1] + + d0x = p0x - qx + d0y = p0y - qy + d1x = p1x - qx + d1y = p1y - qy + d2x = p2x - qx + d2y = p2y - qy + dist0 = d0x * d0x + d0y * d0y + dist1 = d1x * d1x + d1y * d1y + dist2 = d2x * d2x + d2y * d2y + local_vertex = s0 + min_distance = dist0 + if dist1 < min_distance: + local_vertex = s1 + min_distance = dist1 + if dist2 < min_distance: + local_vertex = s2 + min_distance = dist2 + + coordinate_scale = _np_maximum(abs(qx), abs(qy)) + point_scale = _np_maximum(abs(p0x), abs(p0y)) + point_scale = _np_maximum(point_scale, abs(p1x)) + point_scale = _np_maximum(point_scale, abs(p1y)) + point_scale = _np_maximum(point_scale, abs(p2x)) + point_scale = _np_maximum(point_scale, abs(p2y)) + coordinate_scale = _np_maximum(_np_maximum(coordinate_scale, point_scale), 1.0) + vertex_tolerance = 64.0 * EPS * coordinate_scale + vertex_tolerance = vertex_tolerance * vertex_tolerance + on_vertex = min_distance <= vertex_tolerance + + seed_denominator = (p1x - p0x) * (p2y - p0y) - (p1y - p0y) * (p2x - p0x) + if seed_denominator == 0.0: + seed_denominator = 1.0 + bary0 = (d1x * d2y - d1y * d2x) / seed_denominator + bary1 = (d2x * d0y - d2y * d0x) / seed_denominator + bary2 = (d0x * d1y - d0y * d1x) / seed_denominator + edge_tolerance = 128.0 * EPS + min_bary = _np_maximum(-abs(bary0), -abs(bary1)) + min_bary = -_np_maximum(min_bary, -abs(bary2)) + on_edge = (not on_vertex) and (min_bary <= edge_tolerance) + + weight_sum = _pairwise_sum(compact_weights) + degenerate = ( + (not on_vertex) + and (not on_edge) + and ( + inserted_invalid + or (not np.isfinite(weight_sum)) + or abs(weight_sum) <= TINY + ) + ) + safe_sum = weight_sum if abs(weight_sum) > TINY else 1.0 + + if on_vertex: + mappings[q, 0] = local_vertex + weights[q, 0] = 1.0 + size = 1 + elif on_edge: + # Barycentric row of the containing triangle (third weight 0). + if max_neighbors > 0: + mappings[q, 0] = s0 + weights[q, 0] = bary0 + if max_neighbors > 1: + mappings[q, 1] = s1 + weights[q, 1] = bary1 + if max_neighbors > 2: + mappings[q, 2] = s2 + weights[q, 2] = bary2 + size = 3 + else: + size = 0 + for k in range(max_neighbors): + if k >= compact_length or k >= unique_count: + break + weight = compact_weights[k] / safe_sum + if weight > 0.0: + mappings[q, k] = unique_ids[k] + weights[q, k] = weight + size += 1 + + neighbor_overflow = ( + (not on_vertex) and (not on_edge) and unique_count > max_neighbors + ) + failed = cavity_overflow or neighbor_overflow or degenerate + if failed: + for k in range(max_neighbors): + weights[q, k] = np.nan + + sizes[q] = size + cavity_sizes[q] = count + overflow_out[q] = cavity_overflow or neighbor_overflow + degenerate_out[q] = degenerate + + return mappings, sizes, weights, cavity_sizes, overflow_out, degenerate_out diff --git a/autoarray/inversion/mesh/interpolator/sibson.py b/autoarray/inversion/mesh/interpolator/sibson.py index 53b75cc8f..e30e4d9fc 100644 --- a/autoarray/inversion/mesh/interpolator/sibson.py +++ b/autoarray/inversion/mesh/interpolator/sibson.py @@ -556,6 +556,160 @@ def _sibson_single_from_tables( return mappings, size, weights, cavity_size, overflow, degenerate +_SIBSON_NUMBA_CACHE: dict = {} + + +def _sibson_use_numba(): + """Whether the NumPy Sibson path dispatches to the numba kernel. + + True when numba is importable and ``general.yaml`` ``numba: use_numba`` is + true (the default when the config cannot be read). Otherwise the NumPy + path keeps the per-query Python reference loop. + """ + try: + from autonerves import conf + + use_numba = conf.instance["general"]["numba"]["use_numba"] + except Exception: + use_numba = True + + if not use_numba: + return False + + try: + import numba # noqa: F401 + except ImportError: + # Not installed, or installed against an incompatible NumPy. + return False + + return True + + +def sibson_numba_kernel(): + """ + Compile (once) the numba kernel evaluating every Sibson query in parallel. + + The kernel is a line-by-line port of :func:`_sibson_single_from_tables` + (NumPy branch) with ``prange`` over queries: the same cavity walk and + containment tolerance, the same Watson circumcircle arithmetic in the same + operation order, the same on-vertex / on-edge / outside-mesh fallbacks, + degeneracy and overflow diagnostics, and the same padding of unused slots. + Each query owns its scratch arrays and writes only its own output row, so + the parallel loop needs neither a reduction nor a lock. + + Two details are reproduced deliberately rather than approximated: + + * the per-vertex grouping inserts the at-most ``3 * C`` contributions into + a small sorted list (ascending vertex id, contributions summed in cavity + order), which is the reference's ``argsort`` + ``np.add.at`` grouping up + to the order equal vertex ids are summed in -- NumPy's argsort is not + stable, so weights agree to rounding (parity tests: <= 1e-12, observed + ~1e-14 near the hull where cavity contributions cancel), not + bit-for-bit; + * the normalising sum replicates NumPy's pairwise summation over the + length ``3 * C`` compact vector, so ``weight_sum`` (and the degeneracy + test on it) sees the same blocking as ``np.sum``. + + The jitted functions live at module level in ``_sibson_numba``, which is + imported only here, so ``import numba`` stays off the library import path + (``numba_util.jit`` cannot express ``parallel=True``) while numba's + on-disk cache still hits across processes -- closures built inside this + function recompiled in every new process. + + Returns + ------- + The compiled kernel ``(query_points, seeds, fallbacks, points, + simplices_padded, simplex_neighbors, circumcentres, circumradii_squared, + circumcircle_valid, max_cavity_triangles, max_neighbors)`` returning the six + arrays of :func:`sibson_mappings_weights_from_tables`. + """ + if "kernel" in _SIBSON_NUMBA_CACHE: + return _SIBSON_NUMBA_CACHE["kernel"] + + from autoarray.inversion.mesh.interpolator._sibson_numba import ( + _sibson_from_tables_parallel, + ) + + _SIBSON_NUMBA_CACHE["kernel"] = _sibson_from_tables_parallel + return _sibson_from_tables_parallel + + +def _sibson_numba_eligible(query_points, points): + """Float64 inputs only: the kernel hard-codes the float64 tolerances.""" + return ( + np.asarray(query_points).dtype == np.float64 + and np.asarray(points).dtype == np.float64 + ) + + +def _sibson_mappings_weights_numba( + query_points, + points, + simplices_padded, + simplex_neighbors, + simplex_indexes, + outside_fallback_indexes, + circumcircles, + max_cavity_triangles, + max_neighbors, +): + """NumPy-path Sibson tables from the parallel numba kernel.""" + circumcentres, circumradii_squared, circumcircle_valid = circumcircles + kernel = sibson_numba_kernel() + return kernel( + np.ascontiguousarray(query_points, dtype=np.float64), + np.ascontiguousarray(simplex_indexes, dtype=np.int64), + np.ascontiguousarray(outside_fallback_indexes, dtype=np.int64), + np.ascontiguousarray(points, dtype=np.float64), + np.ascontiguousarray(simplices_padded, dtype=np.int64), + np.ascontiguousarray(simplex_neighbors, dtype=np.int64), + np.ascontiguousarray(circumcentres, dtype=np.float64), + np.ascontiguousarray(circumradii_squared, dtype=np.float64), + np.ascontiguousarray(circumcircle_valid, dtype=np.bool_), + int(max_cavity_triangles), + int(max_neighbors), + ) + + +def _sibson_mappings_weights_numpy_reference( + query_points, + points, + simplices_padded, + simplex_neighbors, + simplex_indexes, + outside_fallback_indexes, + circumcircles, + max_cavity_triangles, + max_neighbors, +): + """NumPy-path Sibson tables from the per-query Python reference loop. + + The parity oracle for :func:`sibson_numba_kernel` and the fallback when + numba is unavailable or disabled in ``general.yaml``. + """ + circumcentres, circumradii_squared, circumcircle_valid = circumcircles + rows = [ + _sibson_single_from_tables( + query=query, + seed_simplex=seed, + outside_fallback_index=fallback, + points=points, + simplices_padded=simplices_padded, + simplex_neighbors=simplex_neighbors, + circumcentres=circumcentres, + circumradii_squared=circumradii_squared, + circumcircle_valid=circumcircle_valid, + max_cavity_triangles=max_cavity_triangles, + max_neighbors=max_neighbors, + xp=np, + ) + for query, seed, fallback in zip( + query_points, simplex_indexes, outside_fallback_indexes + ) + ] + return tuple(np.stack(values) for values in zip(*rows)) + + def sibson_mappings_weights_from_tables( query_points, points, @@ -585,7 +739,11 @@ def sibson_mappings_weights_from_tables( latency-bound program rather than saving time, which is why ``query_chunk`` is documented as a memory guard on the ``(C, 3, 2)`` per-cavity intermediates and nothing else (issue #532). The NumPy path - ignores ``query_chunk`` entirely and loops over queries in Python. + ignores ``query_chunk`` entirely: it runs every query through the parallel + numba kernel (:func:`sibson_numba_kernel`) when numba is available and + ``general.yaml`` ``numba: use_numba`` is true, and otherwise loops over + queries in Python (:func:`_sibson_single_from_tables`, the parity + reference). Parameters ---------- @@ -627,13 +785,20 @@ def single(query, seed_simplex, outside_fallback_index): ) if xp is np: - rows = [ - single(query, seed, fallback) - for query, seed, fallback in zip( - query_points, simplex_indexes, outside_fallback_indexes - ) - ] - return tuple(np.stack(values) for values in zip(*rows)) + numpy_kwargs = dict( + query_points=query_points, + points=points, + simplices_padded=simplices_padded, + simplex_neighbors=simplex_neighbors, + simplex_indexes=simplex_indexes, + outside_fallback_indexes=outside_fallback_indexes, + circumcircles=circumcircles, + max_cavity_triangles=max_cavity_triangles, + max_neighbors=max_neighbors, + ) + if _sibson_use_numba() and _sibson_numba_eligible(query_points, points): + return _sibson_mappings_weights_numba(**numpy_kwargs) + return _sibson_mappings_weights_numpy_reference(**numpy_kwargs) import jax diff --git a/test_autoarray/inversion/pixelization/interpolator/test_sibson_numba.py b/test_autoarray/inversion/pixelization/interpolator/test_sibson_numba.py new file mode 100644 index 000000000..d1c2780d9 --- /dev/null +++ b/test_autoarray/inversion/pixelization/interpolator/test_sibson_numba.py @@ -0,0 +1,185 @@ +"""Parity of the numba Sibson kernel against the per-query NumPy reference. + +``sibson_numba_kernel`` is a port of ``_sibson_single_from_tables``; the +integer outputs and diagnostics must agree exactly and the weights to rounding +(the reference groups equal vertex ids with an unstable ``argsort``, so the +order contributions are summed in can differ). +""" + +import numpy as np +import pytest + +pytest.importorskip("numba") + +from autoarray.inversion.mesh.interpolator import sibson +from autoarray.inversion.mesh.interpolator.delaunay import ( + pix_indexes_delaunay_walk_from, + scipy_delaunay_tri_only, +) + + +def _tables(points, query): + simplices, neighbors, vertex_simplex = scipy_delaunay_tri_only(points) + delaunay_mappings, simplex_indexes = pix_indexes_delaunay_walk_from( + query_points=query, + points=points, + simplices_padded=simplices, + simplex_neighbors=neighbors, + vertex_simplex=vertex_simplex, + xp=np, + return_simplex_indexes=True, + ) + return dict( + query_points=query, + points=points, + simplices_padded=simplices, + simplex_neighbors=neighbors, + simplex_indexes=simplex_indexes, + outside_fallback_indexes=delaunay_mappings[:, 0], + circumcircles=sibson.delaunay_circumcircles_from(points, simplices, xp=np), + ) + + +def _reference_and_numba(points, query, max_cavity_triangles=32, max_neighbors=32): + kwargs = _tables(points, query) + kwargs.update( + max_cavity_triangles=max_cavity_triangles, max_neighbors=max_neighbors + ) + reference = sibson._sibson_mappings_weights_numpy_reference(**kwargs) + numba_outputs = sibson._sibson_mappings_weights_numba(**kwargs) + return reference, numba_outputs + + +def _assert_parity(reference, numba_outputs): + names = ("mappings", "sizes", "weights", "cavity_sizes", "overflow", "degenerate") + assert len(reference) == len(numba_outputs) == 6 + for name, expected, actual in zip(names, reference, numba_outputs): + assert expected.dtype == actual.dtype, name + assert expected.shape == actual.shape, name + if name == "weights": + assert np.array_equal(np.isnan(expected), np.isnan(actual)) + np.testing.assert_allclose( + actual, expected, rtol=1.0e-12, atol=1.0e-12, equal_nan=True + ) + else: + assert np.array_equal(expected, actual), name + + +@pytest.mark.parametrize( + "seed, point_count, query_count, query_extent", + [ + (0, 20, 200, 1.0), + (1, 300, 1500, 1.3), + (2, 100, 600, 0.8), + (3, 1000, 1500, 1.2), + ], +) +def test__random_meshes__match_reference(seed, point_count, query_count, query_extent): + rng = np.random.default_rng(seed) + points = rng.uniform(-1.0, 1.0, size=(point_count, 2)) + query = rng.uniform(-query_extent, query_extent, size=(query_count, 2)) + # Queries exactly on vertices and on edge midpoints of the mesh. + simplices, _, _ = scipy_delaunay_tri_only(points) + simplices = simplices[(simplices >= 0).all(axis=1)][:20] + edge_midpoints = 0.5 * (points[simplices[:, 0]] + points[simplices[:, 1]]) + query = np.concatenate([query, points[:20], edge_midpoints]) + + reference, numba_outputs = _reference_and_numba(points, query) + + _assert_parity(reference, numba_outputs) + if query_extent > 1.0: + assert (reference[3] == 0).any() # some queries fall outside the hull + assert (reference[1] == 1).any() # on-vertex / outside fallback rows + assert (reference[1] == 3).any() + + +def test__cocircular_grid__on_vertex_on_edge_and_collinear_queries(): + grid = np.stack(np.meshgrid(np.arange(6.0), np.arange(6.0)), axis=-1) + points = grid.reshape(-1, 2) + simplices, _, _ = scipy_delaunay_tri_only(points) + + query = [points[7], points[14], [2.5, 2.5], [-1.0, 2.0]] + for triangle in simplices: + if (triangle < 0).any(): + continue + for i in range(3): + a, b = points[triangle[i]], points[triangle[(i + 1) % 3]] + for t in (-0.3, 0.25, 0.5, 1.3): + query.append(a + t * (b - a)) + query = np.asarray(query) + + reference, numba_outputs = _reference_and_numba(points, query) + + _assert_parity(reference, numba_outputs) + assert reference[1][0] == 1 and reference[0][0, 0] == 7 + + +def test__thin_triangle__degenerate_watson_circle_matches_reference(): + points = np.array( + [ + [0.0, 0.0], + [1.0, 0.0], + [0.5, 0.01], + [0.5, -100.0], + [0.5, 100.0], + [-50.0, 0.0], + [50.0, 0.0], + ] + ) + query = np.array( + [ + [0.3, 1.0e-15], + [0.3, 0.0], + [0.0, 0.0], + [0.3, 0.004], + [0.7, 1.0e-15], + [0.3, 1.0e-12], + ] + ) + + reference, numba_outputs = _reference_and_numba(points, query) + + _assert_parity(reference, numba_outputs) + assert reference[5].tolist() == [True, False, False, False, True, False] + assert np.isnan(numba_outputs[2][0]).all() + + +@pytest.mark.parametrize( + "max_cavity_triangles, max_neighbors", + [(1, 3), (2, 32), (32, 4), (4, 4), (3, 40)], +) +def test__cavity_and_neighbour_overflow__match_reference( + max_cavity_triangles, max_neighbors +): + rng = np.random.default_rng(5) + points = rng.uniform(-1.0, 1.0, size=(200, 2)) + query = np.concatenate([rng.uniform(-1.1, 1.1, size=(600, 2)), points[:10]]) + + reference, numba_outputs = _reference_and_numba( + points, + query, + max_cavity_triangles=max_cavity_triangles, + max_neighbors=max_neighbors, + ) + + _assert_parity(reference, numba_outputs) + assert reference[4].any() + + +def test__dispatch__numpy_branch_uses_kernel_or_reference(monkeypatch): + rng = np.random.default_rng(7) + points = rng.uniform(-1.0, 1.0, size=(80, 2)) + query = rng.uniform(-1.1, 1.1, size=(300, 2)) + kwargs = _tables(points, query) + del kwargs["circumcircles"] + + reference, numba_outputs = _reference_and_numba(points, query) + + dispatched = sibson.sibson_mappings_weights_from_tables(xp=np, **kwargs) + for expected, actual in zip(numba_outputs, dispatched): + assert np.array_equal(expected, actual, equal_nan=True) + + monkeypatch.setattr(sibson, "_sibson_use_numba", lambda: False) + dispatched = sibson.sibson_mappings_weights_from_tables(xp=np, **kwargs) + for expected, actual in zip(reference, dispatched): + assert np.array_equal(expected, actual, equal_nan=True)