From 58a044eb3b828f428fb5e6c5fb1bddd31f09fc92 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Sat, 17 Oct 2026 14:23:58 +0100 Subject: [PATCH] perf: skip Sibson row padding in adaptive signals and unique mappings (DelaunayNN imaging) DelaunayNN (Sibson) mapping tables are SIBSON_MAX_NEIGHBORS (32) columns wide but hold at most ~15 live entries per row. On the NumPy path: - adaptive_pixel_signals_from: the np.add.at scatter over every padded slot is replaced by a numba kernel (mapper_numba_util.adaptive_pixel_signal_sums_from) visiting only the first pix_sizes[i] entries of every row, with a vectorised np.bincount fallback over the compacted live entries when numba is absent or disabled. Both are bit-identical to the padded scatter (same row-major summation order, same products, skipped entries only ever reached the discarded out-of-bounds slot). The JAX branch is unchanged. - data_slim_to_pixelization_unique_from: reset only the pix_check entries the previous data pixel touched instead of all pix_pixels per data pixel; the mapper passes its tables with np.asarray (no 32-wide copies). Outputs and shapes unchanged, bit-identical. - Sibson: a failed row (overflow / degenerate) with no selected entry used to come out at size 0, so every consumer skipped its NaN weights and the evidence was finite and wrong. Reference, numba kernel and JAX path now keep such a row at size 1 mapped to the nearest seed vertex, so NaN propagates. Tests: bit-parity of both functions against the padded / full-reset algorithms on 32-wide DelaunayNN-shaped, 3-wide Delaunay and variable-width tables and on real Delaunay and DelaunayNN mappers (both NumPy routes), JAX vs NumPy on wide rows, NaN-in-padding ignored, and the size-0 failed-row case (constructed by overflowing the Watson arithmetic) for the reference, numba and JAX Sibson paths plus its NaN reaching signals and unique weights. Closes #629 Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01QBeyBLN7ntDChx9v1fmt6M --- autoarray/inversion/mappers/abstract.py | 8 +- .../inversion/mappers/mapper_numba_util.py | 79 ++++- autoarray/inversion/mappers/mapper_util.py | 208 ++++++++--- .../mesh/interpolator/_sibson_numba.py | 5 + .../inversion/mesh/interpolator/sibson.py | 12 + .../interpolator/test_sibson_failed_rows.py | 110 ++++++ .../mappers/test_mapper_util_wide_rows.py | 328 ++++++++++++++++++ 7 files changed, 704 insertions(+), 46 deletions(-) create mode 100644 test_autoarray/inversion/pixelization/interpolator/test_sibson_failed_rows.py create mode 100644 test_autoarray/inversion/pixelization/mappers/test_mapper_util_wide_rows.py diff --git a/autoarray/inversion/mappers/abstract.py b/autoarray/inversion/mappers/abstract.py index f224a9bcf..e2cee3e57 100644 --- a/autoarray/inversion/mappers/abstract.py +++ b/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, diff --git a/autoarray/inversion/mappers/mapper_numba_util.py b/autoarray/inversion/mappers/mapper_numba_util.py index 79514c2db..2fd20570d 100644 --- a/autoarray/inversion/mappers/mapper_numba_util.py +++ b/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 diff --git a/autoarray/inversion/mappers/mapper_util.py b/autoarray/inversion/mappers/mapper_util.py index 04cc10e2d..149ccf864 100644 --- a/autoarray/inversion/mappers/mapper_util.py +++ b/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) diff --git a/autoarray/inversion/mesh/interpolator/_sibson_numba.py b/autoarray/inversion/mesh/interpolator/_sibson_numba.py index c1f962b0d..2b0a00fd1 100644 --- a/autoarray/inversion/mesh/interpolator/_sibson_numba.py +++ b/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 diff --git a/autoarray/inversion/mesh/interpolator/sibson.py b/autoarray/inversion/mesh/interpolator/sibson.py index e30e4d9fc..11f54fa27 100644 --- a/autoarray/inversion/mesh/interpolator/sibson.py +++ b/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 diff --git a/test_autoarray/inversion/pixelization/interpolator/test_sibson_failed_rows.py b/test_autoarray/inversion/pixelization/interpolator/test_sibson_failed_rows.py new file mode 100644 index 000000000..0fc8e3bc8 --- /dev/null +++ b/test_autoarray/inversion/pixelization/interpolator/test_sibson_failed_rows.py @@ -0,0 +1,110 @@ +""" +Failed Sibson rows keep `size >= 1` so their NaN weights propagate (PyAutoArray#629). + +A degenerate query whose Watson contributions are all non-finite selects no entry. Before #629 such a row came +out at `size == 0`, so every consumer (which reads only the first `size` entries) skipped its NaN weights and +the evidence was finite and wrong. The reference, the numba kernel and the JAX path now all keep the row at +size 1, mapped to the nearest seed vertex, with a NaN weight. + +The degenerate case is built by triangulating a well-scaled point set and then evaluating the Sibson tables +at coordinates scaled by 1e150, where the Watson circle arithmetic overflows. +""" + +import importlib.util + +import numpy as np +import pytest + +from autoarray.inversion.mesh.interpolator import sibson +from autoarray.inversion.mesh.interpolator.delaunay import ( + pix_indexes_delaunay_walk_from, + scipy_delaunay_tri_only, +) + +SCALE = 1.0e150 + + +def _overflowing_tables(): + rng = np.random.default_rng(3) + points = rng.uniform(-1.0, 1.0, size=(30, 2)) + query = rng.uniform(-0.5, 0.5, size=(20, 2)) + 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, + ) + scaled_points = points * SCALE + with np.errstate(all="ignore"): + circumcircles = sibson.delaunay_circumcircles_from( + scaled_points, simplices, xp=np + ) + return dict( + query_points=query * SCALE, + points=scaled_points, + simplices_padded=simplices, + simplex_neighbors=neighbors, + simplex_indexes=simplex_indexes, + outside_fallback_indexes=delaunay_mappings[:, 0], + circumcircles=circumcircles, + max_cavity_triangles=32, + max_neighbors=32, + ) + + +def _assert_failed_rows_kept_at_size_one(outputs, points_count): + mappings, sizes, weights, _, overflow, degenerate = (np.asarray(o) for o in outputs) + failed = overflow | degenerate + + assert failed.all() + assert (sizes == 1).all() + assert ((mappings[:, 0] >= 0) & (mappings[:, 0] < points_count)).all() + assert (mappings[:, 1:] == -1).all() + assert np.isnan(weights[:, 0]).all() + + +def test__reference__degenerate_row_with_no_selected_entry__keeps_size_one_and_nan(): + with np.errstate(all="ignore"): + outputs = sibson._sibson_mappings_weights_numpy_reference( + **_overflowing_tables() + ) + + _assert_failed_rows_kept_at_size_one(outputs, points_count=30) + + +def test__numba__degenerate_row_with_no_selected_entry__matches_reference(): + pytest.importorskip("numba") + + kwargs = _overflowing_tables() + with np.errstate(all="ignore"): + reference = sibson._sibson_mappings_weights_numpy_reference(**kwargs) + numba_outputs = sibson._sibson_mappings_weights_numba(**kwargs) + + _assert_failed_rows_kept_at_size_one(numba_outputs, points_count=30) + for expected, actual in zip(reference, numba_outputs): + assert expected.dtype == actual.dtype + assert np.array_equal(expected, actual, equal_nan=True) + + +@pytest.mark.skipif(importlib.util.find_spec("jax") is None, reason="jax not installed") +def test__jax__degenerate_row_with_no_selected_entry__matches_reference(): + import jax.numpy as jnp + + kwargs = _overflowing_tables() + with np.errstate(all="ignore"): + reference = sibson._sibson_mappings_weights_numpy_reference(**kwargs) + + jax_kwargs = { + key: (jnp.asarray(value) if isinstance(value, np.ndarray) else value) + for key, value in kwargs.items() + if key != "circumcircles" + } + jax_outputs = sibson.sibson_mappings_weights_from_tables(xp=jnp, **jax_kwargs) + + _assert_failed_rows_kept_at_size_one(jax_outputs, points_count=30) + for expected, actual in zip(reference[:3], jax_outputs[:3]): + assert np.array_equal(expected, np.asarray(actual), equal_nan=True) diff --git a/test_autoarray/inversion/pixelization/mappers/test_mapper_util_wide_rows.py b/test_autoarray/inversion/pixelization/mappers/test_mapper_util_wide_rows.py new file mode 100644 index 000000000..372b88a50 --- /dev/null +++ b/test_autoarray/inversion/pixelization/mappers/test_mapper_util_wide_rows.py @@ -0,0 +1,328 @@ +""" +Skipping the padding of wide mapping rows (PyAutoArray#629). + +`DelaunayNN` (Sibson) mapping tables are `SIBSON_MAX_NEIGHBORS` (32) columns wide but hold at most ~15 live +entries per row. `adaptive_pixel_signals_from` and `data_slim_to_pixelization_unique_from` now visit only the +first `pix_sizes[i]` entries of every row. These tests pin that the results are bit-identical to the padded +algorithms they replaced, on wide `DelaunayNN`-shaped tables and on 3-wide `Delaunay` tables alike. +""" + +import importlib.util + +import numpy as np +import pytest + +import autoarray as aa +from autoarray.inversion.mappers import mapper_numba_util, mapper_util + +requires_jax = pytest.mark.skipif( + importlib.util.find_spec("jax") is None, reason="jax not installed" +) + + +def _padded_scatter_reference( + pixels, + pixel_weights, + signal_scale, + pix_indexes_for_sub_slim_index, + pix_size_for_sub_slim_index, + slim_index_for_sub_slim_index, + adapt_data, + exclude_zero_weight_mappings=False, +): + """The NumPy branch of `adaptive_pixel_signals_from` before PyAutoArray#629: an `np.add.at` over every padded + slot, with masked entries sent to a discarded out-of-bounds slot.""" + M_sub, B = pix_indexes_for_sub_slim_index.shape + flat_pixidx = pix_indexes_for_sub_slim_index.reshape(-1) + flat_weights = pixel_weights.reshape(-1) + I_sub = np.repeat(np.arange(M_sub), B) + valid = (np.arange(B)[None, :] < pix_size_for_sub_slim_index[:, None]).reshape(-1) + if exclude_zero_weight_mappings: + valid = valid & (flat_weights != 0) + flat_weights = np.where(valid, flat_weights, 0.0) + flat_pixidx = np.where(valid, flat_pixidx, pixels) + flat_data_vals = np.take(adapt_data[slim_index_for_sub_slim_index], I_sub, axis=0) + flat_contrib = flat_data_vals * flat_weights + pixel_signals = np.zeros((pixels + 1,)) + pixel_counts = np.zeros((pixels + 1,)) + np.add.at(pixel_signals, flat_pixidx, flat_contrib) + np.add.at(pixel_counts, flat_pixidx, valid.astype(float)) + pixel_signals = pixel_signals[:pixels] + pixel_counts = pixel_counts[:pixels] + pixel_counts = np.where(pixel_counts > 0, pixel_counts, 1.0) + pixel_signals = pixel_signals / pixel_counts + max_sig = np.max(pixel_signals) + max_sig = np.where(max_sig > 0, max_sig, 1.0) + pixel_signals = pixel_signals / max_sig + return pixel_signals**signal_scale + + +def _unique_full_reset_reference( + data_pixels, + pix_indexes_for_sub_slim_index, + pix_sizes_for_sub_slim_index, + pix_weights_for_sub_slim_index, + pix_pixels, + sub_size, +): + """`data_slim_to_pixelization_unique_from` before PyAutoArray#629, in plain Python: the whole `pix_check` + array is reset for every data pixel.""" + sub_fraction = 1.0 / (sub_size**2.0) + width = int(np.max(pix_sizes_for_sub_slim_index)) * int(np.max(sub_size)) ** 2 + data_to_pix_unique = -1 * np.ones((data_pixels, width)) + data_weights = np.zeros((data_pixels, width)) + pix_lengths = np.zeros(data_pixels) + ip_sub_start = 0 + for ip in range(data_pixels): + pix_check = -1 * np.ones(pix_pixels) + pix_size = 0 + ip_sub_end = ip_sub_start + sub_size[ip] ** 2 + for ip_sub in range(ip_sub_start, ip_sub_end): + for k in range(pix_sizes_for_sub_slim_index[ip_sub]): + pix = pix_indexes_for_sub_slim_index[ip_sub, k] + weight = pix_weights_for_sub_slim_index[ip_sub, k] + if pix_check[pix] > -0.5: + data_weights[ip, int(pix_check[pix])] += sub_fraction[ip] * weight + else: + data_to_pix_unique[ip, pix_size] = pix + data_weights[ip, pix_size] += sub_fraction[ip] * weight + pix_check[pix] = pix_size + pix_size += 1 + ip_sub_start = ip_sub_end + pix_lengths[ip] = pix_size + return data_to_pix_unique, data_weights, pix_lengths + + +def _padded_tables(width, min_size, max_size, data_pixels=60, sub=2, pixels=40, seed=1): + """Sibson-shaped tables: `width` columns, `size` live entries per row (distinct pixels, positive weights + summing to 1 with a few exact zeros), `-1` / `0.0` padding beyond `size`.""" + rng = np.random.default_rng(seed) + sub_pixels = data_pixels * sub**2 + sizes = rng.integers(min_size, max_size + 1, size=sub_pixels).astype(np.int32) + indexes = np.full((sub_pixels, width), -1, dtype=np.int32) + weights = np.zeros((sub_pixels, width)) + for row, size in enumerate(sizes): + indexes[row, :size] = rng.choice(pixels, size=size, replace=False) + live = rng.uniform(0.0, 1.0, size=size) + live[rng.uniform(size=size) < 0.1] = 0.0 + weights[row, :size] = live / max(live.sum(), 1.0e-300) + return dict( + pixels=pixels, + indexes=indexes, + sizes=sizes, + weights=weights, + slim=np.repeat(np.arange(data_pixels), sub**2), + adapt=rng.uniform(0.0, 2.0, size=data_pixels), + sub_size=np.full(data_pixels, sub), + data_pixels=data_pixels, + ) + + +TABLES = { + "delaunay_nn_32_wide": dict(width=32, min_size=1, max_size=15), + "delaunay_3_wide": dict(width=3, min_size=3, max_size=3), + "knn_10_wide_variable": dict(width=10, min_size=1, max_size=10), +} + + +def _signal_kwargs(tables, exclude): + return dict( + pixels=tables["pixels"], + pixel_weights=tables["weights"], + signal_scale=0.7, + pix_indexes_for_sub_slim_index=tables["indexes"], + pix_size_for_sub_slim_index=tables["sizes"], + slim_index_for_sub_slim_index=tables["slim"], + adapt_data=tables["adapt"], + exclude_zero_weight_mappings=exclude, + ) + + +@pytest.mark.parametrize("use_numba", [True, False]) +@pytest.mark.parametrize("exclude", [True, False]) +@pytest.mark.parametrize("table", sorted(TABLES)) +def test__adaptive_pixel_signals_from__live_entries_only__bit_identical_to_padded_scatter( + table, exclude, use_numba, monkeypatch +): + tables = _padded_tables(**TABLES[table]) + kwargs = _signal_kwargs(tables, exclude) + + monkeypatch.setattr(mapper_util, "_use_numba", lambda: use_numba) + + signals = mapper_util.adaptive_pixel_signals_from(**kwargs) + + assert np.array_equal(signals, _padded_scatter_reference(**kwargs)) + + +def test__adaptive_pixel_signals_from__nan_in_padding_only__ignored_like_padded_scatter( + monkeypatch, +): + tables = _padded_tables(**TABLES["delaunay_nn_32_wide"]) + weights = tables["weights"].copy() + weights[:, -1] = np.nan # padding beyond every row's size + tables["weights"] = weights + kwargs = _signal_kwargs(tables, exclude=True) + + for use_numba in (True, False): + monkeypatch.setattr(mapper_util, "_use_numba", lambda: use_numba) + signals = mapper_util.adaptive_pixel_signals_from(**kwargs) + assert np.isfinite(signals).all() + assert np.array_equal(signals, _padded_scatter_reference(**kwargs)) + + +@requires_jax +def test__adaptive_pixel_signals_from__wide_rows__jax_matches_numpy(): + import jax.numpy as jnp + + tables = _padded_tables(**TABLES["delaunay_nn_32_wide"]) + kwargs = _signal_kwargs(tables, exclude=True) + + numpy_signals = mapper_util.adaptive_pixel_signals_from(**kwargs) + jax_signals = np.asarray(mapper_util.adaptive_pixel_signals_from(**kwargs, xp=jnp)) + + np.testing.assert_allclose(jax_signals, numpy_signals, rtol=1.0e-12, atol=0.0) + + +@pytest.mark.parametrize("table", sorted(TABLES)) +def test__data_slim_to_pixelization_unique_from__touched_reset__bit_identical_to_full_reset( + table, +): + tables = _padded_tables(**TABLES[table]) + kwargs = dict( + data_pixels=tables["data_pixels"], + pix_indexes_for_sub_slim_index=tables["indexes"], + pix_sizes_for_sub_slim_index=tables["sizes"], + pix_weights_for_sub_slim_index=tables["weights"], + pix_pixels=tables["pixels"], + sub_size=tables["sub_size"], + ) + + result = mapper_numba_util.data_slim_to_pixelization_unique_from(**kwargs) + expected = _unique_full_reset_reference(**kwargs) + + for actual, reference in zip(result, expected): + assert actual.shape == reference.shape + assert np.array_equal(actual, reference) + + +def _mappers(sub_size=2): + mask = aa.Mask2D.circular(shape_native=(12, 12), pixel_scales=0.5, radius=2.6) + grid = aa.Grid2D.from_mask(mask=mask, over_sample_size=sub_size) + rng = np.random.default_rng(seed=4) + mesh_grid = aa.Grid2DIrregular(rng.uniform(-3.5, 3.5, size=(30, 2))) + adapt_data = aa.Array2D( + values=rng.uniform(0.0, 1.0, size=mask.pixels_in_mask), mask=mask + ) + mappers = {} + for name, mesh in ( + ("delaunay", aa.mesh.Delaunay(pixels=30)), + ("delaunay_nn", aa.mesh.DelaunayNN(pixels=30)), + ): + interpolator = mesh.interpolator_from( + source_plane_data_grid=grid, + source_plane_mesh_grid=mesh_grid, + adapt_data=adapt_data, + ) + mappers[name] = aa.Mapper( + interpolator=interpolator, + regularization=aa.reg.Constant(coefficient=1.0), + ) + return mappers + + +@pytest.mark.parametrize("name", ["delaunay", "delaunay_nn"]) +def test__mapper__real_delaunay_and_delaunay_nn_tables__unchanged(name, monkeypatch): + """Real `Delaunay` (3-wide) and `DelaunayNN` (32-wide) mapper tables: pixel signals through both NumPy + routes and the unique mappings equal the padded / full-reset algorithms bit for bit. + """ + mapper = _mappers()[name] + + indexes = np.asarray(mapper.pix_indexes_for_sub_slim_index) + sizes = np.asarray(mapper.pix_sizes_for_sub_slim_index) + weights = np.asarray(mapper.pix_weights_for_sub_slim_index) + + if name == "delaunay_nn": + assert indexes.shape[1] == 32 and sizes.max() > 3 + + kwargs = dict( + pixels=mapper.pixels, + pixel_weights=weights, + signal_scale=1.0, + pix_indexes_for_sub_slim_index=indexes, + pix_size_for_sub_slim_index=sizes, + slim_index_for_sub_slim_index=np.asarray(mapper.over_sampler.slim_for_sub_slim), + adapt_data=np.asarray(mapper.adapt_data.array), + exclude_zero_weight_mappings=mapper.interpolator.exclude_zero_weight_mappings, + ) + expected_signals = _padded_scatter_reference(**kwargs) + + for use_numba in (True, False): + monkeypatch.setattr(mapper_util, "_use_numba", lambda: use_numba) + assert np.array_equal( + mapper.pixel_signals_from(signal_scale=1.0), expected_signals + ) + + expected_unique = _unique_full_reset_reference( + data_pixels=mapper.over_sampler.mask.pixels_in_mask, + pix_indexes_for_sub_slim_index=indexes, + pix_sizes_for_sub_slim_index=sizes, + pix_weights_for_sub_slim_index=weights, + pix_pixels=mapper.params, + sub_size=np.asarray(mapper.over_sampler.sub_size).astype(int), + ) + unique = mapper.unique_mappings + for actual, reference in zip( + (unique.data_to_pix_unique, unique.data_weights, unique.pix_lengths), + expected_unique, + ): + assert np.array_equal(actual, reference) + + +def test__failed_row_kept_at_size_one__nan_reaches_signals_and_unique_weights( + monkeypatch, +): + """A failed Sibson row carries NaN weights. At `size == 0` every consumer skipped them (the evidence came + out finite and wrong); at `size == 1` (PyAutoArray#629) the NaN reaches the unique-mapping weights that + build the curvature matrix, and the pixel signals that build the regularization.""" + tables = _padded_tables(**TABLES["delaunay_nn_32_wide"]) + indexes, sizes, weights = ( + tables["indexes"].copy(), + tables["sizes"].copy(), + tables["weights"].copy(), + ) + failed_row = 5 + indexes[failed_row, :] = -1 + indexes[failed_row, 0] = 3 + weights[failed_row, :] = np.nan + + def outputs(size): + sizes[failed_row] = size + unique = mapper_numba_util.data_slim_to_pixelization_unique_from( + data_pixels=tables["data_pixels"], + pix_indexes_for_sub_slim_index=indexes, + pix_sizes_for_sub_slim_index=sizes, + pix_weights_for_sub_slim_index=weights, + pix_pixels=tables["pixels"], + sub_size=tables["sub_size"], + ) + signals = mapper_util.adaptive_pixel_signals_from( + pixels=tables["pixels"], + pixel_weights=weights, + signal_scale=1.0, + pix_indexes_for_sub_slim_index=indexes, + pix_size_for_sub_slim_index=sizes, + slim_index_for_sub_slim_index=tables["slim"], + adapt_data=tables["adapt"], + exclude_zero_weight_mappings=True, + ) + return unique[1], signals + + for use_numba in (True, False): + monkeypatch.setattr(mapper_util, "_use_numba", lambda: use_numba) + + data_weights, signals = outputs(size=0) + assert np.isfinite(data_weights).all() and np.isfinite(signals).all() + + data_weights, signals = outputs(size=1) + assert np.isnan(data_weights[failed_row // 4]).any() + assert np.isnan(signals).any()