fix: exclude zero-weight Delaunay mappings from adaptive pixel signals (NumPy/JAX log-det parity)

autoarray/inversion/mappers/abstract.py

@@ -512,6 +512,7 @@ def pixel_signals_from(self, signal_scale: float, xp=np) -> np.ndarray:
             pix_size_for_sub_slim_index=self.pix_sizes_for_sub_slim_index,
             slim_index_for_sub_slim_index=self.over_sampler.slim_for_sub_slim,
             adapt_data=adapt_data,
+            exclude_zero_weight_mappings=self.interpolator.exclude_zero_weight_mappings,
             xp=xp,
         )
 

autoarray/inversion/mappers/mapper_util.py

@@ -11,6 +11,7 @@ def adaptive_pixel_signals_from(
     pix_size_for_sub_slim_index: np.ndarray,
     slim_index_for_sub_slim_index: np.ndarray,
     adapt_data: np.ndarray,
+    exclude_zero_weight_mappings: bool = False,
     xp=np,
 ) -> np.ndarray:
     """
@@ -19,8 +20,11 @@ def adaptive_pixel_signals_from(
 
     The pixel signals are computed as follows:
 
-    1) Divide by the number of mappe data points in the pixel, to ensure all pixels have the same
+    1) Divide by the number of mapped data points in the pixel, to ensure all pixels have the same
        'relative' signal (i.e. a pixel with 10 pixels doesn't have x2 the signal of one with 5).
+       By default every mapping entry within a data point's size counts, including entries whose
+       weight is zero. With `exclude_zero_weight_mappings=True` (set by Delaunay interpolators),
+       zero-weight entries are not counted -- see that parameter.
 
     2) Divided by the maximum pixel-signal, so that all signals vary between 0 and 1. This ensures that the
        regularization weight_list are defined identically for any data quantity or signal-to-noise_map ratio.
@@ -39,6 +43,20 @@ def adaptive_pixel_signals_from(
         A 1D array util every pixel on the grid to a pixel on the pixelization.
     adapt_data
         The image of the galaxy which is used to compute the weigghted pixel signals.
+    exclude_zero_weight_mappings
+        If `True`, mapping entries whose weight is exactly zero are not counted towards a pixel's number of
+        mapped data points. Delaunay interpolators set this (via their `exclude_zero_weight_mappings` class
+        attribute): a data point lying exactly on a Delaunay vertex has barycentric weights (0, 0, 1) in *any*
+        triangle around that vertex (a point on an edge has a zero weight in either neighbouring triangle), and
+        scipy's `find_simplex` and the JAX Delaunay walk legitimately choose different triangles. Counting those
+        zero-weight duplicates made the pixel signals, and through them the adaptive regularization matrix,
+        differ between the NumPy and JAX backends from identical inputs (PyAutoArray#627). The set of non-zero
+        (index, weight) pairs is unique, so excluding zero weights makes the result independent of that choice.
+
+        It is `False` for every other interpolator. In particular `KNNBarycentric` clips negative barycentric
+        weights to exactly zero for every data point outside the triangle formed by its 3 nearest vertices --
+        whole regions, not rare coincidences -- so excluding them there would change pixel signals broadly
+        and make the count (and hence the likelihood) jump wherever a clipped weight reaches zero.
     """
 
     M_sub, B = pix_indexes_for_sub_slim_index.shape
@@ -50,9 +68,15 @@ def adaptive_pixel_signals_from(
     # 2) Build a matching “parent‐slim” index for each flattened entry:
     I_sub = xp.repeat(xp.arange(M_sub), B)  # (M_sub*B,)
 
-    # 3) Mask out any k >= pix_size_for_sub_slim_index[i]
+    # 3) Mask out any k >= pix_size_for_sub_slim_index[i] and, for Delaunay tables, any entry whose
+    #    weight is zero (see `exclude_zero_weight_mappings` above, PyAutoArray#627). The flag is a
+    #    static Python bool, so this is jit-safe; the mask is not differentiated, so the function
+    #    stays differentiable under JAX.
     valid = (xp.arange(B)[None, :] < pix_size_for_sub_slim_index[:, None]).reshape(-1)
 
+    if exclude_zero_weight_mappings:
+        valid = valid & (flat_weights != 0)
+
     flat_weights = xp.where(valid, flat_weights, 0.0)
     flat_pixidx = xp.where(
         valid, flat_pixidx, pixels

autoarray/inversion/mesh/interpolator/abstract.py

@@ -3,6 +3,11 @@
 
 class AbstractInterpolator:
 
+    # Whether zero-weight mapping entries are excluded when counting the data points mapped to each pixel for
+    # adaptive pixel signals (`mapper_util.adaptive_pixel_signals_from`). Only Delaunay tables set this: their
+    # zero-weight entries are arbitrary duplicates from the triangle choice (PyAutoArray#627).
+    exclude_zero_weight_mappings = False
+
     def __init__(
         self,
         mesh,

autoarray/inversion/mesh/interpolator/delaunay.py

@@ -737,6 +737,12 @@ def splitted_sizes(self):
 
 
 class InterpolatorDelaunay(AbstractInterpolator):
+
+    # A data point on a Delaunay vertex (or edge) has zero barycentric weights whose vertices depend on which
+    # surrounding triangle is chosen, and scipy and the JAX walk choose differently, so zero-weight entries must
+    # not count towards adaptive pixel signals (PyAutoArray#627).
+    exclude_zero_weight_mappings = True
+
     def __init__(
         self,
         mesh,

autoarray/inversion/mesh/interpolator/knn.py

@@ -145,6 +145,11 @@ def kernel_interpolate_points(points, query_chunk, values, k, radius_scale):
 
 class InterpolatorKNearestNeighbor(InterpolatorDelaunay):
 
+    # Not a Delaunay table: `KNNBarycentric` clips negative barycentric weights to exactly zero for every point
+    # outside its 3-nearest triangle, so its zero weights are region-wide, not choice-dependent duplicates.
+    # Excluding them would change pixel signals broadly and add a jump in the likelihood (PyAutoArray#627).
+    exclude_zero_weight_mappings = False
+
     @property
     def dual_areas(self):
         """

test_autoarray/inversion/inversion/test_log_det_sparse.py

@@ -409,3 +409,120 @@ def raise_splu(*args, **kwargs):
     assert inversion.log_det_curvature_reg_matrix_term == pytest.approx(
         _dense_log_det(matrix), rel=1.0e-12
     )
+
+
+def _vertex_coincident_adapt_split_mapper(regularization, xp):
+    """
+    A Delaunay ``AdaptSplit`` mapper whose mesh is the stride-2 subsample of a 16x16 data grid
+    with a smooth ``0.01 * sin(7 * grid)`` distortion, so 64 data points sit exactly on mesh
+    vertices. scipy's ``find_simplex`` and the JAX Delaunay walk pick different (equally valid)
+    triangles for some of those points (PyAutoArray#627).
+    """
+    n = 16
+
+    yy, xx = np.meshgrid(np.linspace(-1, 1, n), np.linspace(-1, 1, n), indexing="ij")
+
+    data_grid = np.stack([yy.ravel(), xx.ravel()], axis=1)
+    data_grid = data_grid + 0.01 * np.sin(7 * data_grid[:, ::-1])
+
+    mesh_grid = data_grid.reshape(n, n, 2)[::2, ::2].reshape(-1, 2)
+
+    adapt_data = np.exp(-np.sum(data_grid**2, axis=1) / 0.3)
+
+    interpolator = aa.mesh.Delaunay(pixels=mesh_grid.shape[0]).interpolator_from(
+        source_plane_data_grid=aa.Grid2D.no_mask(
+            values=data_grid, shape_native=(data_grid.shape[0], 1), pixel_scales=1.0
+        ),
+        source_plane_mesh_grid=aa.Grid2DIrregular(mesh_grid),
+        adapt_data=aa.Array2D.no_mask(
+            values=adapt_data, shape_native=(n, n), pixel_scales=0.1
+        ),
+        xp=xp,
+    )
+
+    return aa.Mapper(interpolator=interpolator, regularization=regularization, xp=xp)
+
+
+def _high_precision_log_det(matrix):
+    """
+    A reference log determinant: a 50-digit mpmath determinant, or, if mpmath is unavailable, a
+    Cholesky carried out in ``np.longdouble`` (NumPy and SciPy linalg cast that back to fp64,
+    so the factorization is written out here).
+    """
+    try:
+        import mpmath
+    except ImportError:
+        mpmath = None
+
+    if mpmath is not None:
+        with mpmath.workdps(50):
+            return float(mpmath.log(mpmath.det(mpmath.matrix(matrix.tolist()))))
+
+    lower = np.array(matrix, dtype=np.longdouble)
+    pixels = lower.shape[0]
+
+    for j in range(pixels):
+        lower[j, j] = np.sqrt(lower[j, j] - np.sum(lower[j, :j] ** 2))
+        lower[j + 1 :, j] = (
+            lower[j + 1 :, j] - lower[j + 1 :, :j] @ lower[j, :j]
+        ) / lower[j, j]
+
+    return float(2 * np.sum(np.log(np.diag(lower))))
+
+
+def test__log_det_regularization_matrix_term__vertex_coincident_adapt_split__numpy_matches_jax():
+    """
+    PyAutoArray#627: on a mesh whose vertices coincide with data points, the NumPy and JAX
+    backends used to build *different* ``AdaptSplit`` regularization matrices (max difference
+    0.74 on max|H| 14.6 here, a log-determinant gap of 0.097), because the zero-weight mapping
+    entries of vertex-coincident points depend on which triangle each backend's point location
+    picks, and those entries were counted in the adaptive pixel signals.
+
+    The log-determinant routes themselves were never the cause: given the same ``H`` both agree
+    with a 50-digit reference. That is pinned here only at moderate coefficients, where ``H`` is
+    well-posed (``cond ~ 2e9``). At high coefficients (``max(coefficient)^4 >~ 1e6``) the absolute
+    1e-8 diagonal jitter puts ``H``'s constant-mode eigenvalue at round-off level, so the log
+    determinant is round-off-determined on either backend and no parity is asserted.
+    """
+    jax = pytest.importorskip("jax")
+    jax.config.update("jax_enable_x64", True)
+    import jax.numpy as jnp
+
+    regularization = aa.reg.AdaptSplit(
+        inner_coefficient=0.4, outer_coefficient=1.7, signal_scale=0.3
+    )
+
+    mapper_np = _vertex_coincident_adapt_split_mapper(regularization, xp=np)
+    mapper_jax = _vertex_coincident_adapt_split_mapper(regularization, xp=jnp)
+
+    matrix_np = np.asarray(
+        regularization.regularization_matrix_from(linear_obj=mapper_np)
+    )
+    matrix_jax = regularization.regularization_matrix_from(
+        linear_obj=mapper_jax, xp=jnp
+    )
+
+    np.testing.assert_allclose(
+        np.asarray(matrix_jax),
+        matrix_np,
+        rtol=1.0e-12,
+        atol=1.0e-12 * np.abs(matrix_np).max(),
+    )
+
+    inversion_np = aa.m.MockInversion(
+        linear_obj_list=[mapper_np], regularization_matrix=matrix_np
+    )
+    inversion_jax = aa.m.MockInversion(
+        linear_obj_list=[mapper_jax], regularization_matrix=matrix_jax
+    )
+    inversion_jax.use_jax = True
+
+    log_det_np = float(inversion_np.log_det_regularization_matrix_term)
+    log_det_jax = float(inversion_jax.log_det_regularization_matrix_term)
+
+    assert abs(log_det_np - log_det_jax) < 1.0e-6
+
+    reference = _high_precision_log_det(matrix_np)
+
+    assert abs(log_det_np - reference) < 1.0e-6
+    assert abs(log_det_jax - reference) < 1.0e-6