perf: skip Sibson row padding in adaptive signals and unique mappings (DelaunayNN imaging)

autoarray/inversion/mappers/abstract.py

@@ -249,11 +249,13 @@ def unique_mappings(self) -> UniqueMappings:
             pix_lengths,
         ) = mapper_numba_util.data_slim_to_pixelization_unique_from(
             data_pixels=self.over_sampler.mask.pixels_in_mask,
-            pix_indexes_for_sub_slim_index=np.array(
+            # `np.asarray`, not `np.array`: the kernel only reads these tables, so a NumPy table is passed
+            # without copying it (DelaunayNN tables are 32 columns wide, PyAutoArray#629).
+            pix_indexes_for_sub_slim_index=np.asarray(
                 self.pix_indexes_for_sub_slim_index
             ),
-            pix_sizes_for_sub_slim_index=np.array(self.pix_sizes_for_sub_slim_index),
-            pix_weights_for_sub_slim_index=np.array(
+            pix_sizes_for_sub_slim_index=np.asarray(self.pix_sizes_for_sub_slim_index),
+            pix_weights_for_sub_slim_index=np.asarray(
                 self.pix_weights_for_sub_slim_index
             ),
             pix_pixels=self.params,

autoarray/inversion/mappers/mapper_numba_util.py

@@ -61,17 +61,21 @@ def data_slim_to_pixelization_unique_from(
 
     # TODO : Work out if we can reduce size from np.max(sub_size) using sub_size of max_pix_mappings.
 
-    data_to_pix_unique = -1 * np.ones(
-        (data_pixels, max_pix_mappings * np.max(sub_size) ** 2)
+    data_to_pix_unique = np.full(
+        (data_pixels, max_pix_mappings * np.max(sub_size) ** 2), -1.0
     )
     data_weights = np.zeros((data_pixels, max_pix_mappings * np.max(sub_size) ** 2))
     pix_lengths = np.zeros(data_pixels)
-    pix_check = -1 * np.ones(shape=pix_pixels)
+    pix_check = np.full(pix_pixels, -1.0)
 
     ip_sub_start = 0
 
     for ip in range(data_pixels):
-        pix_check[:] = -1
+        # Reset only the pixelization pixels the previous data pixel touched (its unique list), rather than
+        # all `pix_pixels` entries: the full reset cost O(data_pixels * pix_pixels) per call (PyAutoArray#629).
+        if ip > 0:
+            for unique_index in range(int(pix_lengths[ip - 1])):
+                pix_check[int(data_to_pix_unique[ip - 1, unique_index])] = -1.0
 
         pix_size = 0
 
@@ -98,3 +102,70 @@ def data_slim_to_pixelization_unique_from(
         pix_lengths[ip] = pix_size
 
     return data_to_pix_unique, data_weights, pix_lengths
+
+
+@numba_util.jit()
+def adaptive_pixel_signal_sums_from(
+    pixels: int,
+    pixel_weights: np.ndarray,
+    pix_indexes_for_sub_slim_index: np.ndarray,
+    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,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """
+    The summed weighted adapt-data signal and the number of counted mapping entries in every pixelization pixel,
+    visiting only the first `pix_size_for_sub_slim_index[i]` (live) entries of every row.
+
+    This is the numba form of steps 1-6 of `mapper_util.adaptive_pixel_signals_from` for the NumPy path. Natural
+    neighbour (`DelaunayNN`) mapping tables are `SIBSON_MAX_NEIGHBORS` (32) columns wide but hold at most ~15 live
+    entries per row, so the padded scatter spent most of its time on padding (PyAutoArray#629).
+
+    The result is bit-identical to the padded `np.add.at` scatter: the live entries are added into each pixel in
+    the same row-major order, starting from the same `0.0`, as the same `adapt_data * weight` products, and every
+    entry skipped here (padding, a `-1` index, and a zero weight when `exclude_zero_weight_mappings`) only ever
+    reached the scatter's discarded out-of-bounds slot.
+
+    Parameters
+    ----------
+    pixels
+        The number of pixelization pixels the signals are summed into.
+    pixel_weights
+        The interpolation weight of every (sub-pixel, mapping) entry, shape [total_sub_pixels, max_mappings].
+    pix_indexes_for_sub_slim_index
+        The pixelization pixel of every (sub-pixel, mapping) entry, padded with -1.
+    pix_size_for_sub_slim_index
+        The number of live mapping entries of every sub-pixel.
+    slim_index_for_sub_slim_index
+        The slim data pixel of every sub-pixel.
+    adapt_data
+        The adapt image whose values are summed.
+    exclude_zero_weight_mappings
+        Whether entries with a zero weight are skipped (Delaunay interpolators, PyAutoArray#627).
+
+    Returns
+    -------
+    The summed signal and the count of every pixel, each of shape [pixels].
+    """
+    pixel_signals = np.zeros(pixels)
+    pixel_counts = np.zeros(pixels)
+
+    for sub_index in range(pix_indexes_for_sub_slim_index.shape[0]):
+        data_value = adapt_data[slim_index_for_sub_slim_index[sub_index]]
+
+        for mapping_index in range(pix_size_for_sub_slim_index[sub_index]):
+            pix = pix_indexes_for_sub_slim_index[sub_index, mapping_index]
+
+            if pix < 0 or pix >= pixels:
+                continue
+
+            weight = pixel_weights[sub_index, mapping_index]
+
+            if exclude_zero_weight_mappings and weight == 0.0:
+                continue
+
+            pixel_signals[pix] += data_value * weight
+            pixel_counts[pix] += 1.0
+
+    return pixel_signals, pixel_counts

autoarray/inversion/mappers/mapper_util.py

@@ -59,47 +59,41 @@ def adaptive_pixel_signals_from(
         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
-
-    # 1) Flatten the per‐mapping tables:
-    flat_pixidx = pix_indexes_for_sub_slim_index.reshape(-1)  # (M_sub*B,)
-    flat_weights = pixel_weights.reshape(-1)  # (M_sub*B,)
-
-    # 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] 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
-    )  # send invalid indices to an out-of-bounds slot
-
-    # 4) Look up data & multiply by mapping weights:
-    flat_data_vals = xp.take(adapt_data[slim_index_for_sub_slim_index], I_sub, axis=0)
-    flat_contrib = flat_data_vals * flat_weights  # (M_sub*B,)
-
-    pixel_signals = xp.zeros((pixels + 1,))
-    pixel_counts = xp.zeros((pixels + 1,))
-
-    # 5) Scatter‐add into signal sums and counts:
     if xp.__name__.startswith("jax"):
-        pixel_signals = pixel_signals.at[flat_pixidx].add(flat_contrib)
-        pixel_counts = pixel_counts.at[flat_pixidx].add(valid.astype(float))
+        pixel_signals, pixel_counts = _pixel_signal_sums_jax(
+            pixels=pixels,
+            pixel_weights=pixel_weights,
+            pix_indexes_for_sub_slim_index=pix_indexes_for_sub_slim_index,
+            pix_size_for_sub_slim_index=pix_size_for_sub_slim_index,
+            slim_index_for_sub_slim_index=slim_index_for_sub_slim_index,
+            adapt_data=adapt_data,
+            exclude_zero_weight_mappings=exclude_zero_weight_mappings,
+            xp=xp,
+        )
+    elif _use_numba():
+        from autoarray.inversion.mappers import mapper_numba_util
+
+        pix_indexes = np.asarray(pix_indexes_for_sub_slim_index)
+
+        pixel_signals, pixel_counts = mapper_numba_util.adaptive_pixel_signal_sums_from(
+            pixels=pixels,
+            pixel_weights=np.asarray(pixel_weights).reshape(pix_indexes.shape),
+            pix_indexes_for_sub_slim_index=pix_indexes,
+            pix_size_for_sub_slim_index=np.asarray(pix_size_for_sub_slim_index),
+            slim_index_for_sub_slim_index=np.asarray(slim_index_for_sub_slim_index),
+            adapt_data=np.asarray(adapt_data),
+            exclude_zero_weight_mappings=bool(exclude_zero_weight_mappings),
+        )
     else:
-        xp.add.at(pixel_signals, flat_pixidx, flat_contrib)
-        xp.add.at(pixel_counts, flat_pixidx, valid.astype(float))
-
-    # 6) Drop the extra “out-of-bounds” slot:
-    pixel_signals = pixel_signals[:pixels]
-    pixel_counts = pixel_counts[:pixels]
+        pixel_signals, pixel_counts = _pixel_signal_sums_numpy(
+            pixels=pixels,
+            pixel_weights=pixel_weights,
+            pix_indexes_for_sub_slim_index=pix_indexes_for_sub_slim_index,
+            pix_size_for_sub_slim_index=pix_size_for_sub_slim_index,
+            slim_index_for_sub_slim_index=slim_index_for_sub_slim_index,
+            adapt_data=adapt_data,
+            exclude_zero_weight_mappings=exclude_zero_weight_mappings,
+        )
 
     # 7) Normalize
     #
@@ -137,6 +131,142 @@ def adaptive_pixel_signals_from(
     return xp.where(zero_signal, 0.0, safe_signals**signal_scale)
 
 
+def _use_numba() -> bool:
+    """
+    Whether the NumPy branch of `adaptive_pixel_signals_from` uses the numba kernel
+    `mapper_numba_util.adaptive_pixel_signal_sums_from`.
+
+    True when numba is importable and `general.yaml` `numba: use_numba` is true (the default when the config
+    cannot be read); otherwise the vectorised `np.bincount` form `_pixel_signal_sums_numpy` is used. Both are
+    bit-identical to the padded scatter.
+    """
+    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:
+        return False
+
+    return True
+
+
+def _pixel_signal_sums_jax(
+    pixels,
+    pixel_weights,
+    pix_indexes_for_sub_slim_index,
+    pix_size_for_sub_slim_index,
+    slim_index_for_sub_slim_index,
+    adapt_data,
+    exclude_zero_weight_mappings,
+    xp,
+):
+    """
+    The JAX branch of `adaptive_pixel_signals_from` steps 1-6: the summed weighted adapt-data signal and the
+    number of counted mapping entries in every pixel, built with fixed shapes over the full padded tables so it
+    traces under `jit` and stays differentiable.
+    """
+    M_sub, B = pix_indexes_for_sub_slim_index.shape
+
+    # 1) Flatten the per-mapping tables:
+    flat_pixidx = pix_indexes_for_sub_slim_index.reshape(-1)  # (M_sub*B,)
+    flat_weights = pixel_weights.reshape(-1)  # (M_sub*B,)
+
+    # 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] 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
+    )  # send invalid indices to an out-of-bounds slot
+
+    # 4) Look up data & multiply by mapping weights:
+    flat_data_vals = xp.take(adapt_data[slim_index_for_sub_slim_index], I_sub, axis=0)
+    flat_contrib = flat_data_vals * flat_weights  # (M_sub*B,)
+
+    pixel_signals = xp.zeros((pixels + 1,))
+    pixel_counts = xp.zeros((pixels + 1,))
+
+    # 5) Scatter-add into signal sums and counts:
+    pixel_signals = pixel_signals.at[flat_pixidx].add(flat_contrib)
+    pixel_counts = pixel_counts.at[flat_pixidx].add(valid.astype(float))
+
+    # 6) Drop the extra "out-of-bounds" slot:
+    return pixel_signals[:pixels], pixel_counts[:pixels]
+
+
+def _pixel_signal_sums_numpy(
+    pixels,
+    pixel_weights,
+    pix_indexes_for_sub_slim_index,
+    pix_size_for_sub_slim_index,
+    slim_index_for_sub_slim_index,
+    adapt_data,
+    exclude_zero_weight_mappings,
+):
+    """
+    The NumPy branch of `adaptive_pixel_signals_from` steps 1-6, computed over the live mapping entries only.
+
+    Natural-neighbour (`DelaunayNN`) tables are `SIBSON_MAX_NEIGHBORS` (32) columns wide but hold at most ~15 live
+    entries per row, so scattering over every padded slot (the JAX branch's fixed-shape form) spent most of its
+    time on padding (PyAutoArray#629). Here the tables are first trimmed to their widest live row, the live
+    entries compacted in row-major order and summed with `np.bincount`.
+
+    The result is bit-identical to the padded scatter: `np.bincount` adds the compacted entries into each bin
+    in the same row-major order that `np.add.at` visited them, starting from the same `0.0`, and every entry it
+    drops (padding, `-1` indexes, and zero weights when `exclude_zero_weight_mappings`) only ever reached the
+    padded scatter's discarded out-of-bounds slot. Counts are integers, so summing them as integers is exact.
+    """
+    pix_indexes = np.asarray(pix_indexes_for_sub_slim_index)
+    weights = np.asarray(pixel_weights).reshape(pix_indexes.shape)
+    sizes = np.asarray(pix_size_for_sub_slim_index)
+
+    # Trim every table to its widest live row (a view; padding beyond it is never read).
+    width = int(np.max(sizes)) if sizes.size else 0
+    width = min(max(width, 0), pix_indexes.shape[1])
+    pix_indexes = pix_indexes[:, :width]
+    weights = weights[:, :width]
+
+    valid = np.arange(width)[None, :] < sizes[:, None]
+
+    if exclude_zero_weight_mappings:
+        valid &= weights != 0
+
+    # A `-1` index inside a row's size only ever landed in the padded scatter's out-of-bounds slot.
+    valid &= pix_indexes >= 0
+
+    rows = np.nonzero(valid)[0]
+    live_indexes = pix_indexes[valid]
+    live_weights = weights[valid]
+
+    data_values = np.asarray(adapt_data)[
+        np.asarray(slim_index_for_sub_slim_index)[rows]
+    ]
+
+    pixel_signals = np.bincount(
+        live_indexes, weights=data_values * live_weights, minlength=pixels
+    )[:pixels]
+    pixel_counts = np.bincount(live_indexes, minlength=pixels)[:pixels].astype(float)
+
+    return pixel_signals, pixel_counts
+
+
 def sparse_triplets_from(
     pix_indexes_for_sub,  # (M_sub, P)
     pix_weights_for_sub,  # (M_sub, P)

autoarray/inversion/mesh/interpolator/_sibson_numba.py

@@ -353,6 +353,11 @@ def _sibson_from_tables_parallel(
         if failed:
             for k in range(max_neighbors):
                 weights[q, k] = np.nan
+            # Keep a failed row at size >= 1 so its NaN weight propagates (see the reference
+            # ``_sibson_single_from_tables``, PyAutoArray#629).
+            if size == 0 and max_neighbors > 0:
+                mappings[q, 0] = local_vertex
+                size = 1
 
         sizes[q] = size
         cavity_sizes[q] = count

autoarray/inversion/mesh/interpolator/sibson.py

@@ -552,6 +552,18 @@ def _sibson_single_from_tables(
     )
     overflow = overflow | neighbor_overflow
     failed = overflow | degenerate
+
+    # A failed row whose Sibson weights were all non-positive or non-finite has no selected entry, so its
+    # size would be 0 and every consumer (which reads only the first `size` entries) would skip the NaN
+    # weights below, silently returning a finite, wrong evidence. Keep such a row at size 1, mapped to the
+    # nearest seed vertex, so its NaN weight reaches the curvature matrix and the evidence is NaN
+    # (PyAutoArray#629).
+    empty_failed = failed & (size == 0)
+    mappings = xp.where(
+        empty_failed & (xp.arange(max_neighbors) == 0), local_vertex, mappings
+    ).astype(xp.int32)
+    size = xp.where(empty_failed, 1, size).astype(xp.int32)
+
     weights = xp.where(failed, xp.nan, weights)
     return mappings, size, weights, cavity_size, overflow, degenerate