From a583f1a6ff7915b3fe03556f5d298378c47a9910 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 28 Aug 2026 17:40:06 -0400 Subject: [PATCH 1/3] perf(numba): oracle the mapper x mapper kernel, then hoist its row gathers (#507) Step 1 of PyAutoArray#507. `curvature_matrix_via_sparse_operator_from` is the single most expensive step of the numba CPU imaging likelihood at HST resolution (0.255 s of a 0.633 s evaluation, measured in step 0) and had **no direct unit test at all** -- only end-to-end inversion fixtures, which cannot tell a restructured kernel from a subtly wrong one. So the tests come first, written and passing against the unmodified kernel, and the restructuring second. Tests (all parametrised over the existing asymmetric, non-square `KERNELS_ODD`): - `test__curvature_matrix_via_sparse_operator_from__matches_dense_psf_precision_operator` pins the kernel to the definition it optimises, `F = M.T @ W @ M`, with `W` the dense [image_pixels, image_pixels] operator from `psf_precision_operator_from` and `M` the dense mapping matrix the unique-mappings triple encodes. The reference knows nothing of the sparse upper-triangle storage, the halved diagonal or the `A + A.T` fold. - `..._is_symmetric_and_the_diagonal_is_not_double_counted` asserts exact symmetry (`np.array_equal`, since the inversion runs no global symmetrizing pass) and isolates the halved-diagonal contract with a single data pixel mapping to a single source pixel, where `F[0, 0]` must be `w**2 * W[0, 0]` and would be twice that if either half of the contract were dropped. Control runs, both on the unmodified kernel, both recorded before the change: A. accumulate `1.0000001 * w0 * w1 * W` instead of `w0 * w1 * W` -> 5 failed, 19 deselected B. fold `range(i + 1, pix_pixels)` instead of `range(i, pix_pixels)` (i.e. drop the diagonal doubling) -> 5 failed, 19 deselected restored -> 5 passed The change itself. The innermost loop re-read `data_weights[data_1, pix_1_index]` and `data_to_pix_unique[data_1, pix_1_index]` on every accumulation, so each 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 every time. Those are now 1-D row views taken once per data pixel and once per stored pair, plus a `curvature_matrix[pix_0]` row view. The accumulated expression is left operand-for-operand identical -- floating-point multiplication is not associative, so folding `psf_precision_value` into the outer weight would have moved the likelihood. The pre-restructuring quadruple loop is retained verbatim as `curvature_matrix_via_sparse_operator_reference_from` (not called by the inversion) and `..._matches_the_reference_kernel_bit_identically` asserts `np.array_equal` between the two -- a tolerance would wave through exactly the reassociation this is guarding against. It is also the fixed reference the step-2 two-stage reformulation will be tested against. Measured, paired in one session on real datasets (B/A/B/A, n_repeats 10, OMP_NUM_THREADS=1, worktree PyAutoArray on PYTHONPATH), reference kernel vs production kernel on identical inputs: hst rectangular (784) 0.2368 s -> 0.2124 s 1.115x euclid rectangular (784) 0.0489 s -> 0.0460 s 1.062x hst delaunay (1250) 0.1051 s -> 0.1015 s 1.036x and in every one of the three cells both the kernel output *and* the assembled `inversion.curvature_matrix` (rebuilt with the reference kernel monkeypatched in) are bit-identical, `np.array_equal` True. This answers one of the plan's marked-undetermined questions: LLVM had already hoisted most, but not all, of the scalar work. 11 % at the HST rectangular fiducial is the low end of the plan's 10-20 % estimate, and the two smaller-u0 / smaller-P cells gain less, as expected. `pytest test_autoarray`: 1305 passed (1296 baseline + 9 new). Refs PyAutoLabs/PyAutoArray#507 Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01Fr6iJ5T1RDARWfxttCuGkK --- .../inversion_imaging_numba_util.py | 81 ++++++ .../imaging/test_inversion_imaging_util.py | 275 ++++++++++++++++++ 2 files changed, 356 insertions(+) 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..67f78e94b 100644 --- a/autoarray/inversion/inversion/imaging_numba/inversion_imaging_numba_util.py +++ b/autoarray/inversion/inversion/imaging_numba/inversion_imaging_numba_util.py @@ -554,6 +554,87 @@ 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_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/test_autoarray/inversion/inversion/imaging/test_inversion_imaging_util.py b/test_autoarray/inversion/inversion/imaging/test_inversion_imaging_util.py index 5f9adf5fa..27ed7e30e 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,278 @@ 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. + + 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_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) From 88e14bc637ab27123d0e6f2fa67eef4441927f41 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 28 Aug 2026 17:51:21 -0400 Subject: [PATCH 2/3] perf(numba): two-stage source-space accumulator for the mapper x mapper block (#507) Step 2 of PyAutoArray#507, and the phase-2 lever. HST rectangular: the mapper x mapper block of `F` goes 0.202 s -> 0.070 s (2.90x) and the whole numba CPU likelihood evaluation 0.511 s -> 0.376 s (1.36x), paired in one session. The direct kernel's innermost work is, per stored (data_0, data_1) pair, `u0 * u1` irregular read-modify-writes scattered across a `pix_pixels ** 2` matrix -- 4.9 MB at the HST fiducial, far outside L2. Step 0 measured 1.77e8 of them at ~1.4 ns each. But `w0` does not depend on `data_1`, so the sum factorises: a[p] = sum over data_0's 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 costs `sum_pairs u1` scatters into an L1-resident `pix_pixels` vector (u1, not u0 * u1); stage 2 is `u0` contiguous, vectorisable AXPYs over whole rows of `F`. The `A + A.T` fold and the halved-diagonal contract are untouched. `curvature_matrix_via_sparse_operator_from` becomes the dispatcher; the step-1 hoisted loop is `..._direct_from` and the new one `..._two_stage_from`, both with the same signature and contract, and `..._reference_from` (step 1) stays as the fixed reference both are tested against. **The measured crossover.** Swept `pix_pixels` on the two production HST geometries -- real 15361-pixel sparse operator, real mapper mappings, only the source-space extent varied -- speed-up of two-stage over direct: pix_pixels 128 784 1250 2048 4096 6144 8192 rectangular u0=4.00 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 Two things this says, neither of them guessable from the op counts. First, the crossover *is* geometry-dependent: it is near 8192 source pixels for the narrow-mapping Delaunay geometry and beyond 8192 for the wide-mapping bilinear one. Second, and contrary to the plan's op-count reasoning, two-stage wins on the Delaunay cell as well (1.53x at its native 1250) even though it executes 2.05x *more* arithmetic there (6.72e7 vs 3.27e7 ops, step 0): a vectorised contiguous AXPY is several times cheaper per op than a scattered RMW into a 12.5 MB matrix, and that dominates the op count. The plan's "valid while n_source <~ P * u1" criterion is too conservative. So there is no crossover anywhere inside the range a PyAuto pixelization is run at. `CURVATURE_TWO_STAGE_MAX_PIX_PIXELS = 4096` is therefore deliberately the conservative end of the measured envelope rather than a fitted crossover -- both geometries still favour two-stage there, and above it the two forms are within a few per cent anyway. It is an explicit, commented module-level constant carrying the sweep above, not a silent heuristic, and it is overridable per call via `two_stage_max_pix_pixels` -- which is how both dispatcher branches are tested without allocating a 134 MB matrix. `prange` is not implemented, per the plan. Correctness. The two forms sum the same products in a different order, so they agree to floating-point reassociation, not bit-identically. Paired on real data, before (direct) vs after (dispatcher), same session: max rel diff on F log likelihood before / after hst rectangular 1.24e-14 27661.910133664110 / 27661.910133664103 euclid rectangular 2.63e-15 6213.306873885871 / 6213.306873885877 hst delaunay 3.90e-13 29090.527210448134 / 29090.527210448126 `np.allclose(rtol=1e-12)` on the assembled `inversion.curvature_matrix` is True in all three, and every log likelihood moves by <1e-15 relative -- nine orders inside the `rtol=1e-6` pins. Tests (18 new, 1305 -> 1323): - `..._two_stage_from__matches_reference_kernel`, over `KERNELS_ODD` x `pix_pixels` in {1, 7, 64} (degenerate source space; smaller than the per-data-pixel mapping footprint; larger than it), at rel=1e-6, plus exact symmetry of the folded result. - `..._two_stage_from__matches_dense_psf_precision_operator`, anchoring it to the `M.T @ W @ M` definition and not only to the other implementation. - `..._from__dispatches_on_the_two_stage_threshold`, asserting each branch *bit-identical* to the implementation it should have called (an rtol comparison could not tell the branches apart), and that the two branches do differ, so neither assertion is trivially satisfied. - `test__curvature_matrix_mapper_diag__matches_reference_kernel` in `test_curvature_matrix_func_list_blocks.py`, one level up: what the inversion assembles into `F`, with the mapper's param range placed correctly and everything outside it untouched. The multi-mapper twin `curvature_matrix_off_diags_via_sparse_operator_from` is not touched and so gets no oracle, per the plan. Control runs, each a deliberately wrong variant, on `-k "two_stage or dispatches or mapper_diag__matches_reference"` (18 selected): C. drop the per-data-pixel accumulator clear -> 17 failed, 1 passed D. scale the stage-2 AXPY by 1.00001 (1e-5 rel) -> 17 failed, 1 passed E. dispatcher never takes the two-stage branch -> 1 failed, 17 passed restored -> 18 passed C and D leave exactly one test passing, and it is the intended one: the dispatcher test asserts the dispatcher is bit-identical to the branch it selects, which stays true when that branch is equally broken. E fails exactly the dispatcher test and nothing else, since the direct branch is correct. `pytest test_autoarray`: 1323 passed. Refs PyAutoLabs/PyAutoArray#507 Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01Fr6iJ5T1RDARWfxttCuGkK --- .../inversion_imaging_numba_util.py | 243 ++++++++++++++++++ .../imaging/test_inversion_imaging_util.py | 220 +++++++++++++++- .../test_curvature_matrix_func_list_blocks.py | 122 +++++++++ 3 files changed, 581 insertions(+), 4 deletions(-) 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 67f78e94b..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` @@ -603,6 +730,122 @@ def curvature_matrix_via_sparse_operator_from( 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, 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 27ed7e30e..cdb247ab3 100644 --- a/test_autoarray/inversion/inversion/imaging/test_inversion_imaging_util.py +++ b/test_autoarray/inversion/inversion/imaging/test_inversion_imaging_util.py @@ -756,6 +756,10 @@ def test__curvature_matrix_via_sparse_operator_from__matches_the_reference_kerne `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 @@ -794,10 +798,8 @@ def test__curvature_matrix_via_sparse_operator_from__matches_the_reference_kerne pix_pixels=pix_pixels, ) - curvature_matrix = ( - aa.util.inversion_imaging_numba.curvature_matrix_via_sparse_operator_from( - **kwargs - ) + 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( @@ -805,3 +807,213 @@ def test__curvature_matrix_via_sparse_operator_from__matches_the_reference_kerne ) 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 + ) From d8bc3bacadf6dc24a5a8454becc3e77ec307d7ef Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 28 Aug 2026 18:09:51 -0400 Subject: [PATCH 3/3] perf(over-sample): cache the non-uniform binning divisor on the OverSampler (#507) `OverSampler.binned_array_2d_from` is called once per numpy light-profile evaluation -- 120 times per likelihood evaluation for a 60-component MGE (the data grid and the blurring grid), and once more for every ordinary light profile. On the non-uniform branch it recomputed, every call, two quantities that depend only on the `sub_size` fixed in the constructor: - `np.bincount(self.segment_ids, ...)`, the number of sub-pixels binned into each pixel, plus the zero-guard `counts[counts == 0] = 1`; - `sub_is_uniform`, an `np.isclose` over the whole `sub_size` array, which selects the branch in the first place. Both are now `cached_property`. Measured on the HST breakdown cell's over sampler (17980 sub-pixels into 15361 pixels): the divisor cost 38.8 us per call and the uniformity check 61.3 us, against 81.5 us for the whole cached call -- so a binned evaluation was more than twice the work it needed to be. The cached divisor is marked read-only (`flags.writeable = False`) and built with `np.maximum(counts, 1)` rather than the old in-place `counts[counts == 0] = 1`, because it is now handed to every caller instead of being a fresh array per call. Values are unchanged: `np.maximum` on the `bincount` output is the same int64 array the in-place guard produced, so the binned arrays are bit-identical. `cached_property` stores on the instance `__dict__`, so the cache pickles to Nautilus worker processes with the rest of the object -- asserted by a round-trip test, since `OverSampler` is pickled once per fit. Tests (`test_autoarray/operators/over_sample/test_over_sampler.py`): `..._non_uniform_counts_are_cached_and_not_mutated` (repeated calls bit-identical, cached array unchanged and not writeable), `..._zero_count_segments_do_not_divide_by_zero` (the guard the in-place mutation used to provide), and `..._pickles_with_cached_properties`. `pytest test_autoarray`: 1326 passed (1323 before, +3). Refs PyAutoLabs/PyAutoArray#507 Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01Fr6iJ5T1RDARWfxttCuGkK --- .../operators/over_sampling/over_sampler.py | 39 ++++++-- .../over_sample/test_over_sampler.py | 91 +++++++++++++++++++ 2 files changed, 122 insertions(+), 8 deletions(-) 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/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), + )