perf(numba): the mapper x mapper block and the MGE operated matrix (#507) by Jammy2211 · Pull Request #508 · PyAutoLabs/PyAutoArray

Jammy2211 · 2026-08-28T22:39:32Z

Summary

Phase 2 of #507: the two terms step 0 measured as carrying the HST rectangular
numba CPU likelihood — the mapper x mapper block of F (46 %) and the MGE
operated mapping matrix (37 %). Three commits here; the MGE half's other lever
is PyAutoLabs/PyAutoGalaxy#590, which must merge after this PR.

Result, paired B/A/B/A in one session on real datasets (three arms, three
rounds per cell in rotating arm order, first round discarded for the numba
recompile, mean of the last two; OMP_NUM_THREADS=1,
AUTOARRAY_NUMBA_OPERATED_MEMO=0, n_repeats=10):

cell	base (`1b89404b`)	after steps 1-2	final	whole phase
hst rectangular, step total	0.6214 s	0.4324 s	0.3334 s	1.86x
hst rectangular, direct eval	0.6184 s	0.4207 s	0.3013 s	2.05x
euclid rectangular	0.2226 s	0.1870 s	0.1417 s	1.57x
hst Delaunay-1250	0.6331 s	0.6076 s	0.4884 s	1.30x
hst: F mapper x mapper	0.2581 s	0.0819 s	0.0825 s	3.13x
hst: MGE total	0.1970 s	0.1937 s	0.0844 s	2.33x

The issue's goal — 0.60 -> ~0.35 s at HST rectangular — is met. Rows the
change does not touch hold still (Blurred image 0.98-1.02x, the batched MGE
PSF convolution 0.96-1.00x), which is the control on the paired session.

Three commits:

a583f1a6 — oracle the mapper x mapper kernel, then hoist its row gathers.
The kernel had no direct unit test; it now has one pinned to F = M.T W M
with a dense W, plus a symmetry / halved-diagonal test, with the pre-change
quadruple loop retained as
curvature_matrix_via_sparse_operator_reference_from. The hoist takes 1-D row
views once per data pixel instead of re-gathering from wide-stride 2-D arrays
u0 times, leaving the accumulated expression operand-for-operand identical.
Bit-identical; 1.04-1.12x.
88e14bc6 — the two-stage source-space accumulator. w0 does not depend
on data_1, so the sum factorises into a per-data-pixel dense accumulator
over source space followed by contiguous AXPYs over whole rows of F. This
replaces ~1.8e8 irregular read-modify-writes into a 4.9 MB matrix with ~4.4e7
L1 scatters plus vectorisable dense adds. 2.90x on the block at HST.
CURVATURE_TWO_STAGE_MAX_PIX_PIXELS = 4096 is an explicit module-level
constant carrying the measured pix_pixels sweep (crossover at or beyond 8192
on both production geometries — outside the range a PyAuto pixelization runs
at), overridable per call, not a silent heuristic.
d8bc3bac — cache the OverSampler's non-uniform binning divisor.
binned_array_2d_from recomputed np.bincount(segment_ids) and the
sub_is_uniform check on every call — 120 calls per evaluation for a
60-Gaussian MGE, and once more for every ordinary numpy light profile. Both
depend only on the constructor's sub_size and are now cached_property. On
the HST over sampler (17980 sub-pixels into 15361 pixels) the divisor cost
38.8 us and the uniformity check 61.3 us per call, against 81.5 us for the
whole cached call. Bit-identical; the cached divisor is read-only, because the
old code guarded zero counts by mutating in place.

API Changes

None — internal changes only. No public signature or default changes. Two
behaviour-preserving additions to OverSampler: sub_is_uniform becomes a
cached_property (same value, computed once per instance) and a new
binned_counts cached property. The new numba kernels
(curvature_matrix_via_sparse_operator_direct_from / ..._two_stage_from /
..._reference_from) live in inversion_imaging_numba_util.py beside the
existing dispatcher, which keeps its name and signature.

See full details below.

Test Plan

pytest test_autoarray — 1326 passed (1296 at the phase-1 baseline;
+12 step 1, +18 step 2, +3 step 3).
Oracle tests written and passing against the unmodified kernel before
any change, with recorded control runs (a deliberately wrong constant and a
dropped diagonal doubling each fail them).
Control runs recorded for every rtol test in step 2 (dropping the
accumulator clear: 17 failed 1 passed; a wrong dispatcher branch: 1 failed).
Bit-identical checks: step-1 kernel and assembled curvature_matrix via
np.array_equal; step-3 binned_array_2d_from across repeated calls, with
the cached array asserted unchanged and non-writeable, and a pickle
round-trip (OverSampler is pickled to Nautilus workers).
Pinned log likelihoods at explicit rtol=1e-6 on all three pinned
profiling cells, on every arm, unchanged to every recorded digit: hst
bilinear 27661.91013366411, hst RTU 27180.70471569685, hst Delaunay
29090.527210448134. euclid (unpinned) 6213.306873885871.
inversion.curvature_matrix, same instance on each arm: bit-identical for
step 3 alone; max rel 1.24e-14 (hst) / 2.63e-15 (euclid) across the
whole phase, np.allclose(rtol=1e-12) True. The MGE operated mapping
matrix is bit-identical base -> final.
Pool run (8 cores, PYAUTO_TEST_MODE=1, PYAUTO_SMALL_DATASETS unset,
2828 masked pixels, 60-Gaussian linear MGE bulge): per evaluation
0.2000 -> 0.1697 s in a pool of 8 and 0.4099 -> 0.3634 s serial, so the
parallel speed-up ratio goes 2.05x -> 2.14x. Flat-to-up is the pass
condition — it is what would fall if a lever had introduced hidden threads.

Merge order

Merge this PR before PyAutoLabs/PyAutoGalaxy#590. PyAutoGalaxy's CI clones
PyAutoArray from source and checks out the same-named branch when it exists, so
#590 is currently testing against this branch; once this merges and the branch is
deleted it falls back to PyAutoArray main, which is why this must land first.
No PyAutoGalaxy test depends on the new OverSampler behaviour, so #590 is green
either way. The workspace PR PyAutoLabs/autolens_profiling#190 (the note and the
re-run artifacts) merges last, behind both.

Heart

Heart is RED at ship time for a pre-existing reason unrelated to this task.
Verbatim from pyauto-heart readiness --json (2026-08-28T21:34:42Z, score 45):

red_reasons: "release validation FAILED (stage integrate)"
yellow_reasons: "workspace validation not passing (2 failed, cloud#33179766004: autolens_test scripts/imaging/rectangular_mge.py, autolens_test scripts/imaging/rectangular_mge_rtu.py)",
"manifest drift: session-start hooks (generated) — 32 mismatch(es) vs PyAutoMind/repos.yaml"

Human authorisation to ship over it, given 2026-08-28, verbatim:

prm and then kick off phase 2, I authorize the heart RED thing

Full API Changes (for automation & release notes)

Added

autoarray.operators.over_sampling.over_sampler.OverSampler.binned_counts — the cached, read-only per-pixel sub-pixel count used as the divisor by binned_array_2d_from on the non-uniform branch.
autoarray.inversion.inversion.imaging_numba.inversion_imaging_numba_util.curvature_matrix_via_sparse_operator_direct_from — the step-1 hoisted direct loop.
...inversion_imaging_numba_util.curvature_matrix_via_sparse_operator_two_stage_from — the step-2 two-stage form.
...inversion_imaging_numba_util.curvature_matrix_via_sparse_operator_reference_from — the pre-change quadruple loop, retained as the test reference; not called by the inversion.
...inversion_imaging_numba_util.CURVATURE_TWO_STAGE_MAX_PIX_PIXELS — the dispatcher threshold constant.

Changed Signature

...curvature_matrix_via_sparse_operator_from(..., two_stage_max_pix_pixels=CURVATURE_TWO_STAGE_MAX_PIX_PIXELS) — one new keyword argument with a default; every existing call is unaffected.

Changed Behaviour

OverSampler.sub_is_uniform — property -> cached_property. Same value; sub_size is fixed in the constructor.
OverSampler.binned_array_2d_from — the non-uniform branch reuses the cached divisor instead of recomputing np.bincount and guarding zeros in place. Values bit-identical.
curvature_matrix_via_sparse_operator_from — dispatches to the two-stage form below two_stage_max_pix_pixels source pixels. Agrees with the direct form to floating-point reassociation (max rel 1.24e-14 on the assembled F), not bit-identically.

Migration

None required.

Generated by the PyAutoLabs agent workflow.

…thers (#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 #507 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Fr6iJ5T1RDARWfxttCuGkK

…er 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 #507 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Fr6iJ5T1RDARWfxttCuGkK

…ampler (#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 #507 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Fr6iJ5T1RDARWfxttCuGkK

Jammy2211 · 2026-08-28T22:42:41Z

Downstream: PyAutoLabs/PyAutoGalaxy#590 (the MGE half) and PyAutoLabs/autolens_profiling#190 (harness + note + artifacts). Merge order: this PR -> #590 -> #190.

Jammy2211 and others added 3 commits August 28, 2026 17:40

Jammy2211 added the pending-release PR queued for the next release build label Aug 28, 2026

This was referenced Aug 28, 2026

perf(linear): share one eccentric-radius grid across an MGE basis (PyAutoArray#507) PyAutoLabs/PyAutoGalaxy#590

Merged

profiling: phase-2 instrumentation and the measured result (#507) PyAutoLabs/autolens_profiling#190

Merged

Jammy2211 mentioned this pull request Aug 28, 2026

perf(numba): HST curvature matrix phase 2 — mapper×mapper block + MGE operated matrix #507

Closed

Jammy2211 merged commit f4b8e75 into main Aug 28, 2026
3 checks passed

Jammy2211 deleted the feature/numba-hst-curvature-matrix-phase2 branch August 28, 2026 22:51

Jammy2211 mentioned this pull request Aug 31, 2026

research: numba positive-only solve vs JAX sparse-operator mode — same linear algebra? #513

Closed

Jammy2211 removed the pending-release PR queued for the next release build label Sep 4, 2026

Navigation Menu

Repository navigation

perf(numba): the mapper x mapper block and the MGE operated matrix (#507) - #508

perf(numba): the mapper x mapper block and the MGE operated matrix (#507)#508
Jammy2211 merged 3 commits into
mainfrom
feature/numba-hst-curvature-matrix-phase2

Jammy2211 commented Aug 28, 2026

Uh oh!

Jammy2211 commented Aug 28, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

Jammy2211 commented Aug 28, 2026

Summary

API Changes

Test Plan

Merge order

Heart

Added

Changed Signature

Changed Behaviour

Migration

Uh oh!

Jammy2211 commented Aug 28, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant