From ab5cc66d2ad5de75d0b688b5336267bd40fb7f8b Mon Sep 17 00:00:00 2001 From: Codex Date: Sun, 9 Aug 2026 22:49:45 +0200 Subject: [PATCH] Add autodifferentiable DelaunayNN mesh --- autoarray/__init__.py | 2 + .../inversion/mesh/interpolator/delaunay.py | 30 +- .../inversion/mesh/interpolator/sibson.py | 830 ++++++++++++++++++ autoarray/inversion/mesh/mesh/__init__.py | 1 + autoarray/inversion/mesh/mesh/abstract.py | 7 +- autoarray/inversion/mesh/mesh/delaunay.py | 11 +- autoarray/inversion/mesh/mesh/delaunay_nn.py | 67 ++ autoarray/inversion/pixelization.py | 5 +- .../inversion/regularization/abstract.py | 3 +- .../inversion/regularization/constant.py | 2 +- .../interpolator/test_delaunay_nn.py | 116 +++ .../pixelization/interpolator/test_sibson.py | 177 ++++ .../test_split_regularization_support.py | 4 + 13 files changed, 1237 insertions(+), 18 deletions(-) create mode 100644 autoarray/inversion/mesh/interpolator/sibson.py create mode 100644 autoarray/inversion/mesh/mesh/delaunay_nn.py create mode 100644 test_autoarray/inversion/pixelization/interpolator/test_delaunay_nn.py create mode 100644 test_autoarray/inversion/pixelization/interpolator/test_sibson.py diff --git a/autoarray/__init__.py b/autoarray/__init__.py index 08f500c8b..40bdadce5 100644 --- a/autoarray/__init__.py +++ b/autoarray/__init__.py @@ -41,6 +41,7 @@ from .inversion.mesh.mesh.abstract import AbstractMesh from .inversion.mesh.interpolator.rectangular import InterpolatorRectangular from .inversion.mesh.interpolator.delaunay import InterpolatorDelaunay +from .inversion.mesh.interpolator.sibson import InterpolatorDelaunayNN from .inversion.inversion.imaging.mapping import InversionImagingMapping from .inversion.inversion.imaging.sparse import InversionImagingSparse from .inversion.inversion.imaging.inversion_imaging_util import ImagingSparseOperator @@ -79,6 +80,7 @@ from .inversion.mesh.mesh_geometry.rectangular import MeshGeometryRectangular from .inversion.mesh.mesh_geometry.delaunay import MeshGeometryDelaunay from .inversion.mesh.interpolator.delaunay import InterpolatorDelaunay +from .inversion.mesh.interpolator.sibson import InterpolatorDelaunayNN from .operators.convolver import Convolver from .operators.interp_2d import interp_2d from .structures.vectors.uniform import VectorYX2D diff --git a/autoarray/inversion/mesh/interpolator/delaunay.py b/autoarray/inversion/mesh/interpolator/delaunay.py index 3c2e36f8d..14621c38c 100644 --- a/autoarray/inversion/mesh/interpolator/delaunay.py +++ b/autoarray/inversion/mesh/interpolator/delaunay.py @@ -133,8 +133,9 @@ def _jax_delaunay_tables(points): NOT an approximation: the callback returns only int32 connectivity tables, which are piecewise-constant in the vertex positions — their true derivative is exactly zero everywhere except the measure-zero - re-wiring (triangle-flip) events, where the likelihood itself is - discontinuous and no gradient exists for any method. Every quantity + re-wiring (triangle-flip) events. The barycentric interpolant is + discontinuous there; the Sibson interpolant in ``DelaunayNN`` instead has + matching limits across an ordinary flip. Every quantity with a non-zero derivative (point location via the visibility walk, barycentric weights, dual areas, split points) is computed in-graph from the traced ``points``, so the frozen-tables gradient is the exact @@ -164,6 +165,7 @@ def pix_indexes_delaunay_walk_from( simplex_neighbors, vertex_simplex, xp=np, + return_simplex_indexes=False, ): """JAX/NumPy point location replacing ``scipy.spatial.Delaunay.find_simplex`` on the JAX likelihood path, via the same visibility-walk algorithm @@ -185,6 +187,11 @@ def pix_indexes_delaunay_walk_from( vmap (JAX path; the NumPy path — used by the unit tests — processes the whole array with early exit). Returns a (Q, 3) int32 mapping array with the same semantics as ``pix_indexes_for_sub_slim_index_delaunay_from``. + + When ``return_simplex_indexes`` is true, also return the containing + simplex index for every query (or -1 outside the convex hull). Sibson + interpolation uses this as the seed of its circumcircle-cavity walk, so + point location is not repeated. """ def cross(u, v): @@ -251,10 +258,15 @@ def locate_chunk(q_chunk): verts = simplices_padded[cur] fallback = xp.stack([seed, -xp.ones_like(seed), -xp.ones_like(seed)], axis=1) - return xp.where(done[:, None], verts, fallback).astype(xp.int32) + mappings = xp.where(done[:, None], verts, fallback).astype(xp.int32) + simplex_indexes = xp.where(done, cur, -1).astype(xp.int32) + return mappings, simplex_indexes if xp is np: - return locate_chunk(query_points) + mappings, simplex_indexes = locate_chunk(query_points) + if return_simplex_indexes: + return mappings, simplex_indexes + return mappings import jax @@ -265,8 +277,14 @@ def locate_chunk(q_chunk): q_padded = xp.concatenate( [query_points, xp.full((pad, 2), 1.0e9, dtype=query_points.dtype)] ) - mappings = jax.lax.map(locate_chunk, q_padded.reshape(-1, chunk, 2)).reshape(-1, 3) - return mappings[:Q] + mappings, simplex_indexes = jax.lax.map( + locate_chunk, q_padded.reshape(-1, chunk, 2) + ) + mappings = mappings.reshape(-1, 3)[:Q] + simplex_indexes = simplex_indexes.reshape(-1)[:Q] + if return_simplex_indexes: + return mappings, simplex_indexes + return mappings def jax_delaunay(points, query_points, areas_factor=0.5): diff --git a/autoarray/inversion/mesh/interpolator/sibson.py b/autoarray/inversion/mesh/interpolator/sibson.py new file mode 100644 index 000000000..65a494734 --- /dev/null +++ b/autoarray/inversion/mesh/interpolator/sibson.py @@ -0,0 +1,830 @@ +"""NumPy and JAX implementations of Sibson natural-neighbour interpolation. + +The Delaunay connectivity still comes from the small qhull callback used by +``delaunay.py``. Everything depending continuously on the mesh and query +coordinates -- circumcircles, cavity traversal, Watson areas, normalization +and stencil compaction -- is array code and therefore differentiable by JAX. + +``InterpolatorDelaunayNN`` exposes the implementation through the public +``mesh.DelaunayNN`` mesh. Overflow and degeneracy diagnostics remain part of +the internal interface: a fixed-shape limit failure produces NaN weights so a +model sample is rejected instead of silently using a truncated stencil. +""" + +import numpy as np +from autonerves import cached_property + +from autoarray.inversion.mesh.interpolator.delaunay import ( + DelaunayInterface, + InterpolatorDelaunay, + _jax_delaunay_tables, + barycentric_dual_area_from, + pix_indexes_delaunay_walk_from, + pix_indexes_for_sub_slim_index_delaunay_from, +) +from autoarray.inversion.regularization.regularization_util import ( + split_points_from, +) + +# Fixed-shape production caps. The workspace mass-model audit (101 traced +# Hilbert meshes, including split points) observed maxima of 25 cavity +# triangles and 27 natural neighbours. Caps 16 and 24 overflowed; 32 was the +# smallest tested safe shape. See +# autolens_workspace_test/scripts/misc/jax_assertions/delaunay_nn_caps.py. +SIBSON_MAX_CAVITY_TRIANGLES = 32 +SIBSON_MAX_NEIGHBORS = 32 +SIBSON_QUERY_CHUNK = 256 + + +def _cross(u, v): + return u[..., 0] * v[..., 1] - u[..., 1] * v[..., 0] + + +def circumcircles_from(a, b, c, xp=np): + """Return circumcentres, squared radii and validity for point triples. + + The leading dimensions of ``a``, ``b`` and ``c`` are broadcast normally; + the final dimension is the two Cartesian coordinates. + """ + ba = b - a + ca = c - a + denominator = 2.0 * _cross(ba, ca) + scale = xp.maximum( + xp.maximum(xp.sum(ba * ba, axis=-1), xp.sum(ca * ca, axis=-1)), + xp.asarray(1.0, dtype=a.dtype), + ) + tolerance = 32.0 * xp.finfo(a.dtype).eps * scale + valid = xp.abs(denominator) > tolerance + safe_denominator = xp.where(valid, denominator, 1.0) + + ba2 = xp.sum(ba * ba, axis=-1) + ca2 = xp.sum(ca * ca, axis=-1) + offset = ( + xp.stack( + [ + ba2 * ca[..., 1] - ca2 * ba[..., 1], + ba[..., 0] * ca2 - ca[..., 0] * ba2, + ], + axis=-1, + ) + / safe_denominator[..., None] + ) + centre = a + offset + radius_squared = xp.sum(offset * offset, axis=-1) + + centre = xp.where(valid[..., None], centre, 0.0) + radius_squared = xp.where(valid, radius_squared, 0.0) + return centre, radius_squared, valid + + +def delaunay_circumcircles_from(points, simplices_padded, xp=np): + """Circumcircles for fixed-shape, -1-padded Delaunay simplices.""" + simplex_valid = (simplices_padded >= 0).all(axis=1) + safe_simplices = simplices_padded.clip(min=0) + vertices = points[safe_simplices] + centres, radii_squared, circle_valid = circumcircles_from( + vertices[:, 0], vertices[:, 1], vertices[:, 2], xp=xp + ) + valid = simplex_valid & circle_valid + centres = xp.where(valid[:, None], centres, 0.0) + radii_squared = xp.where(valid, radii_squared, 0.0) + return centres, radii_squared, valid + + +def _contains_query( + triangle_index, + query, + circumcentres, + circumradii_squared, + circumcircle_valid, + xp, +): + safe_index = xp.maximum(triangle_index, 0) + delta = circumcentres[safe_index] - query + distance_squared = xp.sum(delta * delta) + radius_squared = circumradii_squared[safe_index] + tolerance = 64.0 * xp.finfo(query.dtype).eps * xp.maximum(radius_squared, 1.0) + return ( + (triangle_index >= 0) + & circumcircle_valid[safe_index] + & (distance_squared <= radius_squared + tolerance) + ) + + +def _cavity_triangle_indexes_numpy( + query, + seed_simplex, + simplex_neighbors, + circumcentres, + circumradii_squared, + circumcircle_valid, + max_cavity_triangles, +): + cavity = np.full(max_cavity_triangles, -1, dtype=np.int32) + if seed_simplex < 0: + return cavity, np.int32(0), np.bool_(False) + + cavity[0] = seed_simplex + count = 1 + overflow = False + + for position in range(max_cavity_triangles): + if position >= count: + break + triangle_index = cavity[position] + for candidate in simplex_neighbors[triangle_index]: + if candidate < 0 or candidate in cavity[:count]: + continue + if not _contains_query( + candidate, + query, + circumcentres, + circumradii_squared, + circumcircle_valid, + np, + ): + continue + if count == max_cavity_triangles: + overflow = True + else: + cavity[count] = candidate + count += 1 + + return cavity, np.int32(count), np.bool_(overflow) + + +def _cavity_triangle_indexes_jax( + query, + seed_simplex, + simplex_neighbors, + circumcentres, + circumradii_squared, + circumcircle_valid, + max_cavity_triangles, +): + import jax + import jax.numpy as jnp + + safe_seed = jnp.maximum(seed_simplex, 0) + cavity = -jnp.ones((max_cavity_triangles,), dtype=jnp.int32) + cavity = cavity.at[0].set(safe_seed) + initial_count = jnp.where(seed_simplex >= 0, 1, 0).astype(jnp.int32) + + def process_triangle(position, carry): + cavity, count, overflow = carry + active = position < count + triangle_index = cavity[position].clip(min=0) + candidates = simplex_neighbors[triangle_index] + + def add_candidate(edge, inner_carry): + cavity, count, overflow = inner_carry + candidate = candidates[edge] + unseen = ~jnp.any(cavity == candidate) + contained = _contains_query( + candidate, + query, + circumcentres, + circumradii_squared, + circumcircle_valid, + jnp, + ) + accepted = active & unseen & contained + has_space = count < max_cavity_triangles + add = accepted & has_space + write_index = jnp.minimum(count, max_cavity_triangles - 1) + old_value = cavity[write_index] + cavity = cavity.at[write_index].set(jnp.where(add, candidate, old_value)) + count = count + add.astype(jnp.int32) + overflow = overflow | (accepted & ~has_space) + return cavity, count, overflow + + return jax.lax.fori_loop(0, 3, add_candidate, (cavity, count, overflow)) + + return jax.lax.fori_loop( + 0, + max_cavity_triangles, + process_triangle, + (cavity, initial_count, jnp.asarray(False)), + ) + + +def _sibson_single_from_tables( + query, + seed_simplex, + outside_fallback_index, + points, + simplices_padded, + simplex_neighbors, + circumcentres, + circumradii_squared, + circumcircle_valid, + max_cavity_triangles, + max_neighbors, + xp, +): + if xp is np: + cavity, cavity_size, overflow = _cavity_triangle_indexes_numpy( + query=query, + seed_simplex=int(seed_simplex), + simplex_neighbors=simplex_neighbors, + circumcentres=circumcentres, + circumradii_squared=circumradii_squared, + circumcircle_valid=circumcircle_valid, + max_cavity_triangles=max_cavity_triangles, + ) + else: + cavity, cavity_size, overflow = _cavity_triangle_indexes_jax( + query=query, + seed_simplex=seed_simplex, + simplex_neighbors=simplex_neighbors, + circumcentres=circumcentres, + circumradii_squared=circumradii_squared, + circumcircle_valid=circumcircle_valid, + max_cavity_triangles=max_cavity_triangles, + ) + + cavity_valid = cavity >= 0 + safe_cavity = cavity.clip(min=0) + triangle_vertices = simplices_padded[safe_cavity].clip(min=0) + vertices = points[triangle_vertices] + + # Circle j passes through the query and the edge opposite triangle vertex j. + edge_a = vertices[:, [1, 2, 0], :] + edge_b = vertices[:, [2, 0, 1], :] + query_for_edges = xp.broadcast_to(query, edge_a.shape) + inserted_centres, _, inserted_valid = circumcircles_from( + edge_a, edge_b, query_for_edges, xp=xp + ) + + original_centres = circumcentres[safe_cavity][:, None, :] + relative = inserted_centres - original_centres + contributions = xp.stack( + [ + _cross(relative[:, 1], relative[:, 2]), + _cross(relative[:, 2], relative[:, 0]), + _cross(relative[:, 0], relative[:, 1]), + ], + axis=1, + ) + contribution_valid = cavity_valid[:, None] & xp.stack( + [ + inserted_valid[:, 1] & inserted_valid[:, 2], + inserted_valid[:, 2] & inserted_valid[:, 0], + inserted_valid[:, 0] & inserted_valid[:, 1], + ], + axis=1, + ) + contributions = xp.where(contribution_valid, contributions, 0.0) + + flat_vertices = triangle_vertices.reshape(-1) + flat_contributions = contributions.reshape(-1) + flat_valid = xp.broadcast_to(cavity_valid[:, None], contributions.shape).reshape(-1) + sentinel = points.shape[0] + sortable_vertices = xp.where(flat_valid, flat_vertices, sentinel) + sortable_contributions = xp.where(flat_valid, flat_contributions, 0.0) + + # Compact the at-most 3*C Watson contributions without allocating an + # N-vertex vector per query. Equal vertex ids become adjacent after the + # small fixed-size sort and are scatter-added into consecutive groups. + # This is substantially cheaper than a global N-way top-k for production + # meshes (N ~= 1200, while 3*C <= 96). + order = xp.argsort(sortable_vertices) + sorted_vertices = sortable_vertices[order] + sorted_contributions = sortable_contributions[order] + sorted_valid = sorted_vertices < sentinel + group_start = sorted_valid & xp.concatenate( + [ + xp.ones((1,), dtype=bool), + sorted_vertices[1:] != sorted_vertices[:-1], + ] + ) + group = xp.cumsum(group_start.astype(xp.int32)) - 1 + unique_count = xp.sum(group_start).astype(xp.int32) + + compact_weights = xp.zeros(flat_vertices.shape[0], dtype=points.dtype) + if xp is np: + np.add.at( + compact_weights, + group, + np.where(sorted_valid, sorted_contributions, 0.0), + ) + start_positions = np.flatnonzero(group_start) + compact_indexes = np.full(flat_vertices.shape[0], -1, dtype=np.int32) + compact_indexes[: len(start_positions)] = sorted_vertices[start_positions] + else: + compact_weights = compact_weights.at[group].add( + xp.where(sorted_valid, sorted_contributions, 0.0) + ) + start_positions = xp.nonzero( + group_start, size=flat_vertices.shape[0], fill_value=0 + )[0] + compact_indexes = xp.where( + xp.arange(flat_vertices.shape[0]) < unique_count, + sorted_vertices[start_positions], + -1, + ).astype(xp.int32) + + inside = seed_simplex >= 0 + seed_vertices = simplices_padded[xp.maximum(seed_simplex, 0)].clip(min=0) + seed_points = points[seed_vertices] + seed_distances_squared = xp.sum((seed_points - query[None, :]) ** 2, axis=1) + local_vertex = seed_vertices[xp.argmin(seed_distances_squared)] + coordinate_scale = xp.maximum( + xp.maximum(xp.max(xp.abs(query)), xp.max(xp.abs(points[seed_vertices]))), + 1.0, + ) + vertex_tolerance = (64.0 * xp.finfo(points.dtype).eps * coordinate_scale) ** 2 + on_vertex = inside & (xp.min(seed_distances_squared) <= vertex_tolerance) + + seed_denominator = _cross( + seed_points[1] - seed_points[0], seed_points[2] - seed_points[0] + ) + safe_seed_denominator = xp.where(seed_denominator != 0.0, seed_denominator, 1.0) + seed_barycentric = ( + xp.stack( + [ + _cross(seed_points[1] - query, seed_points[2] - query), + _cross(seed_points[2] - query, seed_points[0] - query), + _cross(seed_points[0] - query, seed_points[1] - query), + ] + ) + / safe_seed_denominator + ) + edge_tolerance = 128.0 * xp.finfo(points.dtype).eps + on_edge = ( + inside & (~on_vertex) & (xp.min(xp.abs(seed_barycentric)) <= edge_tolerance) + ) + + fallback_index = xp.where(on_vertex, local_vertex, outside_fallback_index) + use_fallback = (~inside) | on_vertex + + weight_sum = xp.sum(compact_weights) + degenerate = ( + inside + & (~on_vertex) + & (~on_edge) + & ( + xp.any(cavity_valid[:, None] & ~inserted_valid) + | (~xp.isfinite(weight_sum)) + | (xp.abs(weight_sum) <= xp.finfo(points.dtype).tiny) + ) + ) + safe_sum = xp.where( + xp.abs(weight_sum) > xp.finfo(points.dtype).tiny, weight_sum, 1.0 + ) + compact_weights = compact_weights / safe_sum + + compact_length = flat_vertices.shape[0] + if max_neighbors <= compact_length: + mappings = compact_indexes[:max_neighbors] + weights = compact_weights[:max_neighbors] + else: + padding = max_neighbors - compact_length + mappings = xp.pad(compact_indexes, (0, padding), constant_values=-1) + weights = xp.pad(compact_weights, (0, padding)) + + selected_valid = (xp.arange(max_neighbors) < unique_count) & (weights > 0.0) + mappings = xp.where(selected_valid, mappings, -1).astype(xp.int32) + weights = xp.where(selected_valid, weights, 0.0) + + fallback_mappings = -xp.ones((max_neighbors,), dtype=xp.int32) + fallback_weights = xp.zeros((max_neighbors,), dtype=points.dtype) + if xp is np: + fallback_mappings[0] = int(fallback_index) + fallback_weights[0] = 1.0 + else: + fallback_mappings = fallback_mappings.at[0].set(fallback_index) + fallback_weights = fallback_weights.at[0].set(1.0) + mappings = xp.where(use_fallback, fallback_mappings, mappings) + weights = xp.where(use_fallback, fallback_weights, weights) + size = xp.where(use_fallback, 1, xp.sum(selected_valid)).astype(xp.int32) + + # Inserting a query exactly on an existing edge makes one of the Watson + # circumcircles collinear. Sibson coordinates there are the ordinary + # linear coordinates of the two edge endpoints, which the containing + # triangle's barycentric row represents exactly (the third weight is 0). + edge_mappings = -xp.ones((max_neighbors,), dtype=xp.int32) + edge_weights = xp.zeros((max_neighbors,), dtype=points.dtype) + if xp is np: + edge_mappings[:3] = seed_vertices + edge_weights[:3] = seed_barycentric + else: + edge_mappings = edge_mappings.at[:3].set(seed_vertices) + edge_weights = edge_weights.at[:3].set(seed_barycentric) + mappings = xp.where(on_edge, edge_mappings, mappings) + weights = xp.where(on_edge, edge_weights, weights) + size = xp.where(on_edge, 3, size).astype(xp.int32) + + neighbor_overflow = ( + inside & (~on_vertex) & (~on_edge) & (unique_count > max_neighbors) + ) + overflow = overflow | neighbor_overflow + failed = overflow | degenerate + weights = xp.where(failed, xp.nan, weights) + return mappings, size, weights, cavity_size, overflow, degenerate + + +def sibson_mappings_weights_from_tables( + query_points, + points, + simplices_padded, + simplex_neighbors, + simplex_indexes, + outside_fallback_indexes, + max_cavity_triangles=SIBSON_MAX_CAVITY_TRIANGLES, + max_neighbors=SIBSON_MAX_NEIGHBORS, + query_chunk=SIBSON_QUERY_CHUNK, + xp=np, +): + """Calculate fixed-shape Sibson mappings and weights from Delaunay tables. + + Returns + ------- + mappings, sizes, weights + Mapper-compatible arrays with at most ``max_neighbors`` unique source + vertices per query. + cavity_sizes, cavity_overflow, degenerate + Prototype diagnostics. Overflow or a Watson edge degeneracy produces + NaN weights rather than silently returning an approximate stencil. + """ + circumcentres, circumradii_squared, circumcircle_valid = ( + delaunay_circumcircles_from(points, simplices_padded, xp=xp) + ) + + def single(query, seed_simplex, outside_fallback_index): + return _sibson_single_from_tables( + query=query, + seed_simplex=seed_simplex, + outside_fallback_index=outside_fallback_index, + 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=xp, + ) + + 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)) + + import jax + + query_count = query_points.shape[0] + pad = (-query_count) % query_chunk + padded_queries = xp.concatenate( + [query_points, xp.zeros((pad, 2), dtype=query_points.dtype)] + ) + padded_simplex_indexes = xp.concatenate( + [simplex_indexes, -xp.ones((pad,), dtype=xp.int32)] + ) + padded_fallback_indexes = xp.concatenate( + [outside_fallback_indexes, xp.zeros((pad,), dtype=xp.int32)] + ) + + chunk_inputs = ( + padded_queries.reshape(-1, query_chunk, 2), + padded_simplex_indexes.reshape(-1, query_chunk), + padded_fallback_indexes.reshape(-1, query_chunk), + ) + batched_single = jax.vmap(single) + outputs = jax.lax.map(lambda args: batched_single(*args), chunk_inputs) + return tuple( + output.reshape((-1,) + output.shape[2:])[:query_count] for output in outputs + ) + + +def jax_sibson( + points, + query_points, + max_cavity_triangles=SIBSON_MAX_CAVITY_TRIANGLES, + max_neighbors=SIBSON_MAX_NEIGHBORS, + query_chunk=SIBSON_QUERY_CHUNK, +): + """Prototype end-to-end JAX Sibson interpolation-table construction. + + Qhull supplies only integer connectivity through ``pure_callback``. The + returned floating-point weights retain gradients with respect to both + ``points`` and ``query_points`` everywhere the natural-neighbour geometry + is differentiable. + """ + import jax.numpy as jnp + + simplices_padded, simplex_neighbors, vertex_simplex = _jax_delaunay_tables(points) + delaunay_mappings, simplex_indexes = pix_indexes_delaunay_walk_from( + query_points=query_points, + points=points, + simplices_padded=simplices_padded, + simplex_neighbors=simplex_neighbors, + vertex_simplex=vertex_simplex, + xp=jnp, + return_simplex_indexes=True, + ) + outputs = sibson_mappings_weights_from_tables( + query_points=query_points, + points=points, + simplices_padded=simplices_padded, + simplex_neighbors=simplex_neighbors, + simplex_indexes=simplex_indexes, + outside_fallback_indexes=delaunay_mappings[:, 0], + max_cavity_triangles=max_cavity_triangles, + max_neighbors=max_neighbors, + query_chunk=query_chunk, + xp=jnp, + ) + return (points, simplices_padded) + outputs + + +def scipy_delaunay_nn( + points_np, + query_points_np, + areas_factor=0.5, + max_cavity_triangles=SIBSON_MAX_CAVITY_TRIANGLES, + max_neighbors=SIBSON_MAX_NEIGHBORS, + query_chunk=SIBSON_QUERY_CHUNK, +): + """Build all NumPy tables required by ``InterpolatorDelaunayNN``.""" + from scipy.spatial import Delaunay + + triangle = Delaunay(points_np) + points = triangle.points.astype(points_np.dtype) + simplex_count = triangle.simplices.shape[0] + + simplices_padded = -np.ones((2 * points.shape[0], 3), dtype=np.int32) + neighbors_padded = -np.ones_like(simplices_padded) + simplices_padded[:simplex_count] = triangle.simplices.astype(np.int32) + neighbors_padded[:simplex_count] = triangle.neighbors.astype(np.int32) + + def mappings_weights_for(query_points): + simplex_indexes = triangle.find_simplex(query_points).astype(np.int32) + delaunay_mappings = pix_indexes_for_sub_slim_index_delaunay_from( + data_grid=query_points, + simplex_index_for_sub_slim_index=simplex_indexes, + pix_indexes_for_simplex_index=triangle.simplices, + delaunay_points=points, + ) + return sibson_mappings_weights_from_tables( + query_points=query_points, + points=points, + simplices_padded=simplices_padded, + simplex_neighbors=neighbors_padded, + simplex_indexes=simplex_indexes, + outside_fallback_indexes=delaunay_mappings[:, 0], + max_cavity_triangles=max_cavity_triangles, + max_neighbors=max_neighbors, + query_chunk=query_chunk, + xp=np, + ) + + mappings, sizes, weights, cavity_sizes, overflow, degenerate = mappings_weights_for( + query_points_np + ) + + areas = barycentric_dual_area_from( + mesh_grid=points, + simplices=triangle.simplices, + xp=np, + ) + split_points = split_points_from( + points=points, + area_weights=areas_factor * np.sqrt(areas), + xp=np, + ) + ( + splitted_mappings, + splitted_sizes, + splitted_weights, + split_cavity_sizes, + split_overflow, + split_degenerate, + ) = mappings_weights_for(split_points) + + return ( + points, + simplices_padded, + mappings, + sizes, + weights, + split_points, + splitted_mappings, + splitted_sizes, + splitted_weights, + cavity_sizes, + overflow, + degenerate, + split_cavity_sizes, + split_overflow, + split_degenerate, + ) + + +def jax_delaunay_nn( + points, + query_points, + areas_factor=0.5, + max_cavity_triangles=SIBSON_MAX_CAVITY_TRIANGLES, + max_neighbors=SIBSON_MAX_NEIGHBORS, + query_chunk=SIBSON_QUERY_CHUNK, +): + """Build all JAX tables required by ``InterpolatorDelaunayNN``. + + Qhull returns only fixed-shape integer connectivity through a stopped + ``pure_callback``. Point location, circumcircles, Sibson weights, dual + areas, and split-cross coordinates all remain in the JAX graph. + """ + import jax.numpy as jnp + + simplices_padded, simplex_neighbors, vertex_simplex = _jax_delaunay_tables(points) + + def mappings_weights_for(query): + delaunay_mappings, simplex_indexes = pix_indexes_delaunay_walk_from( + query_points=query, + points=points, + simplices_padded=simplices_padded, + simplex_neighbors=simplex_neighbors, + vertex_simplex=vertex_simplex, + xp=jnp, + return_simplex_indexes=True, + ) + return sibson_mappings_weights_from_tables( + query_points=query, + points=points, + simplices_padded=simplices_padded, + simplex_neighbors=simplex_neighbors, + simplex_indexes=simplex_indexes, + outside_fallback_indexes=delaunay_mappings[:, 0], + max_cavity_triangles=max_cavity_triangles, + max_neighbors=max_neighbors, + query_chunk=query_chunk, + xp=jnp, + ) + + mappings, sizes, weights, cavity_sizes, overflow, degenerate = mappings_weights_for( + query_points + ) + + valid = simplices_padded[:, 0] >= 0 + simplices = simplices_padded.clip(min=0) + p0 = points[simplices[:, 0]] + p1 = points[simplices[:, 1]] + p2 = points[simplices[:, 2]] + triangle_cross = (p1[:, 0] - p0[:, 0]) * (p2[:, 1] - p0[:, 1]) - ( + p1[:, 1] - p0[:, 1] + ) * (p2[:, 0] - p0[:, 0]) + contribution = jnp.where(valid, 0.5 * jnp.abs(triangle_cross) / 3.0, 0.0) + areas = jnp.zeros(points.shape[0], dtype=points.dtype) + for vertex_index in range(3): + areas = areas.at[simplices[:, vertex_index]].add(contribution) + + split_points = split_points_from( + points=points, + area_weights=areas_factor * jnp.sqrt(areas), + xp=jnp, + ) + ( + splitted_mappings, + splitted_sizes, + splitted_weights, + split_cavity_sizes, + split_overflow, + split_degenerate, + ) = mappings_weights_for(split_points) + + return ( + points, + simplices_padded, + mappings, + sizes, + weights, + split_points, + splitted_mappings, + splitted_sizes, + splitted_weights, + cavity_sizes, + overflow, + degenerate, + split_cavity_sizes, + split_overflow, + split_degenerate, + ) + + +class DelaunayNNInterface(DelaunayInterface): + """Fixed-shape natural-neighbour interpolation tables and diagnostics.""" + + def __init__( + self, + points, + simplices, + mappings, + sizes, + weights, + split_points, + splitted_mappings, + splitted_sizes, + splitted_weights, + cavity_sizes, + overflow, + degenerate, + split_cavity_sizes, + split_overflow, + split_degenerate, + xp=np, + ): + super().__init__( + points=points, + simplices=simplices, + mappings=mappings, + split_points=split_points, + splitted_mappings=splitted_mappings, + xp=xp, + ) + self._sizes = sizes + self.weights = weights + self._splitted_sizes = splitted_sizes + self.splitted_weights = splitted_weights + self.cavity_sizes = cavity_sizes + self.overflow = overflow + self.degenerate = degenerate + self.split_cavity_sizes = split_cavity_sizes + self.split_overflow = split_overflow + self.split_degenerate = split_degenerate + + @cached_property + def sizes(self): + return self._sizes.astype(np.int32) + + @cached_property + def splitted_sizes(self): + return self._splitted_sizes.astype(np.int32) + + +class InterpolatorDelaunayNN(InterpolatorDelaunay): + """Delaunay-mesh interpolator using Sibson natural-neighbour weights. + + Connectivity is obtained from the same qhull callback as + ``InterpolatorDelaunay``. In JAX mode all floating-point geometry is + autodifferentiable with respect to both mesh and query coordinates while + the piecewise-constant connectivity is frozen. Unlike triangle-local + barycentric interpolation, Sibson coordinates agree through a Delaunay + diagonal flip, removing the finite sign/jump seen when the mass model + changes the triangulation. + """ + + @cached_property + def delaunay(self): + kwargs = { + "areas_factor": self.mesh.areas_factor, + "max_cavity_triangles": self.mesh.max_cavity_triangles, + "max_neighbors": self.mesh.max_neighbors, + "query_chunk": self.mesh.query_chunk, + } + if self._xp.__name__.startswith("jax"): + outputs = jax_delaunay_nn( + points=self.mesh_grid_xy, + query_points=self.data_grid.over_sampled.array, + **kwargs, + ) + else: + outputs = scipy_delaunay_nn( + points_np=self.mesh_grid_xy, + query_points_np=self.data_grid.over_sampled.array, + **kwargs, + ) + + return DelaunayNNInterface(*outputs, xp=self._xp) + + @cached_property + def _mappings_sizes_weights(self): + return ( + self.delaunay.mappings.astype(self._xp.int32), + self.delaunay.sizes.astype(self._xp.int32), + self.delaunay.weights, + ) + + @cached_property + def _mappings_sizes_weights_split(self): + # ``reg_split_from`` may need to insert the centre pixel when it is not + # already in a split-point stencil, so retain one spare padded column. + row_count = self.delaunay.splitted_mappings.shape[0] + mappings = self._xp.hstack( + ( + self.delaunay.splitted_mappings.astype(self._xp.int32), + -self._xp.ones((row_count, 1), dtype=self._xp.int32), + ) + ) + weights = self._xp.hstack( + ( + self.delaunay.splitted_weights, + self._xp.zeros((row_count, 1), dtype=self.delaunay.weights.dtype), + ) + ) + return mappings, self.delaunay.splitted_sizes.astype(self._xp.int32), weights diff --git a/autoarray/inversion/mesh/mesh/__init__.py b/autoarray/inversion/mesh/mesh/__init__.py index b14718292..808df01d2 100644 --- a/autoarray/inversion/mesh/mesh/__init__.py +++ b/autoarray/inversion/mesh/mesh/__init__.py @@ -3,5 +3,6 @@ from .rectangular_adapt_image import RectangularAdaptImage from .rectangular_uniform import RectangularUniform from .delaunay import Delaunay +from .delaunay_nn import DelaunayNN from .knn import KNearestNeighbor from .knn import KNNBarycentric diff --git a/autoarray/inversion/mesh/mesh/abstract.py b/autoarray/inversion/mesh/mesh/abstract.py index 8399ac51a..657398948 100644 --- a/autoarray/inversion/mesh/mesh/abstract.py +++ b/autoarray/inversion/mesh/mesh/abstract.py @@ -15,9 +15,10 @@ class AbstractMesh: ``AdaptSplitZeroth``). Split schemes regularize using a split-cross calculation, which requires the mesh's interpolator to - provide ``_mappings_sizes_weights_split``. The adaptive meshes (``Delaunay``, ``KNNBarycentric``) - compute this; the rectangular meshes do not, and set this to ``False`` so that ``Pixelization`` - rejects the combination at construction rather than failing deep inside the inversion. + provide ``_mappings_sizes_weights_split``. The adaptive meshes (``Delaunay``, + ``DelaunayNN``, ``KNNBarycentric``) compute this; the rectangular meshes do + not, and set this to ``False`` so that ``Pixelization`` rejects the + combination at construction rather than failing deep inside the inversion. """ def __eq__(self, other): diff --git a/autoarray/inversion/mesh/mesh/delaunay.py b/autoarray/inversion/mesh/mesh/delaunay.py index 1bf818c45..d43a47172 100644 --- a/autoarray/inversion/mesh/mesh/delaunay.py +++ b/autoarray/inversion/mesh/mesh/delaunay.py @@ -37,15 +37,16 @@ def __init__( In what sense "autodifferentiable"? The same sense as a ReLU network: the likelihood is piecewise-smooth — perfectly smooth within each triangulation topology, with measure-zero jump discontinuities at the - triangle-flip (re-wiring) boundaries, where no gradient exists for - any method. A sampler almost surely never lands on a seam, and on + triangle-flip (re-wiring) boundaries, where the barycentric interpolant + jumps and has no gradient. A sampler almost surely never lands on a seam, and on either side autodiff returns the exact gradient of the branch the evaluation point is on (FD comparisons show the seams, autodiff does not — individual finite-difference steps can straddle a flip). This contrasts the adaptive rectangular (kernel-CDF) meshes, which are - C-infinity by construction with no seams at all — the cleanest choice - for gradient-based inference, with Delaunay a scientifically exact - piecewise-smooth alternative. + C-infinity by construction with no seams at all. ``DelaunayNN`` is the + local, linearly precise alternative when continuity through Delaunay + flips is required: its Sibson natural-neighbour weights approach the + same value from either triangulation. Caveat for batched samplers: the callback is ``vmap_method="sequential"`` (one host qhull call per vmap lane) — diff --git a/autoarray/inversion/mesh/mesh/delaunay_nn.py b/autoarray/inversion/mesh/mesh/delaunay_nn.py new file mode 100644 index 000000000..31198b0e6 --- /dev/null +++ b/autoarray/inversion/mesh/mesh/delaunay_nn.py @@ -0,0 +1,67 @@ +"""Sibson natural-neighbour interpolation on a Delaunay mesh.""" + +from autoarray.inversion.mesh.interpolator.sibson import ( + SIBSON_MAX_CAVITY_TRIANGLES, + SIBSON_MAX_NEIGHBORS, + SIBSON_QUERY_CHUNK, +) +from autoarray.inversion.mesh.mesh.delaunay import Delaunay + + +class DelaunayNN(Delaunay): + """An irregular source mesh with Sibson natural-neighbour interpolation. + + This mesh has the same vertices, Delaunay connectivity, area calculation, + and split-regularization support as :class:`Delaunay`. Its interpolation + weights differ: every coordinate uses all natural neighbours in the local + circumcircle cavity instead of only the three vertices of its containing + triangle. The resulting interpolant is local and linearly precise, while + remaining continuous when the Delaunay diagonal flips as a mass model + moves the source-plane mesh. + + Under JAX, qhull supplies only integer connectivity via ``pure_callback``. + The callback is stopped under differentiation because that connectivity is + piecewise constant. Circumcircles, natural-neighbour weights, split points, + and the inversion remain in JAX and are autodifferentiable with respect to + the mesh and query coordinates. At an exactly degenerate topology the + connectivity itself still has no derivative, but the Sibson interpolant + approaches the same value through an ordinary diagonal flip. + + The fixed cavity and neighbour caps make mapper shapes static for JIT + compilation. If either cap is exceeded, the affected weights become NaN so + the sample is rejected instead of using a truncated interpolation. The + class constants have headroom over the observed production Hilbert meshes. + + Parameters + ---------- + pixels + Number of active mesh vertices used by the inversion. + zeroed_pixels + Number of final boundary vertices fixed to zero. + areas_factor + Scale applied to the area-derived split-cross offsets. + """ + + max_cavity_triangles = SIBSON_MAX_CAVITY_TRIANGLES + max_neighbors = SIBSON_MAX_NEIGHBORS + query_chunk = SIBSON_QUERY_CHUNK + + def __init__( + self, + pixels: int, + zeroed_pixels: int | None = 0, + areas_factor: float = 0.5, + ): + super().__init__( + pixels=pixels, + zeroed_pixels=zeroed_pixels, + areas_factor=areas_factor, + ) + + @property + def interpolator_cls(self): + from autoarray.inversion.mesh.interpolator.sibson import ( + InterpolatorDelaunayNN, + ) + + return InterpolatorDelaunayNN diff --git a/autoarray/inversion/pixelization.py b/autoarray/inversion/pixelization.py index 1b65400b4..aa9431f01 100644 --- a/autoarray/inversion/pixelization.py +++ b/autoarray/inversion/pixelization.py @@ -169,8 +169,9 @@ def __init__( Use either: - - an adaptive mesh which supports split regularization, e.g. `Delaunay` or - `KNNBarycentric`, with `{type(regularization).__name__}`; or + - an adaptive mesh which supports split regularization, e.g. `Delaunay`, + `DelaunayNN`, or `KNNBarycentric`, with + `{type(regularization).__name__}`; or - a non-split regularization scheme with `{type(mesh).__name__}`, e.g. `Constant` instead of `ConstantSplit`, or `Adapt` instead of `AdaptSplit`. """ diff --git a/autoarray/inversion/regularization/abstract.py b/autoarray/inversion/regularization/abstract.py index c149f1966..bff7cb3a5 100644 --- a/autoarray/inversion/regularization/abstract.py +++ b/autoarray/inversion/regularization/abstract.py @@ -13,7 +13,8 @@ class AbstractRegularization: calculation of the mesh's mappings rather than the mappings themselves. Split schemes require the mesh's interpolator to provide `_mappings_sizes_weights_split`, which - only the adaptive meshes (e.g. ``Delaunay``, ``KNNBarycentric``) do. ``Pixelization`` uses this + only the adaptive meshes (e.g. ``Delaunay``, ``DelaunayNN``, + ``KNNBarycentric``) do. ``Pixelization`` uses this flag together with ``AbstractMesh.supports_split_regularization`` to reject unsupported combinations at construction. """ diff --git a/autoarray/inversion/regularization/constant.py b/autoarray/inversion/regularization/constant.py index c872bb3af..f1e01815f 100644 --- a/autoarray/inversion/regularization/constant.py +++ b/autoarray/inversion/regularization/constant.py @@ -89,7 +89,7 @@ def __init__(self, coefficient: float = 1.0): **JAX & gradient support** (2026-07 gradient sweep): on meshes with an analytic neighbor structure (the rectangular family) this scheme is JAX-differentiable and FD-certified. On the Delaunay mesh family - (``Delaunay``, ``KNearestNeighbor``, ``KNNBarycentric``) the neighbors + (``Delaunay``, ``DelaunayNN``, ``KNearestNeighbor``, ``KNNBarycentric``) the neighbors come from a direct ``scipy.spatial.Delaunay`` call on the traced source-plane mesh grid, so it raises ``TracerArrayConversionError`` under ``jax.jit`` / ``jax.grad`` — use a split-family scheme diff --git a/test_autoarray/inversion/pixelization/interpolator/test_delaunay_nn.py b/test_autoarray/inversion/pixelization/interpolator/test_delaunay_nn.py new file mode 100644 index 000000000..696221ff1 --- /dev/null +++ b/test_autoarray/inversion/pixelization/interpolator/test_delaunay_nn.py @@ -0,0 +1,116 @@ +import numpy as np + +import autoarray as aa +from autoarray.inversion.mesh.interpolator.delaunay import ( + pixel_weights_delaunay_from, + scipy_delaunay, +) +from autoarray.inversion.mesh.interpolator.sibson import ( + InterpolatorDelaunayNN, + scipy_delaunay_nn, +) + + +def test__mesh_is_public_and_selects_natural_neighbor_interpolator(): + mesh = aa.mesh.DelaunayNN( + pixels=100, + zeroed_pixels=8, + areas_factor=0.4, + ) + + assert mesh.pixels == 108 + np.testing.assert_array_equal(mesh.zeroed_pixels, np.arange(100, 108)) + assert mesh.areas_factor == 0.4 + assert mesh.max_cavity_triangles == 32 + assert mesh.max_neighbors == 32 + assert mesh.query_chunk == 256 + assert mesh.interpolator_cls is InterpolatorDelaunayNN + assert aa.InterpolatorDelaunayNN is InterpolatorDelaunayNN + + +def test__interpolator__partition_of_unity_linear_precision_and_split_padding( + grid_2d_sub_1_7x7, +): + mesh_grid = aa.Grid2DIrregular( + values=[ + [-1.0, -1.0], + [-1.0, 1.0], + [1.0, -1.0], + [1.0, 1.0], + [-0.25, 0.1], + [0.35, -0.2], + ] + ) + interpolator = aa.InterpolatorDelaunayNN( + mesh=aa.mesh.DelaunayNN(pixels=6), + mesh_grid=mesh_grid, + data_grid=grid_2d_sub_1_7x7, + ) + + mappings, sizes, weights = interpolator._mappings_sizes_weights + np.testing.assert_allclose(weights.sum(axis=1), 1.0, atol=1.0e-12) + assert mappings.shape == (grid_2d_sub_1_7x7.over_sampled.shape[0], 32) + assert sizes.max() <= 6 + assert not interpolator.delaunay.overflow.any() + assert not interpolator.delaunay.degenerate.any() + + inside = interpolator.delaunay.cavity_sizes > 0 + reconstructed_coordinates = np.sum( + mesh_grid.array[mappings.clip(min=0)] * weights[..., None], axis=1 + ) + np.testing.assert_allclose( + reconstructed_coordinates[inside], + grid_2d_sub_1_7x7.over_sampled.array[inside], + atol=1.0e-11, + ) + + split_mappings, split_sizes, split_weights = ( + interpolator._mappings_sizes_weights_split + ) + assert split_mappings.shape == (4 * mesh_grid.shape[0], 33) + assert split_weights.shape == split_mappings.shape + np.testing.assert_array_equal(split_mappings[:, -1], -1) + np.testing.assert_array_equal(split_weights[:, -1], 0.0) + np.testing.assert_array_equal(split_sizes, interpolator.delaunay.splitted_sizes) + np.testing.assert_allclose(split_weights[:, :-1].sum(axis=1), 1.0, atol=1.0e-12) + assert not interpolator.delaunay.split_overflow.any() + assert not interpolator.delaunay.split_degenerate.any() + + +def test__smooth_source_mapping_is_numerically_close_to_delaunay(): + rng = np.random.default_rng(4) + axis = np.linspace(-1.0, 1.0, 12) + y, x = np.meshgrid(axis, axis, indexing="ij") + points = np.stack([y.ravel(), x.ravel()], axis=1) + points += rng.normal(0.0, 0.025, size=points.shape) + + query_axis = np.linspace(-0.9, 0.9, 45) + query_y, query_x = np.meshgrid(query_axis, query_axis, indexing="ij") + query = np.stack([query_y.ravel(), query_x.ravel()], axis=1) + + _, _, delaunay_mappings, _, _ = scipy_delaunay(points, query, areas_factor=0.5) + delaunay_weights = pixel_weights_delaunay_from( + data_grid=query, + mesh_grid=points, + pix_indexes_for_sub_slim_index=delaunay_mappings, + ) + delaunay_nn = scipy_delaunay_nn(points, query, areas_factor=0.5) + nn_mappings, nn_weights = delaunay_nn[2], delaunay_nn[4] + + source = np.exp(-((points[:, 0] - 0.12) ** 2 + (points[:, 1] + 0.08) ** 2) / 0.22) + delaunay_image = np.sum( + source[delaunay_mappings.clip(min=0)] * delaunay_weights, axis=1 + ) + nn_image = np.sum(source[nn_mappings.clip(min=0)] * nn_weights, axis=1) + + relative_l2 = np.linalg.norm(nn_image - delaunay_image) / np.linalg.norm( + delaunay_image + ) + correlation = np.corrcoef(nn_image, delaunay_image)[0, 1] + + assert relative_l2 < 0.02 + assert correlation > 0.999 + assert not delaunay_nn[10].any() + assert not delaunay_nn[11].any() + assert not delaunay_nn[13].any() + assert not delaunay_nn[14].any() diff --git a/test_autoarray/inversion/pixelization/interpolator/test_sibson.py b/test_autoarray/inversion/pixelization/interpolator/test_sibson.py new file mode 100644 index 000000000..418a41e2e --- /dev/null +++ b/test_autoarray/inversion/pixelization/interpolator/test_sibson.py @@ -0,0 +1,177 @@ +import numpy as np + +from autoarray.inversion.mesh.interpolator.delaunay import ( + pix_indexes_delaunay_walk_from, + scipy_delaunay_tri_only, +) +from autoarray.inversion.mesh.interpolator.sibson import ( + delaunay_circumcircles_from, + sibson_mappings_weights_from_tables, +) + + +def _sibson_numpy(points, query, max_cavity_triangles=32, max_neighbors=32): + 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 sibson_mappings_weights_from_tables( + query_points=query, + points=points, + simplices_padded=simplices, + simplex_neighbors=neighbors, + simplex_indexes=simplex_indexes, + outside_fallback_indexes=delaunay_mappings[:, 0], + max_cavity_triangles=max_cavity_triangles, + max_neighbors=max_neighbors, + xp=np, + ) + + +def test__circumcircles__known_right_triangle_and_padding(): + points = np.array([[0.0, 0.0], [2.0, 0.0], [0.0, 2.0]]) + simplices = np.array([[0, 1, 2], [-1, -1, -1]], dtype=np.int32) + + centres, radii_squared, valid = delaunay_circumcircles_from( + points, simplices, xp=np + ) + + np.testing.assert_allclose(centres[0], [1.0, 1.0]) + np.testing.assert_allclose(radii_squared[0], 2.0) + assert valid.tolist() == [True, False] + + +def test__single_triangle__matches_barycentric_and_outside_fallback(): + points = np.array([[0.0, 0.0], [2.0, 0.0], [0.0, 2.0]]) + query = np.array([[0.4, 0.6], [0.2, 0.2], [3.0, 3.0]]) + + mappings, sizes, weights, cavity_sizes, overflow, degenerate = _sibson_numpy( + points, query, max_cavity_triangles=8, max_neighbors=4 + ) + + dense_weights = np.zeros((query.shape[0], points.shape[0])) + for row in range(query.shape[0]): + valid = mappings[row] >= 0 + np.add.at(dense_weights[row], mappings[row, valid], weights[row, valid]) + + np.testing.assert_allclose(dense_weights[0], [0.5, 0.2, 0.3]) + np.testing.assert_allclose(dense_weights[1], [0.8, 0.1, 0.1]) + np.testing.assert_allclose(dense_weights[2], [0.0, 1.0, 0.0]) + assert sizes.tolist() == [3, 3, 1] + assert cavity_sizes.tolist() == [1, 1, 0] + assert not overflow.any() + assert not degenerate.any() + + +def test__watson_weights__match_historical_c_natural_neighbor_reference(): + points = np.array( + [ + [-1.0, -0.8], + [-0.2, -1.1], + [0.9, -0.7], + [1.2, 0.4], + [0.5, 1.1], + [-0.6, 0.9], + [-0.1, 0.0], + [0.55, 0.2], + ] + ) + query = np.array([[-0.3, -0.2], [0.2, 0.35], [0.75, -0.2], [-0.45, 0.55]]) + expected = np.array( + [ + [ + 0.1874290255921386, + 0.0833454011106457, + 0.0002366866860013, + 0.0, + 0.0, + 0.0464320470840236, + 0.6825568395271908, + 0.0, + ], + [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.1673987581070807, + 0.0991207867537038, + 0.3502171651198253, + 0.3832632900193901, + ], + [ + 0.0, + 0.0009216096884568, + 0.4570373759794244, + 0.1185942343731840, + 0.0, + 0.0, + 0.0559355794908082, + 0.3675112004681266, + ], + [ + 0.0568307353360843, + 0.0, + 0.0, + 0.0, + 0.0229991789644525, + 0.6323181584385680, + 0.2824561831946050, + 0.0053957440662901, + ], + ] + ) + + mappings, _, weights, _, overflow, degenerate = _sibson_numpy(points, query) + dense_weights = np.zeros_like(expected) + for row in range(query.shape[0]): + valid = mappings[row] >= 0 + np.add.at(dense_weights[row], mappings[row, valid], weights[row, valid]) + + np.testing.assert_allclose(dense_weights, expected, atol=1.0e-13) + assert not overflow.any() + assert not degenerate.any() + + +def test__random_mesh__partition_of_unity_and_linear_precision(): + rng = np.random.default_rng(10) + points = rng.uniform(-1.0, 1.0, size=(100, 2)) + query = rng.uniform(-0.8, 0.8, size=(500, 2)) + + mappings, sizes, weights, cavity_sizes, overflow, degenerate = _sibson_numpy( + points, query + ) + + reconstructed = np.zeros_like(query) + for row in range(query.shape[0]): + valid = mappings[row] >= 0 + reconstructed[row] = np.sum( + points[mappings[row, valid]] * weights[row, valid, None], axis=0 + ) + + np.testing.assert_allclose(weights.sum(axis=1), 1.0, atol=1.0e-12) + inside = cavity_sizes > 0 + np.testing.assert_allclose(reconstructed[inside], query[inside], atol=1.0e-11) + assert sizes.max() <= 12 + assert cavity_sizes.max() <= 12 + assert not overflow.any() + assert not degenerate.any() + + +def test__cavity_cap__reports_overflow_instead_of_silent_approximation(): + points = np.array([[0.0, 0.0], [2.0, 0.0], [0.1, 1.8], [1.8, 2.1], [1.0, 0.7]]) + query = np.array([[0.9, 0.9]]) + + _, _, weights, cavity_sizes, overflow, _ = _sibson_numpy( + points, query, max_cavity_triangles=1, max_neighbors=5 + ) + + assert cavity_sizes[0] == 1 + assert overflow[0] + assert np.isnan(weights[0]).all() diff --git a/test_autoarray/inversion/pixelization/test_split_regularization_support.py b/test_autoarray/inversion/pixelization/test_split_regularization_support.py index 23fc724f4..28192f1d2 100644 --- a/test_autoarray/inversion/pixelization/test_split_regularization_support.py +++ b/test_autoarray/inversion/pixelization/test_split_regularization_support.py @@ -36,6 +36,7 @@ ADAPTIVE_MESHES = [ aa.mesh.Delaunay, + aa.mesh.DelaunayNN, aa.mesh.KNNBarycentric, ] @@ -99,6 +100,7 @@ def test__capability_flags(): assert aa.mesh.RectangularAdaptDensity(shape=(15, 15)).supports_split_regularization is False assert aa.mesh.RectangularAdaptImage(shape=(15, 15)).supports_split_regularization is False assert aa.mesh.Delaunay(pixels=100).supports_split_regularization is True + assert aa.mesh.DelaunayNN(pixels=100).supports_split_regularization is True assert aa.mesh.KNNBarycentric(pixels=100).supports_split_regularization is True assert aa.reg.ConstantSplit().is_split_regularization is True @@ -119,10 +121,12 @@ def test__interpolator_rectangular_has_no_split_mappings(): InterpolatorRectangularUniform, ) from autoarray.inversion.mesh.interpolator.delaunay import InterpolatorDelaunay + from autoarray.inversion.mesh.interpolator.sibson import InterpolatorDelaunayNN assert not hasattr(InterpolatorRectangular, "_mappings_sizes_weights_split") assert not hasattr(InterpolatorRectangularUniform, "_mappings_sizes_weights_split") assert hasattr(InterpolatorDelaunay, "_mappings_sizes_weights_split") + assert hasattr(InterpolatorDelaunayNN, "_mappings_sizes_weights_split") def test__pixelization_exception_is_not_a_fit_exception():