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()