perf: numba CPU curvature matrix F at HST resolution (phase 1: split + FFT mapper×func block)

[Jammy2211]

Overview

Successor to the numba-cpu-likelihood epic (COMPLETE 2026-08-28; 21.3 s → 1.25 s per HST evaluation). The one step left on the numba CPU (use_jax=False, apply_sparse_operator_cpu()) path is the curvature matrix F: 0.997 s = 78 % of a 1.28 s HST rectangular evaluation and 55 % on Delaunay-1250. Target: ≥ 2× on the F step (HST rectangular evaluation ≤ ~0.7 s) with the log-likelihood unchanged to pinned tolerance (hst rectangular pin 27661.910133664103).

Code exploration found the prompt's candidate list partly stale: the mapper block already exploits symmetry (half-stored PSF pairs, halved diagonal, fold + mirror in curvature_matrix_via_sparse_operator_from); the mapper × linear-func block is dense — curvature_matrix_off_diags_via_mapper_and_linear_func_curvature_vector_from expands the 60 MGE curvature weights onto the full native grid and runs a direct sliding-window correlation (convolve_with_kernel_native) over ny × nx × ky × kx × 60 every evaluation; and there are redundant per-evaluation passes (a global re-mirror of an already-symmetric matrix with a fresh O(P²) allocation, np.array copies of the unique-mapping arrays, per-pair / noise_map**2 of memoized operated matrices).

This is phase 1 of a split (Brain FeatureDecision: split-into-phases). Threading is not a lever: Nautilus samples with Python multiprocessing, one process per core, so prange would oversubscribe (N procs × N threads). Phase 1 uses single-thread algorithmic wins only; any threading lever would need to be measured under the real pool as well as OMP_NUM_THREADS=1 before it is even considered.

Plan

• Instrument F: expose its sub-blocks (mapper diag, mapper × linear-func, func × func, mirror) as separately timeable steps in the autolens_profiling breakdown harness; record the split at hst + euclid, rectangular (Bilinear + RTU) + Delaunay-1250, with OMP_NUM_THREADS=1 AUTOARRAY_NUMBA_OPERATED_MEMO=0. Checkpoint: the split decides whether the FFT step runs.
• Commit the 2026-08-28 re-baseline artifacts (regenerated, not hand-copied) — the committed *_v2026.8.17.1 artifacts pre-date PyAutoArray#462 and still show an 18.8 s "sparse triplets" step.
• Remove the redundant per-evaluation passes (global mirror, array copies, repeated noise-weighting). Bit-identical output expected.
• If the split confirms the dense mapper × linear-func convolution dominates: route it through the existing batched FFT Convolver machinery (phase-1 lever, PyAutoArray#497) with the flipped kernel (the current kernel computes a correlation), keeping a numba kernel for the scatter/accumulate only. Parity pinned at rtol 1e-6 with a brute-force test on an asymmetric PSF.
• test_autoarray green; smoke cpu_fast_modeling.py; measure after on the same four cells plus one Nautilus pool run (catches any implicit FFT/BLAS thread pool oversubscribing); record before/after in autolens_profiling results + a note; post the numbers here.
• Phase 2 is filed at close-out only if phase 1 misses 2× (mapper-independent hoisting the split reveals).

Detailed implementation plan

Work Classification

Both — library (PyAutoArray) first, workspace (autolens_profiling) follow-up once the private helpers exist.

Affected Repositories

• PyAutoArray (primary)
• autolens_profiling

Branch Survey

Repository	Current Branch	Dirty?	Claim
./PyAutoArray	main	clean	none
./autolens_profiling	main	clean	claimed by nuts-warm-start-driver-and-a100-probe — disjoint file set (scripts/misc/searches/, scripts/imaging/searches/nuts/, results/notes/inference/), so a parallel worktree is used

Suggested branch: feature/numba-hst-curvature-matrix-speedup
Worktree root: ~/Code/PyAutoLabs-wt/numba-hst-curvature-matrix-speedup/ (created by /start_library)

Implementation Steps

Step 0 — instrument

1. PyAutoArray/autoarray/inversion/inversion/imaging_numba/sparse.py: split _curvature_matrix_func_list_and_mapper into private helpers that each return their block — _curvature_matrix_mapper_diag (exists), _curvature_matrix_mapper_func_blocks_from(curvature_matrix), _curvature_matrix_func_func_blocks_from(curvature_matrix). No behaviour change.
2. autolens_profiling/scripts/imaging/likelihood_breakdown/pixelization_numba.py + delaunay_numba.py: add STEP_ACCESSORS entries for each sub-block (mapper diag → mapper×func → func×func → mirror) ahead of the existing Curvature matrix F entry (incremental-cost scheme attributes the time); keep the F entry as the total.
3. Run the four cells (--instrument hst|euclid, --rect-mesh bilinear|rtu; _profile_cli.py:128-140) with OMP_NUM_THREADS=1 AUTOARRAY_NUMBA_OPERATED_MEMO=0; commit JSON+PNG under results/breakdown/imaging/; post the split table here. Checkpoint.

Step 1 — redundant passes (bit-identical)
4. sparse.py:347: drop the global curvature_matrix_mirrored_from — the x1-mapper path is already folded+mirrored (util.py:596-602); in the func-list path write the mapper×func block's transpose into [func, mapper] at sparse.py:540; multi-mapper path keeps off_diag_0 + off_diag_1.T and places both blocks. Grep other callers of curvature_matrix_mirrored_from before removing anything.
5. sparse.py:389-393: pass mapper.unique_mappings.* arrays directly (verify dtypes; drop the np.array copies unless a cast is needed).
6. sparse.py:527-530 / 552: form operated / noise_map**2 once per linear func, hoisted out of the mapper loop.
7. Fix the curvature_matrix docstring (sparse.py:336-338, says "not a cached property" under @cached_property).

Step 2 — FFT for the mapper × linear-func block (gated by step 0)
8. In sparse.py, compute blurred_slim for curvature_weights ([n_unmasked, n_funcs]) via the existing FFT Convolver path (autoarray/operators/convolver.py; reuse the dataset's preloaded ConvolverState, Convolver.state_from(mask) / is_for_mask) with the flipped PSF kernel — the current kernel computes a correlation (util.py:806-812). Zero outside the mask. scipy.fft workers=1.
9. Split the numba kernel: keep only the scatter/accumulate loop (util.py:770-777) as curvature_matrix_off_diags_via_mapper_and_blurred_curvature_weights_from(...); keep the old dense kernel + convolve_with_kernel_native as the reference for tests.
10. Parity test in test_autoarray/inversion/inversion/imaging/test_inversion_imaging_util.py: FFT-path block vs dense kernel, small masked grid, asymmetric PSF (catches the flip), rtol=1e-6; extend test_curvature_matrix_func_list_blocks.py::test__curvature_matrix_func_list_blocks__matches_brute_force if it lacks an asymmetric kernel.

Step 3 — measure + ship
11. Re-run the four cells + one Nautilus pool run (cpu_fast_modeling.py, number_of_cores = machine cores; per-iteration wall-clock must improve by ~the single-thread factor); commit artifacts; add autolens_profiling/results/notes/numba_curvature_matrix_f_split.md with the before/after table.
12. ship_library (PyAutoArray) first; ship_workspace (autolens_profiling) after the library PR's API summary.

Key Files

• PyAutoArray/autoarray/inversion/inversion/imaging_numba/sparse.py — F assembly (curvature_matrix @319, mapper diag @362, mapper×func @520-545, func×func @550-577, global mirror @347)
• PyAutoArray/autoarray/inversion/inversion/imaging_numba/inversion_imaging_numba_util.py — kernels: sparse-op diag @515-604, off-diag @607-694, dense mapper×func @697-778, convolve_with_kernel_native @781-815, mirror @495-512
• PyAutoArray/autoarray/operators/convolver.py — batched FFT Convolver / ConvolverState (existing lever)
• PyAutoArray/test_autoarray/inversion/inversion/imaging/test_inversion_imaging_util.py, test_autoarray/inversion/inversion/test_curvature_matrix_func_list_blocks.py — parity tests
• autolens_profiling/scripts/imaging/likelihood_breakdown/{pixelization_numba,delaunay_numba}.py — STEP_ACCESSORS harness; autolens_profiling/_profile_cli.py — instrument / --rect-mesh selection; results/breakdown/imaging/ — artifacts

Verification

• test_autoarray green in the worktree; new parity tests pass.
• Harness pins unchanged (hst rectangular 27661.910133664103) on the unperturbed instance.
• autolens_workspace/scripts/imaging/features/pixelization/cpu_fast_modeling.py smoke profile.
• Before/after table hst/euclid × bilinear/rtu + Delaunay-1250 and the pool run posted here; F ≥ 2× or phase 2 filed with the measured residual.

Original Prompt

Click to expand starting prompt

Numba CPU likelihood at HST resolution: speed up the curvature matrix F (the 78% step)

Type: feature
Epic: none (successor to numba-cpu-likelihood, COMPLETE 2026-08-28)
Target: autoarray
Repos:

• @PyAutoArray
• @autolens_profiling
  Themes:
• numba-cpu
• pixelization
• profiling
  Difficulty: large
  Autonomy: supervised
  Priority: high
  Status: formalised
  Filed: 2026-08-28

Successor to epic numba-cpu-likelihood (epics.md; records complete/2026/08/numba-cpu-*.md).
That epic took the numba CPU (use_jax=False, apply_sparse_operator_cpu()) likelihood from
21.3 s → 1.25 s per evaluation at HST resolution on the default RectangularBilinearAdaptDensity
mesh and 4.5 → ~1.6 s on the Delaunay-1250 fiducial. What is left is one step.

Where the time goes now (2026-08-28, pixelization_numba.py / delaunay_numba.py breakdowns, OMP_NUM_THREADS=1, AUTOARRAY_NUMBA_OPERATED_MEMO=0, PyAutoArray main 1f5c636)

step	hst rectangular (1.28 s)	euclid rectangular (0.36 s)	hst Delaunay-1250 (3.22 s)
Curvature matrix F [numba sparse-op]	0.997 s (78%)	0.211 s (59%)	1.77 s (55%)
MGE operated mapping matrix (60 funcs)	0.189 s	0.077 s	—
Reconstruction solve (warm-started NNLS)	0.008 s	0.008 s	0.54 s (warm: ~0.07 s)
everything else	≤ 0.03 s each	≤ 0.02 s each	inversion build ~0.5 s

F is the curvature_matrix of InversionImagingSparseNumba
(autoarray/inversion/inversion/imaging_numba/sparse.py, kernels in
inversion_imaging_numba_util.py): the mapper × PSF-precision-operator × mapper contraction over the
sparse-operator (w-tilde heritage) representation, plus the linear-func blocks (already noise-weighted
once and mirrored, phase 1). It is rebuilt every evaluation because the mapper changes with the mass
model — no cross-evaluation memo applies.

Goal

Make F substantially cheaper on the numba CPU path at HST resolution (target: ≥ 2× on the F step,
i.e. an HST rectangular evaluation ≤ ~0.7 s; Delaunay similar), with the log-likelihood unchanged to
pinned tolerance (delaunay_numba.py / pixelization_numba.py pins; hst rectangular pin
27661.910133664103).

1. Step 0 — decompose F itself. The breakdown cells time fit.inversion.curvature_matrix as one
   step. Instrument the kernel(s) behind it (mapper-mapper block, mapper-linear-func blocks,
   linear-func-linear-func block, any dense fill / symmetrisation / noise weighting) and record the
   split at hst + euclid, rectangular (default) + Delaunay. Commit the re-baseline artifacts from the
   2026-08-28 re-profile too (pixelization_numba_breakdown_{euclid,hst} incl. the _rtu variants —
   numbers in complete/2026/08/numba-cpu-kernel-cdf-fast-path.md; regenerate, do not hand-copy).
2. Pick the lever from the split, not from doctrine. Candidates to evaluate, cheapest first:
   • the unique-mappings compression (data_slim_to_pixelization_unique_from) — is F's inner loop
     iterating over data pixels × PSF footprint × source pixels where a per-source-pixel / per-unique
     -mapping formulation would be smaller;
   • symmetry — compute the upper triangle only and mirror (the linear-func blocks already do);
   • numba.prange over the outer loop with the OMP thread count (the campaign runs one process per
     core under Nautilus, so gains must be measured with OMP_NUM_THREADS=1 AND with the pool);
   • reuse across evaluations of whatever does NOT depend on the mapper (PSF-precision-operator
     products are already preloaded — verify nothing mapper-independent is recomputed);
   • the MGE operated-matrix term is second (0.19 s at hst) — the phase-1 batching left the
     per-evaluation convolution of the varying profiles; check whether it is FFT- or scatter-bound.
3. Ship behind exact parity (bit-level where the summation order is unchanged, pinned rtol 1e-6
   otherwise); test_autoarray green; smoke autolens_workspace/scripts/imaging/features/pixelization/ cpu_fast_modeling.py on the smoke profile.
4. Measure before/after on the four cells above; record results + a note in autolens_profiling.

Out of scope: the JAX path; the RTU / kernel-CDF meshes (GPU-only by decision, 2026-08-28); the NNLS
solve (done — nnls_warm_start_memo).

[Jammy2211]

Step 0 — F split (OMP_NUM_THREADS=1, memo off, n_repeats=10, v2026.8.17.1)

cell	total eval	mapper×mapper	mapper×linear-func (dense conv)	func×func	residual	F total	next non-F
hst-bilinear	1.74 s	0.316	0.953	0.022	−0.016	1.275	MGE operated matrix 0.260
euclid-bilinear	0.444 s	0.063	0.198	0.004	0.018	0.284	MGE operated matrix 0.095
hst-rtu	8.50 s	0.304	0.929	0.022	−0.005	1.250	sparse triplets 7.88 (RTU = GPU-only by decision; separate prompt)
hst-delaunay-1250	1.46 s	0.130	0.899	0.021	0.004	1.055	MGE operated matrix 0.235

The dense native-grid mapper × linear-func convolution is 70–85 % of F on every cell and scales with image pixels (3841 → 15361: 0.198 → 0.95 s), not with the source mesh. Checkpoint verdict: the FFT lever (step 2) runs. Pins passed on all three pinned cells. Delaunay-1250 now measures 1.46 s (the issue's opening table quoted 3.22 s, pre-#501). Commits: PyAutoArray dedde5ac, autolens_profiling 6f7d20a.

[Jammy2211]

Measured: F is 3.3x faster, the HST evaluation 2.6x

Two changes on feature/numba-hst-curvature-matrix-speedup:

1. Drop the redundant passes in F assembly — the three blocks each rebuilt
   operated_mapping_matrix / noise_map ** 2 and re-walked the mapping matrix.
   Split into per-block helpers over one shared weighted copy. Bit-identical
   (np.array_equal True on both euclid and hst).
2. FFT the mapper x linear-func block — its dense sliding window is a
   correlation with the PSF, i.e. a convolution with the PSF reversed along
   both axes. It now goes through the existing batched Convolver (new cached
   Convolver.reversed_kernel), once per linear func instead of once per
   (mapper, linear func) pair; only the sparse scatter onto source pixels stays
   in numba.

Before / after (s per evaluation, OMP_NUM_THREADS=1, memo off, n_repeats 10)

The before column is a fresh re-measurement of the pre-change commit taken
back-to-back with the after run — this host carries 20-30% session-to-session
variance, so the original baseline table is not a safe comparator.

cell	eval	F total	mapper x mapper	mapper x l-func
hst, rectangular bilinear	1.562 -> 0.595	1.195 -> 0.359	0.295 -> 0.277	0.858 -> 0.054
euclid, rectangular bilinear	0.349 -> 0.249	0.256 -> 0.097	0.060 -> 0.065	0.182 -> 0.023
hst, Delaunay-1250	1.367 -> 0.758	1.077 -> 0.184	0.122 -> 0.123	0.854 -> 0.052
hst, rectangular RTU	(8.501) -> 7.479	(1.250) -> 0.334	(0.304) -> 0.249	(0.929) -> 0.061

RTU has no same-session before column — it is an 8.5 s cell dominated by its
7.1 s "mapper sparse triplets" step and is GPU-only by the 2026-08-28 decision,
so it was re-run once for currency rather than paired; its numbers are shown in
parentheses.

~16x on the mapper x linear-func block, 3.3x on F, 2.6x on the whole HST
evaluation (5.9x on F for Delaunay, where F was 74% of the evaluation). The
issue's 2x goal on F is met on both HST cells.

Correctness

• F agrees with the sliding-window result to 3e-18 relative on euclid and
  hst (the FFT is not bit-identical to the direct sum, as expected).
• All 3 pinned log-likelihoods PASSED: hst bilinear 27661.91013366411
  (pin 27661.910133665442), hst RTU 27180.70471569685 (pin 27180.704715698186),
  hst Delaunay 29090.527210448134 (pin 29090.527192092646). euclid has no pin;
  it measured 6213.306873885871, unchanged to every recorded digit.
• pytest test_autoarray: 1296 passed (1290 before, 6 new). The new tests
  assert the FFT path against the retained dense kernel using asymmetric
  non-square PSFs, so a missing axis reversal fails.

Pool run — no thread oversubscription

A single-thread win is only real if it survives Nautilus's one-process-per-core
pool, and an FFT that quietly spins up its own threads would show as a pool
regression. Checked with autolens_workspace's
imaging/features/pixelization/cpu_fast_modeling.py (first, non-SLaM fit) on an
8-core host, against canonical main and against the worktree via PYTHONPATH
(autoarray.__file__ asserted both ways), number_of_cores raised 2 -> 8.

Two forced departures from the smoke profile: PYAUTO_TEST_MODE=1 rather than
the default 2 (level 2 bypasses the sampler, so the smoke profile as written
runs no pool at all; level 1 sets n_like_max=1 = one 100-evaluation
exploration batch), and PYAUTO_SMALL_DATASETS left unset (under the 15x15 cap
the fit sees 256 masked pixels and F is too small to time; unset it is 2828).

The script's own model is mass + shear with no lens light, so its
linear_obj_list holds the mapper alone and the FFT'd block does no work —
that config is flat, as it should be (median 0.0625 -> 0.0677 s/eval, noise).
Adding a 60-component linear MGE bulge (the multi_gaussian_expansion/modeling.py
recipe) makes the block non-empty:

	pool of 8 (median of 3)	serial (number_of_cores=1)
main	0.1919 s/eval	0.4845 s/eval
branch	0.1747 s/eval	0.4489 s/eval

The pool improves 9% and the single process 7% — the pool gain tracks the
single-thread gain rather than eroding it — and the parallel speed-up ratio is
flat across the change: 2.52x -> 2.57x on 8 cores. That ratio is what would
drop if hidden threads had appeared. Consistent with the code: Convolver's FFT
path uses np.fft.rfft2 (no workers=, single-threaded); scipy.fft appears
only as next_fast_len, a shape calculation.

Read that as a regression check, not a speed-up measurement — the workspace
dataset is a 2828-pixel 0.1"/pixel simulation with a small mesh, so F is a much
smaller share of the evaluation than in the HST harness cells.

What leads now

F is no longer the dominant term at HST resolution. On hst bilinear the largest
remaining steps are the mapper x mapper sparse-operator block (0.277 s) and
the MGE operated mapping matrix (~0.224 s); the mapper x linear-func block
has fallen from 72% of F to 15%. The mapper x mapper block is the natural
phase-2 candidate. Since the 2x goal on F is met, no phase-2 prompt is filed by
this task — flagging it here so the next performance pass starts from the right
term.

Artifacts and the full method are in autolens_profiling:
results/notes/numba_curvature_matrix_f_split.md.

[Jammy2211]

Library PR Created

PyAutoArray: #506 (pending-release)

pytest test_autoarray/ in the task worktree: 1296 passed.

API changes: internal to the numba imaging-inversion path. The numba-only
inversion_imaging_numba_util.curvature_matrix_mirrored_from is removed (F now
fills both triangles directly) — the identically-named functions in
inversion_util and inversion_imaging_util are untouched, and the only
downstream workspace reference uses that untouched non-numba namespace. Added:
inversion_imaging_numba_util.curvature_matrix_off_diags_via_mapper_and_blurred_curvature_weights_from
and the cached Convolver.reversed_kernel property (convolver.py is purely
additive, 48 added / 0 removed). No user-facing class, signature or default
changed, so no workspace migration is needed.

Heart readiness — human RED override. pyauto-heart readiness graded RED
score 45 at ship time (ts 2026-08-28T15:02:11Z). The diff repairs none of the
named reasons, so this is a plain human override, not the corrective-PR
exception. Verbatim reasons:

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, 2026-08-28, verbatim: "I authorize things to override the
heart RED." Scope: push + PR-open for this task only; merge and release stay
human, and the override does not extend to any other task.

The RED is pre-existing on main and unrelated to this branch — every failing
script is on the JAX likelihood path or in autofit, in other repos, while this
branch touches only the numba path plus an additive-only Convolver.

Next: /ship_workspace for autolens_profiling (breakdown artifacts +
results/notes/numba_curvature_matrix_f_split.md), behind the library-first
merge gate.

[Jammy2211]

Workspace PR Created

autolens_profiling: PyAutoLabs/autolens_profiling#189 (pending-release)

Harness split of F into its three sub-blocks, re-baselined breakdown artifacts
(stale since #462), and results/notes/numba_curvature_matrix_f_split.md with
the method, the paired before/after table, the correctness evidence and the
pool-run section. Only this task's claimed file sets are touched
(scripts/imaging/likelihood_breakdown/, results/breakdown/imaging/,
results/notes/), committed with explicit pathspecs per the parallel claim with
nuts-warm-start-driver-and-a100-probe.

One extra commit (e936aae): the lint job's build_readme.py --check
dashboard-idempotence leg failed because the re-baselined artifacts change three
generated headline rows. Regenerated with build_readme.py; only this task's
three cells move (imaging/delaunay_numba hst, imaging/pixelization_numba
euclid and hst), so no other claimant's rows are affected.

CI — all green

PR	check	result
PyAutoArray#506	Tests / unittest (3.12)	SUCCESS
PyAutoArray#506	Tests / unittest (3.13)	SUCCESS
PyAutoArray#506	Tests / unittest-nojax	SUCCESS
autolens_profiling#189	lint	SUCCESS

Both PRs are OPEN / MERGEABLE.

Library-first merge gate

PyAutoArray#506 merges first. autolens_profiling#189 must not merge before
it — the profiling repo's "after" columns and re-baselined artifacts were
measured against that branch. Neither PR has been merged; merge stays human.

[Jammy2211]

Shipped

Phase 1 merged.

• Library: perf(numba): speed up the curvature matrix F at HST resolution (#505) #506 — MERGED (merge commit 1b89404b, head d43e1aca)
• Workspace: profiling: split and re-baseline the numba curvature-matrix F breakdown (#505) autolens_profiling#189 — MERGED (merge commit b3fa632a, head e936aae7)

CI green on every run and every leg before merge: PyAutoArray Tests [pull_request] — unittest (3.12), unittest (3.13), unittest-nojax all success; autolens_profiling lint [pull_request] success.

Result

Measured at OMP_NUM_THREADS=1, memo off, n_repeats 10, before column re-measured back-to-back:

cell	eval	F total	mapper x l-func
hst, rectangular bilinear	1.562 -> 0.595	1.195 -> 0.359	0.858 -> 0.054
euclid, rectangular bilinear	0.349 -> 0.249	0.256 -> 0.097	0.182 -> 0.023
hst, Delaunay-1250	1.367 -> 0.758	1.077 -> 0.184	0.854 -> 0.052

~16x on the mapper x linear-func block, 3.3x on F, 2.6x on the whole HST evaluation — the issue's >=2x goal on F is met on both HST cells. All 3 pinned log-likelihoods PASSED; test_autoarray 1296 passed (6 new); the Nautilus pool run showed no thread oversubscription.

Follow-up

F is no longer the dominant term at HST resolution. Phase 2 (the mapper x mapper sparse-operator block and the MGE operated mapping matrix) is filed as PyAutoArray#507.

Merged over the pre-existing Heart RED "release validation FAILED (stage integrate)" under explicit human authorisation dated 2026-08-28 ("prm and then kick off phase 2, I authorize the heart RED thing"), scoped to this task only. The RED is unrelated to this diff (JAX-path and autofit scripts in other repos).