diff --git a/autoarray/inversion/inversion/imaging_numba/inversion_imaging_numba_util.py b/autoarray/inversion/inversion/imaging_numba/inversion_imaging_numba_util.py index aef50fb9f..7d3a86b13 100644 --- a/autoarray/inversion/inversion/imaging_numba/inversion_imaging_numba_util.py +++ b/autoarray/inversion/inversion/imaging_numba/inversion_imaging_numba_util.py @@ -492,6 +492,44 @@ def curvature_matrix_with_added_to_diag_from( return curvature_matrix +# --------------------------------------------------------------------------- +# Two-stage mapper x mapper threshold (PyAutoArray#507 step 2) +# --------------------------------------------------------------------------- +# +# `curvature_matrix_via_sparse_operator_from` has two implementations with +# identical contracts (see their docstrings): the direct quadruple loop, whose +# cost is `sum_pairs u0 * u1` irregular read-modify-writes and does not depend +# on `pix_pixels`; and the two-stage source-space accumulator, whose cost grows +# linearly in `pix_pixels` but is dominated by contiguous, vectorisable AXPYs. +# The two-stage form therefore must lose eventually, and this is the measured +# point at which it is switched off. +# +# Measured by sweeping `pix_pixels` on the two production HST geometries -- the +# real 15361-pixel sparse operator and real mapper mappings, with only the +# source-space extent varied (autolens_profiling, OMP_NUM_THREADS=1). Speed-up +# of two-stage over direct: +# +# pix_pixels 128 784 1250 2048 4096 6144 8192 +# rectangular u0=4.0 3.10 2.93 2.47 2.10 1.36 1.12 1.04 +# delaunay u0=1.55 1.79 1.72 1.53 1.33 1.04 1.01 0.98 +# +# So the crossover is geometry-dependent -- the narrow-mapping Delaunay +# geometry reaches parity near 8192 source pixels while the wide-mapping +# bilinear one is still ahead there -- and there is *no* crossover anywhere +# inside the range any PyAuto pixelization is run at. The threshold below is +# deliberately the conservative end of the measured envelope rather than a +# fitted crossover: at 4096 both geometries still favour the two-stage kernel, +# and above it the two forms are within a few per cent of each other anyway, so +# picking the direct one costs nothing. It is *not* a silent heuristic: it is a +# single explicit constant, and it can be overridden per call through the +# `two_stage_max_pix_pixels` argument (which is how both branches are tested). +# +# For scale: a 4096-pixel source makes `F` a 134 MB dense matrix whose Cholesky +# factorisation alone is ~2e10 flops, so the inversion is long past being +# dominated by this kernel by the time the threshold is in play. +CURVATURE_TWO_STAGE_MAX_PIX_PIXELS = 4096 + + @numba_util.jit() def curvature_matrix_via_sparse_operator_from( psf_precision_operator: np.ndarray, @@ -501,6 +539,95 @@ def curvature_matrix_via_sparse_operator_from( data_weights: np.ndarray, pix_lengths: np.ndarray, pix_pixels: int, + two_stage_max_pix_pixels: int = CURVATURE_TWO_STAGE_MAX_PIX_PIXELS, +) -> np.ndarray: + """ + Returns the mapper x mapper block of the curvature matrix `F` (see Warren & Dye + 2003) from the sparse PSF precision operator. + + This is the entry point the inversion calls. It selects between two + implementations with identical inputs, outputs and contracts: + + - `curvature_matrix_via_sparse_operator_two_stage_from`, which accumulates a + dense source-space vector per data pixel and then writes whole rows of `F` + with contiguous AXPYs. Faster for every source-space size any PyAuto + pixelization uses (2.9x at the HST rectangular fiducial), because it + replaces `sum_pairs u0 * u1` irregular read-modify-writes into a multi-MB + matrix with `sum_pairs u1` L1-resident scatters plus vectorisable dense + adds. + - `curvature_matrix_via_sparse_operator_direct_from`, the quadruple loop, + whose cost does not grow with `pix_pixels` and which therefore takes over + above `two_stage_max_pix_pixels`. + + The two agree to floating-point reassociation, not bit-identically: the + two-stage form sums the same products in a different order. The measured + maximum relative difference on the production HST geometries is ~4e-13, + three orders inside the `rtol=1e-6` the likelihood is pinned at. + + Parameters + ---------- + psf_precision_operator + A matrix that precomputes the values for fast computation of the curvature matrix in a memory efficient way. + psf_precision_indexes + The image-pixel indexes of the values stored in the w tilde preload matrix, which are used to compute + the weights of the data values when computing the curvature matrix. + psf_precision_lengths + The number of image pixels in every row of `psf_precision_operator`, which is iterated over when computing the + curvature matrix. + data_to_pix_unique + An array that maps every data pixel index (e.g. the masked image pixel indexes in 1D) to its unique set of + pixelization pixel indexes (see `data_slim_to_pixelization_unique_from`). + data_weights + For every unique mapping between a set of data sub-pixels and a pixelization pixel, the weight of these mapping + based on the number of sub-pixels that map to pixelization pixel. + pix_lengths + A 1D array describing how many unique pixels each data pixel maps too, which is used to iterate over + `data_to_pix_unique` and `data_weights`. + pix_pixels + The total number of pixels in the pixelization that reconstructs the data. + two_stage_max_pix_pixels + The largest `pix_pixels` for which the two-stage implementation is used; + above it the direct loop is. Defaults to the measured + `CURVATURE_TWO_STAGE_MAX_PIX_PIXELS`; exposed so that both branches can + be exercised without building a source space of that size. + + Returns + ------- + ndarray + The curvature matrix `F` (see Warren & Dye 2003). + """ + + if pix_pixels <= two_stage_max_pix_pixels: + return curvature_matrix_via_sparse_operator_two_stage_from( + psf_precision_operator=psf_precision_operator, + psf_precision_indexes=psf_precision_indexes, + psf_precision_lengths=psf_precision_lengths, + data_to_pix_unique=data_to_pix_unique, + data_weights=data_weights, + pix_lengths=pix_lengths, + pix_pixels=pix_pixels, + ) + + return curvature_matrix_via_sparse_operator_direct_from( + psf_precision_operator=psf_precision_operator, + psf_precision_indexes=psf_precision_indexes, + psf_precision_lengths=psf_precision_lengths, + data_to_pix_unique=data_to_pix_unique, + data_weights=data_weights, + pix_lengths=pix_lengths, + pix_pixels=pix_pixels, + ) + + +@numba_util.jit() +def curvature_matrix_via_sparse_operator_direct_from( + psf_precision_operator: np.ndarray, + psf_precision_indexes: np.ndarray, + psf_precision_lengths: np.ndarray, + data_to_pix_unique: np.ndarray, + data_weights: np.ndarray, + pix_lengths: np.ndarray, + pix_pixels: int, ) -> np.ndarray: """ Returns the curvature matrix `F` (see Warren & Dye 2003) by computing it using `psf_precision_operator` @@ -554,6 +681,203 @@ def curvature_matrix_via_sparse_operator_from( curvature_index = 0 + for data_0 in range(data_pixels): + + # Hoisted once per data pixel and once per stored pair (PyAutoArray#507 + # step 1). The unhoisted form re-read `data_weights[data_1, pix_1_index]` + # and `data_to_pix_unique[data_1, pix_1_index]` inside the innermost + # loop, so every one of `data_1`'s u1 mappings was gathered u0 times from + # a wide-stride 2-D array; and `curvature_matrix[pix_0, pix_1]` re-did the + # row offset multiply on every accumulation. Taking 1-D row views and the + # `curvature_matrix[pix_0]` row view leaves the arithmetic untouched -- + # the accumulated expression below is operand-for-operand the one it + # replaces, so the result is bit-identical, which + # `curvature_matrix_via_sparse_operator_reference_from` and its test + # assert directly. + pix_lengths_0 = pix_lengths[data_0] + pix_row_0 = data_to_pix_unique[data_0] + weight_row_0 = data_weights[data_0] + + for data_1_index in range(psf_precision_lengths[data_0]): + data_1 = psf_precision_indexes[curvature_index] + psf_precision_value = psf_precision_operator[curvature_index] + + curvature_index += 1 + + pix_lengths_1 = pix_lengths[data_1] + pix_row_1 = data_to_pix_unique[data_1] + weight_row_1 = data_weights[data_1] + + for pix_0_index in range(pix_lengths_0): + data_0_weight = weight_row_0[pix_0_index] + curvature_row = curvature_matrix[pix_row_0[pix_0_index]] + + for pix_1_index in range(pix_lengths_1): + curvature_row[pix_row_1[pix_1_index]] += ( + data_0_weight + * weight_row_1[pix_1_index] + * psf_precision_value + ) + + for i in range(pix_pixels): + for j in range(i, pix_pixels): + curvature_matrix[i, j] += curvature_matrix[j, i] + + for i in range(pix_pixels): + for j in range(i, pix_pixels): + curvature_matrix[j, i] = curvature_matrix[i, j] + + return curvature_matrix + + +@numba_util.jit() +def curvature_matrix_via_sparse_operator_two_stage_from( + psf_precision_operator: np.ndarray, + psf_precision_indexes: np.ndarray, + psf_precision_lengths: np.ndarray, + data_to_pix_unique: np.ndarray, + data_weights: np.ndarray, + pix_lengths: np.ndarray, + pix_pixels: int, +) -> np.ndarray: + """ + The mapper x mapper block of `F`, computed in two stages via a dense + source-space accumulator (PyAutoArray#507 step 2). Same inputs, same outputs + and same halved-diagonal / `A + A.T` contract as + `curvature_matrix_via_sparse_operator_direct_from`; only the order of the + summation differs, so the two agree to floating-point reassociation + (`rtol=1e-6` is the pinned tolerance; the observed difference is far smaller) + rather than bit-identically. + + The direct kernel's innermost work is + + for each stored pair (data_0, data_1): + for each of data_0's u0 mappings (pix_0, w0): + for each of data_1's u1 mappings (pix_1, w1): + F[pix_0, pix_1] += w0 * w1 * W(data_0, data_1) + + -- `u0 * u1` irregular read-modify-writes per stored pair, scattered across a + `pix_pixels ** 2` matrix that at HST resolution is 4.9 MB, far outside L2. + Measured at the HST rectangular fiducial that is 1.77e8 accumulations at + ~1.4 ns each. + + The same sum factorises, because `w0` does not depend on `data_1`: + + a[p] = sum over data_0's stored pairs of W(data_0, data_1) * w1(data_1, p) + F[pix_0, :] += w0 * a[:] for each of data_0's u0 mappings + + Stage 1 accumulates `a` -- one scatter per (stored pair, mapping of data_1), + i.e. `u1` rather than `u0 * u1` -- into an L1-resident `pix_pixels` vector. + Stage 2 is `u0` contiguous, vectorisable AXPYs of length `pix_pixels` into + whole rows of `F`. + + The trade is `sum_pairs u0 * u1` scattered RMWs for + `sum_pairs u1` L1 scatters plus `(sum_data u0 + data_pixels) * pix_pixels` + dense ops. It is therefore *not* unconditionally faster: it wins when the + source space is small relative to the PSF overlap and the mappings are wide + (bilinear, u0 = 4), and loses when the source space is large and the mappings + narrow (barycentric Delaunay, u0 ~ 1.55). `curvature_matrix_via_sparse_operator_from` + chooses between the two from the measured geometry -- see + `CURVATURE_TWO_STAGE_COST_RATIO_THRESHOLD`. + + Parameters and return value are exactly those of + `curvature_matrix_via_sparse_operator_direct_from`. + """ + + data_pixels = psf_precision_lengths.shape[0] + + curvature_matrix = np.zeros((pix_pixels, pix_pixels)) + + # The per-data-pixel source-space accumulator, allocated once and zeroed + # after every data pixel (it is left clean by the loop below, so the + # allocation's own zeros are correct for the first data pixel). + source_accumulator = np.zeros(pix_pixels) + + curvature_index = 0 + + for data_0 in range(data_pixels): + + pair_length_0 = psf_precision_lengths[data_0] + pix_lengths_0 = pix_lengths[data_0] + + if pix_lengths_0 == 0: + # No mappings for this data pixel: it contributes nothing, but its + # stored pairs must still be stepped over. + curvature_index += pair_length_0 + continue + + # -- stage 1: accumulate the dense source-space vector --------------- + for data_1_index in range(pair_length_0): + data_1 = psf_precision_indexes[curvature_index] + psf_precision_value = psf_precision_operator[curvature_index] + + curvature_index += 1 + + pix_row_1 = data_to_pix_unique[data_1] + weight_row_1 = data_weights[data_1] + + for pix_1_index in range(pix_lengths[data_1]): + source_accumulator[pix_row_1[pix_1_index]] += ( + weight_row_1[pix_1_index] * psf_precision_value + ) + + # -- stage 2: one contiguous AXPY per mapping of data_0 -------------- + pix_row_0 = data_to_pix_unique[data_0] + weight_row_0 = data_weights[data_0] + + for pix_0_index in range(pix_lengths_0): + data_0_weight = weight_row_0[pix_0_index] + curvature_row = curvature_matrix[pix_row_0[pix_0_index]] + + for pix_1 in range(pix_pixels): + curvature_row[pix_1] += data_0_weight * source_accumulator[pix_1] + + for pix_1 in range(pix_pixels): + source_accumulator[pix_1] = 0.0 + + for i in range(pix_pixels): + for j in range(i, pix_pixels): + curvature_matrix[i, j] += curvature_matrix[j, i] + + for i in range(pix_pixels): + for j in range(i, pix_pixels): + curvature_matrix[j, i] = curvature_matrix[i, j] + + return curvature_matrix + + +@numba_util.jit() +def curvature_matrix_via_sparse_operator_reference_from( + psf_precision_operator: np.ndarray, + psf_precision_indexes: np.ndarray, + psf_precision_lengths: np.ndarray, + data_to_pix_unique: np.ndarray, + data_weights: np.ndarray, + pix_lengths: np.ndarray, + pix_pixels: int, +) -> np.ndarray: + """ + The straightforward quadruple loop that `curvature_matrix_via_sparse_operator_from` + is an optimisation of, retained verbatim as the reference the optimised kernel is + tested against. + + It is not called by the inversion. It exists so that the restructurings applied to + the production kernel -- the hoisted row views of PyAutoArray#507 step 1, and the + two-stage source-space accumulator of step 2 -- have a fixed, unrestructured + definition to be asserted against, rather than only the dense `M.T W M` oracle + (which is a far smaller problem than the kernel is ever run on, since it + materializes the [image_pixels, image_pixels] precision operator). + + Parameters and return value are exactly those of + `curvature_matrix_via_sparse_operator_from`. + """ + + data_pixels = psf_precision_lengths.shape[0] + + curvature_matrix = np.zeros((pix_pixels, pix_pixels)) + + curvature_index = 0 + for data_0 in range(data_pixels): for data_1_index in range(psf_precision_lengths[data_0]): data_1 = psf_precision_indexes[curvature_index] diff --git a/autoarray/operators/over_sampling/over_sampler.py b/autoarray/operators/over_sampling/over_sampler.py index b840e1b6e..f28cf0fd9 100644 --- a/autoarray/operators/over_sampling/over_sampler.py +++ b/autoarray/operators/over_sampling/over_sampler.py @@ -167,13 +167,41 @@ def __init__(self, mask: Mask2D, sub_size: Union[int, Array2D]): ): self.segment_ids[start:end] = seg_id - @property + @cached_property def sub_is_uniform(self) -> bool: """ Returns True if the sub_size is uniform across all pixels in the mask. + + The `sub_size` is fixed when the `OverSampler` is created, so this is cached: it is evaluated once per + instance instead of on every call to `binned_array_2d_from` (which is called once per light profile + evaluation, and therefore many hundreds of times per likelihood evaluation). """ return np.all(np.isclose(self.sub_size, self.sub_size[0])) + @cached_property + def binned_counts(self) -> np.ndarray: + """ + The number of sub-pixels which are binned up into every pixel of the mask, used as the divisor when + binning an over sampled array via `binned_array_2d_from` for a non-uniform `sub_size`. + + Segments which contain no sub-pixels have their count set to 1, so that the division by this array does + not raise a divide-by-zero (the corresponding sum is zero, so the binned value is zero either way). + + This is a function of the `segment_ids` alone, which are fixed when the `OverSampler` is created, so it is + computed once per instance instead of on every call to `binned_array_2d_from`. + + The returned array is read-only, because it is shared by every call and must never be mutated in place. + """ + counts = np.bincount(self.segment_ids, minlength=self.mask.pixels_in_mask) + + # Avoid division by zero (on a copy, because `np.bincount` output is not reused elsewhere but the cached + # array below is handed out to every caller). + + counts = np.maximum(counts, 1) + counts.flags.writeable = False + + return counts + def tree_flatten(self): return (self.mask, self.sub_size), () @@ -263,13 +291,8 @@ def binned_array_2d_from(self, array: Array2D, xp=np) -> "Array2D": self.segment_ids, weights=array, minlength=self.mask.pixels_in_mask ) - # Count number of items per segment - counts = np.bincount( - self.segment_ids, minlength=self.mask.pixels_in_mask - ) - - # Avoid division by zero - counts[counts == 0] = 1 + # Count number of items per segment (cached, as it depends only on the segment ids) + counts = self.binned_counts binned_array_2d = sums / counts diff --git a/test_autoarray/inversion/inversion/imaging/test_inversion_imaging_util.py b/test_autoarray/inversion/inversion/imaging/test_inversion_imaging_util.py index 5f9adf5fa..cdb247ab3 100644 --- a/test_autoarray/inversion/inversion/imaging/test_inversion_imaging_util.py +++ b/test_autoarray/inversion/inversion/imaging/test_inversion_imaging_util.py @@ -530,3 +530,490 @@ def test__convolver_reversed_kernel__is_the_kernel_reversed_and_convolves_as_a_c # Cached, so the reversed kernel and its FFT geometry are built once. assert convolver.reversed_kernel is convolver.reversed_kernel + + +# ------------------------------------------------------------------ +# The mapper x mapper block of `F` (`curvature_matrix_via_sparse_operator_from`). +# +# This is the single most expensive step of the numba CPU imaging likelihood at +# HST resolution (0.255 s of a 0.63 s evaluation, PyAutoArray#507 step 0), and +# until these tests it had no direct unit test at all -- only the end-to-end +# inversion fixtures, which cannot distinguish a restructured kernel from a +# subtly wrong one. +# +# The reference the kernel is pinned against is the definition it is an +# optimisation of: +# +# F = M.T @ W @ M +# +# with `W` the *dense* [image_pixels, image_pixels] precision operator from +# `psf_precision_operator_from` and `M` the dense [image_pixels, pix_pixels] +# mapping matrix that `(data_to_pix_unique, data_weights, pix_lengths)` encodes. +# Nothing in the reference re-derives the kernel: it does not know about the +# sparse upper-triangle storage, the halved diagonal or the `A + A.T` fold. +# +# The kernels are asymmetric and non-square (`KERNELS_ODD`), so a transposed +# axis in the operator build cannot pass. +# ------------------------------------------------------------------ + + +def _mapping_matrix_from(data_to_pix_unique, data_weights, pix_lengths, pix_pixels): + """The dense mapping matrix `M` that the unique-mappings triple encodes.""" + mapping_matrix = np.zeros((pix_lengths.shape[0], pix_pixels)) + + for data_index in range(pix_lengths.shape[0]): + for pix_index in range(pix_lengths[data_index]): + mapping_matrix[ + data_index, data_to_pix_unique[data_index, pix_index] + ] += data_weights[data_index, pix_index] + + return mapping_matrix + + +def _sparse_operator_mappings(seed=507, shape=(6, 6), pix_pixels=7, max_lengths=3): + """A small fully-unmasked dataset plus random non-uniform unique mappings.""" + rng = np.random.default_rng(seed) + + noise_map = 0.5 + rng.random(shape) * 2.0 + native_index_for_slim_index = np.array( + [[y, x] for y in range(shape[0]) for x in range(shape[1])] + ) + + data_pixels = native_index_for_slim_index.shape[0] + + pix_lengths = rng.integers(1, max_lengths + 1, size=data_pixels).astype("int") + data_to_pix_unique = rng.integers( + 0, pix_pixels, size=(data_pixels, max_lengths) + ).astype("int") + data_weights = rng.random(size=(data_pixels, max_lengths)) + + return ( + noise_map, + native_index_for_slim_index, + data_to_pix_unique, + data_weights, + pix_lengths, + ) + + +@pytest.mark.parametrize("kernel", KERNELS_ODD, ids=KERNEL_IDS) +def test__curvature_matrix_via_sparse_operator_from__matches_dense_psf_precision_operator( + kernel, +): + pix_pixels = 7 + + ( + noise_map, + native_index_for_slim_index, + data_to_pix_unique, + data_weights, + pix_lengths, + ) = _sparse_operator_mappings(pix_pixels=pix_pixels) + + ( + psf_precision_operator, + psf_precision_indexes, + psf_precision_lengths, + ) = aa.util.inversion_imaging_numba.psf_precision_operator_sparse_from( + noise_map_native=noise_map, + kernel_native=kernel, + native_index_for_slim_index=native_index_for_slim_index, + ) + + curvature_matrix = ( + aa.util.inversion_imaging_numba.curvature_matrix_via_sparse_operator_from( + psf_precision_operator=psf_precision_operator, + psf_precision_indexes=psf_precision_indexes.astype("int"), + psf_precision_lengths=psf_precision_lengths.astype("int"), + data_to_pix_unique=data_to_pix_unique, + data_weights=data_weights, + pix_lengths=pix_lengths, + pix_pixels=pix_pixels, + ) + ) + + psf_precision_operator_dense = ( + aa.util.inversion_imaging_numba.psf_precision_operator_from( + noise_map_native=noise_map, + kernel_native=kernel, + native_index_for_slim_index=native_index_for_slim_index, + ) + ) + + mapping_matrix = _mapping_matrix_from( + data_to_pix_unique=data_to_pix_unique, + data_weights=data_weights, + pix_lengths=pix_lengths, + pix_pixels=pix_pixels, + ) + + curvature_matrix_dense = np.dot( + mapping_matrix.T, np.dot(psf_precision_operator_dense, mapping_matrix) + ) + + assert curvature_matrix == pytest.approx(curvature_matrix_dense, rel=1.0e-10) + + +def test__curvature_matrix_via_sparse_operator_from__is_symmetric_and_the_diagonal_is_not_double_counted(): + """ + Two contracts the kernel's `A + A.T` fold depends on, asserted separately from + the dense oracle above so a future restructuring cannot satisfy one by breaking + the other. + + 1. The returned matrix is *exactly* symmetric -- not to a tolerance. The + inversion's `curvature_matrix` runs no global symmetrizing pass over the + assembled F, so this block must come out symmetric on its own. + + 2. `psf_precision_operator_sparse_from` halves the `ip0 == ip1` entries because + the fold doubles the diagonal of the accumulated matrix. A single data pixel + mapping to a single source pixel isolates exactly that pair: `F[0, 0]` must + be `w**2 * W[0, 0]`, and would come out twice that if either half of the + contract were dropped. + """ + kernel = np.arange(1.0, 16.0).reshape(3, 5) + + # -- 1. exact symmetry ------------------------------------------------ + pix_pixels = 7 + + ( + noise_map, + native_index_for_slim_index, + data_to_pix_unique, + data_weights, + pix_lengths, + ) = _sparse_operator_mappings(pix_pixels=pix_pixels) + + ( + psf_precision_operator, + psf_precision_indexes, + psf_precision_lengths, + ) = aa.util.inversion_imaging_numba.psf_precision_operator_sparse_from( + noise_map_native=noise_map, + kernel_native=kernel, + native_index_for_slim_index=native_index_for_slim_index, + ) + + curvature_matrix = ( + aa.util.inversion_imaging_numba.curvature_matrix_via_sparse_operator_from( + psf_precision_operator=psf_precision_operator, + psf_precision_indexes=psf_precision_indexes.astype("int"), + psf_precision_lengths=psf_precision_lengths.astype("int"), + data_to_pix_unique=data_to_pix_unique, + data_weights=data_weights, + pix_lengths=pix_lengths, + pix_pixels=pix_pixels, + ) + ) + + assert np.array_equal(curvature_matrix, curvature_matrix.T) + + # -- 2. the halved diagonal -------------------------------------------- + noise_map_1 = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]) + native_index_for_slim_index_1 = np.array([[1, 1]]) + + weight = 0.75 + + ( + psf_precision_operator_1, + psf_precision_indexes_1, + psf_precision_lengths_1, + ) = aa.util.inversion_imaging_numba.psf_precision_operator_sparse_from( + noise_map_native=noise_map_1, + kernel_native=kernel, + native_index_for_slim_index=native_index_for_slim_index_1, + ) + + curvature_matrix_1 = ( + aa.util.inversion_imaging_numba.curvature_matrix_via_sparse_operator_from( + psf_precision_operator=psf_precision_operator_1, + psf_precision_indexes=psf_precision_indexes_1.astype("int"), + psf_precision_lengths=psf_precision_lengths_1.astype("int"), + data_to_pix_unique=np.array([[0]]), + data_weights=np.array([[weight]]), + pix_lengths=np.array([1]), + pix_pixels=1, + ) + ) + + psf_precision_operator_dense_1 = ( + aa.util.inversion_imaging_numba.psf_precision_operator_from( + noise_map_native=noise_map_1, + kernel_native=kernel, + native_index_for_slim_index=native_index_for_slim_index_1, + ) + ) + + assert curvature_matrix_1 == pytest.approx( + np.array([[weight**2.0 * psf_precision_operator_dense_1[0, 0]]]), rel=1.0e-12 + ) + + +@pytest.mark.parametrize("kernel", KERNELS_ODD, ids=KERNEL_IDS) +def test__curvature_matrix_via_sparse_operator_from__matches_the_reference_kernel_bit_identically( + kernel, +): + """ + `curvature_matrix_via_sparse_operator_reference_from` is the unrestructured + quadruple loop the production kernel is an optimisation of. + + `curvature_matrix_via_sparse_operator_direct_from` is the branch of the public + `curvature_matrix_via_sparse_operator_from` taken above + `CURVATURE_TWO_STAGE_MAX_PIX_PIXELS`. + + The PyAutoArray#507 step-1 restructuring hoists `data_1`'s `(pix, weight)` pairs + into 1-D row views (they were re-gathered u0 times per stored pair from + wide-stride 2-D arrays) and takes `curvature_matrix[pix_0]` as a row view. It + leaves the accumulated expression operand-for-operand unchanged, so the result + must be *bit-identical*, not merely close: floating-point addition is not + associative, and `np.array_equal` is what catches a reassociation that a + tolerance would wave through. + """ + pix_pixels = 7 + + ( + noise_map, + native_index_for_slim_index, + data_to_pix_unique, + data_weights, + pix_lengths, + ) = _sparse_operator_mappings(pix_pixels=pix_pixels) + + ( + psf_precision_operator, + psf_precision_indexes, + psf_precision_lengths, + ) = aa.util.inversion_imaging_numba.psf_precision_operator_sparse_from( + noise_map_native=noise_map, + kernel_native=kernel, + native_index_for_slim_index=native_index_for_slim_index, + ) + + kwargs = dict( + psf_precision_operator=psf_precision_operator, + psf_precision_indexes=psf_precision_indexes.astype("int"), + psf_precision_lengths=psf_precision_lengths.astype("int"), + data_to_pix_unique=data_to_pix_unique, + data_weights=data_weights, + pix_lengths=pix_lengths, + pix_pixels=pix_pixels, + ) + + curvature_matrix = aa.util.inversion_imaging_numba.curvature_matrix_via_sparse_operator_direct_from( + **kwargs + ) + + curvature_matrix_reference = aa.util.inversion_imaging_numba.curvature_matrix_via_sparse_operator_reference_from( + **kwargs + ) + + assert np.array_equal(curvature_matrix, curvature_matrix_reference) + + +# ------------------------------------------------------------------ +# The two-stage mapper x mapper kernel (PyAutoArray#507 step 2). +# +# `curvature_matrix_via_sparse_operator_two_stage_from` sums exactly the same +# products as the direct loop but in a different order (a dense source-space +# accumulator per data pixel, then one contiguous AXPY per mapping), so it agrees +# to floating-point reassociation rather than bit-identically. rtol=1e-6 is the +# tolerance the likelihood itself is pinned at; the observed difference on the +# production HST geometries is ~4e-13. +# +# Control runs recorded against these tests, on the correct kernel, in the +# commit message. +# ------------------------------------------------------------------ + + +@pytest.mark.parametrize("kernel", KERNELS_ODD, ids=KERNEL_IDS) +@pytest.mark.parametrize("pix_pixels", [1, 7, 64]) +def test__curvature_matrix_via_sparse_operator_two_stage_from__matches_reference_kernel( + kernel, pix_pixels +): + """ + `pix_pixels = 1` is the degenerate source space (every data pixel maps to the + same pixel, so the accumulator is a scalar and the fold is the diagonal alone); + 7 is smaller than the per-data-pixel mapping footprint; 64 is larger than it, + so the accumulator is sparse within itself. All three sit below + `CURVATURE_TWO_STAGE_MAX_PIX_PIXELS`, which the dispatcher test below covers + from both sides. + """ + ( + noise_map, + native_index_for_slim_index, + data_to_pix_unique, + data_weights, + pix_lengths, + ) = _sparse_operator_mappings(pix_pixels=pix_pixels) + + ( + psf_precision_operator, + psf_precision_indexes, + psf_precision_lengths, + ) = aa.util.inversion_imaging_numba.psf_precision_operator_sparse_from( + noise_map_native=noise_map, + kernel_native=kernel, + native_index_for_slim_index=native_index_for_slim_index, + ) + + kwargs = dict( + psf_precision_operator=psf_precision_operator, + psf_precision_indexes=psf_precision_indexes.astype("int"), + psf_precision_lengths=psf_precision_lengths.astype("int"), + data_to_pix_unique=data_to_pix_unique, + data_weights=data_weights, + pix_lengths=pix_lengths, + pix_pixels=pix_pixels, + ) + + curvature_matrix_two_stage = aa.util.inversion_imaging_numba.curvature_matrix_via_sparse_operator_two_stage_from( + **kwargs + ) + + curvature_matrix_reference = aa.util.inversion_imaging_numba.curvature_matrix_via_sparse_operator_reference_from( + **kwargs + ) + + assert curvature_matrix_two_stage == pytest.approx( + curvature_matrix_reference, rel=1.0e-6 + ) + + # The fold contract must survive the reformulation: exactly symmetric, and the + # diagonal not double counted (both are what `A + A.T` on the halved-diagonal + # accumulation buys). + assert np.array_equal(curvature_matrix_two_stage, curvature_matrix_two_stage.T) + + +@pytest.mark.parametrize("kernel", KERNELS_ODD, ids=KERNEL_IDS) +def test__curvature_matrix_via_sparse_operator_two_stage_from__matches_dense_psf_precision_operator( + kernel, +): + """ + The two-stage kernel against the same dense `M.T @ W @ M` oracle the direct + kernel is pinned to, so it is anchored to the definition and not only to the + other implementation. + """ + pix_pixels = 7 + + ( + noise_map, + native_index_for_slim_index, + data_to_pix_unique, + data_weights, + pix_lengths, + ) = _sparse_operator_mappings(pix_pixels=pix_pixels) + + ( + psf_precision_operator, + psf_precision_indexes, + psf_precision_lengths, + ) = aa.util.inversion_imaging_numba.psf_precision_operator_sparse_from( + noise_map_native=noise_map, + kernel_native=kernel, + native_index_for_slim_index=native_index_for_slim_index, + ) + + curvature_matrix = aa.util.inversion_imaging_numba.curvature_matrix_via_sparse_operator_two_stage_from( + psf_precision_operator=psf_precision_operator, + psf_precision_indexes=psf_precision_indexes.astype("int"), + psf_precision_lengths=psf_precision_lengths.astype("int"), + data_to_pix_unique=data_to_pix_unique, + data_weights=data_weights, + pix_lengths=pix_lengths, + pix_pixels=pix_pixels, + ) + + psf_precision_operator_dense = ( + aa.util.inversion_imaging_numba.psf_precision_operator_from( + noise_map_native=noise_map, + kernel_native=kernel, + native_index_for_slim_index=native_index_for_slim_index, + ) + ) + + mapping_matrix = _mapping_matrix_from( + data_to_pix_unique=data_to_pix_unique, + data_weights=data_weights, + pix_lengths=pix_lengths, + pix_pixels=pix_pixels, + ) + + curvature_matrix_dense = np.dot( + mapping_matrix.T, np.dot(psf_precision_operator_dense, mapping_matrix) + ) + + assert curvature_matrix == pytest.approx(curvature_matrix_dense, rel=1.0e-6) + + +def test__curvature_matrix_via_sparse_operator_from__dispatches_on_the_two_stage_threshold(): + """ + Both branches of the public entry point, exercised from either side of + `CURVATURE_TWO_STAGE_MAX_PIX_PIXELS` without building a 4096-pixel source + space: the threshold is a per-call argument defaulting to the module constant, + precisely so that this test does not have to allocate a 134 MB matrix to reach + the other branch. + + Each branch is asserted *bit-identical* to the implementation it is supposed to + have called -- so a dispatcher that quietly ran the wrong one, or ran neither + faithfully, cannot pass, whereas an rtol comparison between the two branches + would (they agree to ~1e-13). + """ + pix_pixels = 7 + + ( + noise_map, + native_index_for_slim_index, + data_to_pix_unique, + data_weights, + pix_lengths, + ) = _sparse_operator_mappings(pix_pixels=pix_pixels) + + ( + psf_precision_operator, + psf_precision_indexes, + psf_precision_lengths, + ) = aa.util.inversion_imaging_numba.psf_precision_operator_sparse_from( + noise_map_native=np.arange(1.0, 37.0).reshape(6, 6), + kernel_native=np.arange(1.0, 16.0).reshape(3, 5), + native_index_for_slim_index=native_index_for_slim_index, + ) + + kwargs = dict( + psf_precision_operator=psf_precision_operator, + psf_precision_indexes=psf_precision_indexes.astype("int"), + psf_precision_lengths=psf_precision_lengths.astype("int"), + data_to_pix_unique=data_to_pix_unique, + data_weights=data_weights, + pix_lengths=pix_lengths, + pix_pixels=pix_pixels, + ) + + # Default threshold (4096) -> pix_pixels of 7 is below it -> two-stage. + assert np.array_equal( + aa.util.inversion_imaging_numba.curvature_matrix_via_sparse_operator_from( + **kwargs + ), + aa.util.inversion_imaging_numba.curvature_matrix_via_sparse_operator_two_stage_from( + **kwargs + ), + ) + + # Threshold pushed below pix_pixels -> the direct loop. + assert np.array_equal( + aa.util.inversion_imaging_numba.curvature_matrix_via_sparse_operator_from( + two_stage_max_pix_pixels=pix_pixels - 1, **kwargs + ), + aa.util.inversion_imaging_numba.curvature_matrix_via_sparse_operator_direct_from( + **kwargs + ), + ) + + # ... and the two branches are not the same array, so the assertions above are + # not both trivially satisfied by a single implementation. + assert not np.array_equal( + aa.util.inversion_imaging_numba.curvature_matrix_via_sparse_operator_two_stage_from( + **kwargs + ), + aa.util.inversion_imaging_numba.curvature_matrix_via_sparse_operator_direct_from( + **kwargs + ), + ) diff --git a/test_autoarray/inversion/inversion/test_curvature_matrix_func_list_blocks.py b/test_autoarray/inversion/inversion/test_curvature_matrix_func_list_blocks.py index 585216dcb..680c0da60 100644 --- a/test_autoarray/inversion/inversion/test_curvature_matrix_func_list_blocks.py +++ b/test_autoarray/inversion/inversion/test_curvature_matrix_func_list_blocks.py @@ -290,3 +290,125 @@ def test__curvature_matrix_mapper_func_blocks__matches_dense_kernel_and_places_t assert curvature_matrix[: mapper.params, : mapper.params] == pytest.approx( np.zeros((mapper.params, mapper.params)), abs=1.0e-12 ) + + +# The mapper x mapper block of the numba sparse imaging inversion is computed by +# `curvature_matrix_via_sparse_operator_from`, which since PyAutoArray#507 step 2 +# dispatches to a two-stage source-space accumulator for every source size a +# PyAuto pixelization uses. The kernel-level tests live beside the kernel +# (`imaging/test_inversion_imaging_util.py`); this asserts the same thing one +# level up -- that what the inversion actually assembles into `F` still matches +# the retained reference loop, with the mapper's param range placed correctly. + + +class StubSparseOperator: + def __init__(self, psf_precision_operator_sparse, indexes, lengths): + self.psf_precision_operator_sparse = psf_precision_operator_sparse + self.indexes = indexes + self.lengths = lengths + + +class StubInversionMapperDiag(InversionImagingSparseNumba): + """Bypasses the real constructor: `_curvature_matrix_mapper_diag` needs only the + mapper list, its param range, the total params and the sparse operator.""" + + def __init__(self, mapper, sparse_operator, total_params, mapper_offset): + self._mapper = mapper + self._sparse_operator = sparse_operator + self._total_params = total_params + self._mapper_offset = mapper_offset + + @property + def sparse_operator(self): + return self._sparse_operator + + @property + def total_params(self): + return self._total_params + + def has(self, cls): + return cls is Mapper + + def total(self, cls): + return 1 if cls is Mapper else 0 + + def cls_list_from(self, cls): + return [self._mapper] if cls is Mapper else [] + + def param_range_list_from(self, cls): + if cls is not Mapper: + return [] + return [[self._mapper_offset, self._mapper_offset + self._mapper.params]] + + +def test__curvature_matrix_mapper_diag__matches_reference_kernel(): + pix_pixels = 9 + mapper_offset = 4 + total_params = mapper_offset + pix_pixels + + rng = np.random.default_rng(507) + + noise_map = 0.5 + rng.random((6, 6)) * 2.0 + native_index_for_slim_index = np.array( + [[y, x] for y in range(6) for x in range(6)] + ) + data_pixels = native_index_for_slim_index.shape[0] + + ( + psf_precision_operator, + indexes, + lengths, + ) = inversion_imaging_numba_util.psf_precision_operator_sparse_from( + noise_map_native=noise_map, + kernel_native=ASYMMETRIC_KERNEL, + native_index_for_slim_index=native_index_for_slim_index, + ) + + max_lengths = 3 + unique_mappings = UniqueMappings( + data_to_pix_unique=rng.integers( + 0, pix_pixels, size=(data_pixels, max_lengths) + ).astype("int"), + data_weights=rng.random(size=(data_pixels, max_lengths)), + pix_lengths=rng.integers(1, max_lengths + 1, size=data_pixels).astype("int"), + ) + + mapper = FakeMapper(params=pix_pixels, unique_mappings=unique_mappings) + + inversion = StubInversionMapperDiag( + mapper=mapper, + sparse_operator=StubSparseOperator( + psf_precision_operator_sparse=psf_precision_operator, + indexes=indexes.astype("int"), + lengths=lengths.astype("int"), + ), + total_params=total_params, + mapper_offset=mapper_offset, + ) + + curvature_matrix = inversion._curvature_matrix_mapper_diag + + diag_reference = inversion_imaging_numba_util.curvature_matrix_via_sparse_operator_reference_from( + psf_precision_operator=psf_precision_operator, + psf_precision_indexes=indexes.astype("int"), + psf_precision_lengths=lengths.astype("int"), + data_to_pix_unique=unique_mappings.data_to_pix_unique, + data_weights=unique_mappings.data_weights, + pix_lengths=unique_mappings.pix_lengths, + pix_pixels=pix_pixels, + ) + + assert curvature_matrix.shape == (total_params, total_params) + + assert curvature_matrix[ + mapper_offset:, mapper_offset: + ] == pytest.approx(diag_reference, rel=1.0e-6) + + # The mapper block is this property's only output: everything outside the + # mapper's param range must be untouched. + assert curvature_matrix[:mapper_offset, :] == pytest.approx( + np.zeros((mapper_offset, total_params)), abs=1.0e-12 + ) + assert curvature_matrix[:, :mapper_offset] == pytest.approx( + np.zeros((total_params, mapper_offset)), abs=1.0e-12 + ) diff --git a/test_autoarray/operators/over_sample/test_over_sampler.py b/test_autoarray/operators/over_sample/test_over_sampler.py index f7c9f4980..1fe137d16 100644 --- a/test_autoarray/operators/over_sample/test_over_sampler.py +++ b/test_autoarray/operators/over_sample/test_over_sampler.py @@ -107,3 +107,94 @@ def test__slim_index_for_sub_slim_index(): ) assert (over_sampling.slim_for_sub_slim == slim_index_for_sub_slim_index_util).all() + + +def test__over_sampler_binned_array_2d__non_uniform_counts_are_cached_and_not_mutated(): + """ + The divisor used to bin a non-uniform over sampled array depends only on the `segment_ids`, so it is cached on + the `OverSampler`. Repeated calls must therefore be bit-identical and must never mutate the cached array. + """ + mask = aa.Mask2D( + mask=[[False, False], [True, True]], + pixel_scales=1.0, + ) + + over_sampler = aa.OverSampler( + mask=mask, sub_size=aa.Array2D(values=[1, 2], mask=mask) + ) + + assert not over_sampler.sub_is_uniform + + arr = np.array([1.0, 5.0, 7.0, 10.0, 10.0]) + + binned_0 = over_sampler.binned_array_2d_from(array=arr) + + counts = over_sampler.binned_counts + + assert counts == pytest.approx(np.array([1, 4]), 1.0e-4) + assert not counts.flags.writeable + + binned_1 = over_sampler.binned_array_2d_from(array=arr) + binned_2 = over_sampler.binned_array_2d_from(array=2.0 * arr) + + assert np.array_equal(np.array(binned_0.slim), np.array(binned_1.slim)) + assert np.array_equal(np.array(binned_2.slim), 2.0 * np.array(binned_0.slim)) + + # The cached array is handed to every caller, so it must be unchanged after the calls above. + + assert over_sampler.binned_counts is counts + assert np.array_equal(counts, np.array([1, 4])) + + +def test__over_sampler_binned_array_2d__zero_count_segments_do_not_divide_by_zero(): + """ + A `sub_size` of zero produces a segment with no sub-pixels, whose count is guarded to 1 so the binned value is + zero rather than a divide-by-zero NaN. + """ + mask = aa.Mask2D( + mask=[[False, False], [True, True]], + pixel_scales=1.0, + ) + + with np.errstate(divide="ignore"): + over_sampler = aa.OverSampler( + mask=mask, sub_size=aa.Array2D(values=[0, 2], mask=mask) + ) + + arr = np.array([1.0, 2.0, 3.0, 4.0]) + + binned = over_sampler.binned_array_2d_from(array=arr) + + assert over_sampler.binned_counts == pytest.approx(np.array([1, 4]), 1.0e-4) + assert binned.slim == pytest.approx(np.array([0.0, 2.5]), 1.0e-4) + + +def test__over_sampler__pickles_with_cached_properties(): + """ + `OverSampler` instances are pickled to Nautilus worker processes, so the cached divisor must survive a + round trip. + """ + import pickle + + mask = aa.Mask2D( + mask=[[False, False], [True, True]], + pixel_scales=1.0, + ) + + over_sampler = aa.OverSampler( + mask=mask, sub_size=aa.Array2D(values=[1, 2], mask=mask) + ) + + arr = np.array([1.0, 5.0, 7.0, 10.0, 10.0]) + + binned = over_sampler.binned_array_2d_from(array=arr) + + over_sampler_pickled = pickle.loads(pickle.dumps(over_sampler)) + + assert np.array_equal( + over_sampler_pickled.binned_counts, over_sampler.binned_counts + ) + assert np.array_equal( + np.array(over_sampler_pickled.binned_array_2d_from(array=arr).slim), + np.array(binned.slim), + )