Repository navigation
Description
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_profilingbreakdown harness; record the split at hst + euclid, rectangular (Bilinear + RTU) + Delaunay-1250, withOMP_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.1artifacts 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
Convolvermachinery (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_autoarraygreen; smokecpu_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 inautolens_profilingresults + 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
PyAutoArray/autoarray/inversion/inversion/imaging_numba/sparse.py: split_curvature_matrix_func_list_and_mapperinto 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.autolens_profiling/scripts/imaging/likelihood_breakdown/pixelization_numba.py+delaunay_numba.py: addSTEP_ACCESSORSentries for each sub-block (mapper diag → mapper×func → func×func → mirror) ahead of the existingCurvature matrix Fentry (incremental-cost scheme attributes the time); keep the F entry as the total.- Run the four cells (
--instrument hst|euclid,--rect-mesh bilinear|rtu;_profile_cli.py:128-140) withOMP_NUM_THREADS=1 AUTOARRAY_NUMBA_OPERATED_MEMO=0; commit JSON+PNG underresults/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-512PyAutoArray/autoarray/operators/convolver.py— batched FFTConvolver/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 testsautolens_profiling/scripts/imaging/likelihood_breakdown/{pixelization_numba,delaunay_numba}.py—STEP_ACCESSORSharness;autolens_profiling/_profile_cli.py— instrument /--rect-meshselection;results/breakdown/imaging/— artifacts
Verification
test_autoarraygreen 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.pysmoke 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; recordscomplete/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 defaultRectangularBilinearAdaptDensity
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).
- Step 0 — decompose F itself. The breakdown cells time
fit.inversion.curvature_matrixas 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_rtuvariants —
numbers incomplete/2026/08/numba-cpu-kernel-cdf-fast-path.md; regenerate, do not hand-copy). - 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.prangeover the outer loop with the OMP thread count (the campaign runs one process per
core under Nautilus, so gains must be measured withOMP_NUM_THREADS=1AND 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.
- the unique-mappings compression (
- Ship behind exact parity (bit-level where the summation order is unchanged, pinned rtol 1e-6
otherwise);test_autoarraygreen; smokeautolens_workspace/scripts/imaging/features/pixelization/ cpu_fast_modeling.pyon the smoke profile. - 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).
Activity
Jammy2211 commentedon Aug 28, 2026
Step 0 — F split (OMP_NUM_THREADS=1, memo off, n_repeats=10, v2026.8.17.1)
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_profiling6f7d20a.Jammy2211 commentedon Aug 28, 2026
Measured: F is 3.3x faster, the HST evaluation 2.6x
Two changes on
feature/numba-hst-curvature-matrix-speedup:operated_mapping_matrix / noise_map ** 2and re-walked the mapping matrix.Split into per-block helpers over one shared weighted copy. Bit-identical
(
np.array_equalTrue on both euclid and hst).correlation with the PSF, i.e. a convolution with the PSF reversed along
both axes. It now goes through the existing batched
Convolver(new cachedConvolver.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.
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
hst (the FFT is not bit-identical to the direct sum, as expected).
(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 testsassert 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'simaging/features/pixelization/cpu_fast_modeling.py(first, non-SLaM fit) on an8-core host, against canonical
mainand against the worktree viaPYTHONPATH(
autoarray.__file__asserted both ways),number_of_coresraised 2 -> 8.Two forced departures from the smoke profile:
PYAUTO_TEST_MODE=1rather thanthe default
2(level 2 bypasses the sampler, so the smoke profile as writtenruns no pool at all; level 1 sets
n_like_max=1= one 100-evaluationexploration batch), and
PYAUTO_SMALL_DATASETSleft unset (under the 15x15 capthe 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_listholds 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.pyrecipe) makes the block non-empty:
number_of_cores=1)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 FFTpath uses
np.fft.rfft2(noworkers=, single-threaded);scipy.fftappearsonly 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 commentedon Aug 28, 2026
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_fromis removed (F nowfills both triangles directly) — the identically-named functions in
inversion_utilandinversion_imaging_utilare untouched, and the onlydownstream workspace reference uses that untouched non-numba namespace. Added:
inversion_imaging_numba_util.curvature_matrix_off_diags_via_mapper_and_blurred_curvature_weights_fromand the cached
Convolver.reversed_kernelproperty (convolver.pyis purelyadditive, 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 readinessgraded REDscore 45 at ship time (ts
2026-08-28T15:02:11Z). The diff repairs none of thenamed reasons, so this is a plain human override, not the corrective-PR
exception. Verbatim reasons:
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
mainand unrelated to this branch — every failingscript is on the JAX likelihood path or in
autofit, in other repos, while thisbranch touches only the numba path plus an additive-only
Convolver.Next:
/ship_workspaceforautolens_profiling(breakdown artifacts +results/notes/numba_curvature_matrix_f_split.md), behind the library-firstmerge gate.
Jammy2211 commentedon Aug 28, 2026
Workspace PR Created
autolens_profiling: PyAutoLabs/autolens_profiling#189 (
pending-release)Harness split of
Finto its three sub-blocks, re-baselined breakdown artifacts(stale since #462), and
results/notes/numba_curvature_matrix_f_split.mdwiththe 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 withnuts-warm-start-driver-and-a100-probe.One extra commit (
e936aae): thelintjob'sbuild_readme.py --checkdashboard-idempotence leg failed because the re-baselined artifacts change three
generated headline rows. Regenerated with
build_readme.py; only this task'sthree cells move (
imaging/delaunay_numbahst,imaging/pixelization_numbaeuclid and hst), so no other claimant's rows are affected.
CI — all green
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.
Merge pull request #506 from PyAutoLabs/feature/numba-hst-curvature-m…
Jammy2211 commentedon Aug 28, 2026
Shipped
Phase 1 merged.
1b89404b, headd43e1aca)b3fa632a, heade936aae7)CI green on every run and every leg before merge: PyAutoArray
Tests [pull_request]—unittest (3.12),unittest (3.13),unittest-nojaxall success; autolens_profilinglint [pull_request]success.Result
Measured at
OMP_NUM_THREADS=1, memo off, n_repeats 10, before column re-measured back-to-back:~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_autoarray1296 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).