perf: batch MGE convolution + cache operated matrices (numba CPU likelihood, phase 1)

[Jammy2211]

Overview

Phase 1 of the numba CPU sparse-operator likelihood speed restoration (epic numba-cpu-likelihood; profiling in autolens_profiling#151). On the cpu_fast_modeling.py route (MGE-60 linear lens light + rectangular pixelization, use_jax=False) the 60 MGE operated images cost ~19% of a euclid evaluation because operated_mapping_matrix_override convolves each Gaussian separately — and each call rebuilds a ConvolverState (mask resize + rfft2) because Convolver.state_from never reuses the precomputed state. This task batches the convolution through the existing convolved_mapping_matrix_via_real_space_np_from, fixes state reuse, caches the base linear_func_operated_mapping_matrix_dict, and hoists redundant work out of the O(N²) linear-func pair loop. Likelihood values must be unchanged (rtol 1e-6, expected bit-comparable).

Plan

• Batch the 60 MGE convolutions into one convolved_mapping_matrix_via_real_space_np_from call on the numpy path of operated_mapping_matrix_override; JAX and oversampled branches untouched.
• Fix Convolver.state_from so a precomputed state is reused when its mask matches (removes the per-call rfft2).
• Make the base linear_func_operated_mapping_matrix_dict a cached_property so the non-numba sparse inversion also builds it once per evaluation (the numba subclass already does).
• Hoist the per-pixel noise division out of the linear-func pair loop in both sparse inversions and compute only the upper triangle, mirroring it (exact for any per-pixel noise map — the middle factor is diag(1/σ²)).
• Tests: blurring-mask ordering parity, batched == per-Gaussian columns to 1e-14, state reuse, pair-loop equality against brute force with a random spatially varying noise map.
• Re-run the autolens_profiling pixelization_numba breakdown + runtime at hst (pins) and euclid (compare to v2026.8.17.1 JSON); record the MGE-step gain here.

Detailed implementation plan

Work Classification

Library

Affected Repositories

• PyAutoArray (primary)
• PyAutoGalaxy

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoArray	main	clean
./PyAutoGalaxy	main	clean

No worktree claims on either repo (worktree_check_conflict exit 0).

Suggested branch: feature/numba-cpu-mge-batch-convolve-cache
Worktree root: ~/Code/PyAutoLabs-wt/numba-cpu-mge-batch-convolve-cache/

Implementation Steps

PyAutoGalaxy — autogalaxy/profiles/light/linear/abstract.py (operated_mapping_matrix_override, lines 319-397)

1. After the psf is None and LightProfileOperated early returns, add a branch if self._xp is np and self.psf.convolve_over_sample_size == 1:
   • stack light_profile.image_2d_from(grid=self.grid) columns into an (N_pix, N_gauss) matrix and image_2d_from(grid=self.blurring_grid) into (N_blur, N_gauss);
   • return self.psf.convolved_mapping_matrix_via_real_space_np_from(mapping_matrix=..., mask=self.grid.mask, blurring_mapping_matrix=..., blurring_mask=self.blurring_grid.mask).
   • The existing per-profile loop remains for the JAX and oversampled paths.
2. Risk: convolved_mapping_matrix_via_real_space_np_from ignores its blurring_mask argument and scatters on state.blurring_mask (derived from the resized mask + kernel shape, convolver.py:1437-1560). The linear func's blurring_grid comes from mask.derive_mask.blurring_from(kernel_shape_native, allow_padding=True). Verify slim ordering is identical; if not, raise a clear error (no silent guard) rather than reorder silently.
3. Tests in test_autogalaxy/profiles/light/linear/test_abstract.py: extend test__operated_mapping_matrix__columns_match_individual_blurred_images (bright Gaussian near the mask edge so blurring flux is exercised; abs 1e-14 vs per-profile convolved_image_from).

PyAutoArray — autoarray/operators/convolver.py
4. state_from (lines 308-326): the shape test compares the mask to the kernel, so _state is never returned for convolve_over_sample_size == 1. Return the precomputed state when its (pre-resize) mask matches the requested mask; otherwise build a fresh one. Test in test_autoarray/operators/test_convolver.py: identical object returned on repeat call; output bit-identical to a fresh state.
5. Add a direct test of convolved_mapping_matrix_via_real_space_np_from with a blurring_mapping_matrix, against N × convolved_image_via_real_space_np_from (currently no test covers that argument on the numpy path).

PyAutoArray — autoarray/inversion/inversion/imaging/abstract.py
6. from autonerves import cached_property; linear_func_operated_mapping_matrix_dict (line 183-217) → cached_property. Per-evaluation lifetime (one inversion per FitImaging), consistent with the numba subclass override at imaging_numba/sparse.py:103.

PyAutoArray — imaging_numba/sparse.py:545-563 and imaging/sparse.py:452-472 (linear-func × linear-func curvature blocks)
7. Precompute weighted = [dict[f] / noise_map[:, None] for f in linear_func_list] once; loop index_1 >= index_0, np.dot(weighted[i].T, weighted[j]), assign the block and its transpose. Exact for any per-pixel (non-constant, non-symmetric) RMS noise map since block(i,j) = M_iᵀ diag(1/σ²) M_j; only ulp-level BLAS ordering differences in the mirrored half.
8. Test (host in test_sparse_numba_operated_memo.py's StubInversion or a new sibling): random spatially varying noise map, ≥3 linear funcs, mirrored result equals brute-force double loop to 1e-12.

Measurement (autolens_profiling, no source edits)
9. Run scripts/imaging/likelihood_breakdown/pixelization_numba.py and scripts/imaging/likelihood_runtime/pixelization_numba.py at hst (pins bilinear 27661.910133665442 / rtu 27180.704715698186 must hold, rtol 1e-6) and euclid (compare to results/breakdown/imaging/pixelization_numba_breakdown_euclid_v2026.8.17.1.json). Record the "MGE operated mapping matrix (60 funcs)" step before/after on this issue (expect 0.42 s → ~0.05-0.1 s at euclid).

Key Files

• PyAutoGalaxy/autogalaxy/profiles/light/linear/abstract.py — per-Gaussian convolution loop to batch
• PyAutoArray/autoarray/operators/convolver.py — state_from, convolved_mapping_matrix_via_real_space_np_from, ConvolverState
• PyAutoArray/autoarray/inversion/inversion/imaging/abstract.py — base linear_func_operated_mapping_matrix_dict
• PyAutoArray/autoarray/inversion/inversion/imaging_numba/sparse.py, .../imaging/sparse.py — pair loop
• autolens_profiling/scripts/imaging/likelihood_breakdown/pixelization_numba.py — measurement + pins

Epic

numba-cpu-likelihood — profiling (done, autolens_profiling#151), first-call garbage bug (done), phase 1 = this, phase 2 = kernel-CDF numba fast path (deferred), phase 3 = numba fnnls solver restoration (not yet filed).

Original Prompt

Click to expand starting prompt

Numba CPU likelihood phase 1: batched MGE convolution + operated-matrix caching

Type: feature
Target: autoarray
Repos:

• @PyAutoArray
• @PyAutoGalaxy
  Difficulty: medium
  Autonomy: supervised
  Priority: high
  Status: formalised
  Filed: 2026-08-20 (backfilled from git)

Phase 1 of the CPU-likelihood speed restoration
(autolens_profiling#151 profiling; user request 2026-08-20 recorded verbatim
in the phase-2 prompt numba_cpu_likelihood_kernel_cdf_fast_path.md).
Exact-identical wins on files disjoint from phase 2's kernel-CDF work.

Context (from the 2026-08-20 profiling + source hunt)

On the numba CPU sparse-operator likelihood (apply_sparse_operator_cpu() +
use_jax=False, MGE-60 linear lens light + rectangular pixelization —
the cpu_fast_modeling.py production route):

1. The 60 MGE linear-Gaussian operated images cost ~19% of a euclid evaluation
   (0.42 s of 2.15 s; 0.87 s at hst): AbstractLinearObjFuncList. operated_mapping_matrix_override
   (PyAutoGalaxy autogalaxy/profiles/light/linear/abstract.py:319-382) loops
   the Gaussians and calls psf.convolved_image_from 60 separate times,
   each re-padding to fft_shape and re-transforming the PSF. A batched exact
   equivalent already exists and handles the blurring region:
   Convolver.convolved_mapping_matrix_via_real_space_np_from
   (PyAutoArray autoarray/operators/convolver.py:1437) — one scipy FFT
   convolution amortized over all 60 columns.
   AbstractLinearObjFuncList.mapping_matrix (linear/abstract.py:291)
   already produces the stacked unblurred matrix; only the blurring-grid stack
   is missing.
2. linear_func_operated_mapping_matrix_dict
   (PyAutoArray autoarray/inversion/inversion/imaging/abstract.py:184) is an
   uncached @property rebuilt on every access; the numba sparse inversion
   accesses it ~5 times per evaluation (imaging_numba/sparse.py:194,419,443, 451,509), including inside an O(60^2) loop that also repeats a
   (N_pix, 60) noise-map division per pair. Cache it (cached_property,
   consistent with the inversion's per-evaluation lifetime) and hoist the
   noise division out of the pair loop.

Goal

• Batch the MGE/linear-func operated mapping matrix construction through the
  existing batched convolver call (numpy path; JAX path untouched).
• Cache linear_func_operated_mapping_matrix_dict and hoist repeated
  per-pair work in imaging_numba/sparse.py.
• Likelihood values unchanged: pinned euclid/hst log-likelihoods in
  autolens_profiling's pixelization_numba cells must pass (rtol 1e-6; expect
  bit-comparable), plus the existing unit suites in both repos.
• Re-run the autolens_profiling runtime + breakdown cells to record the gain
  (expect the "MGE operated mapping matrix" step 0.42 s -> ~0.05-0.1 s at
  euclid).

[Jammy2211]

Implementation complete on feature/numba-cpu-mge-batch-convolve-cache (worktree, uncommitted — Heart RED gate pending human ack)

PyAutoArray (1237 passed): Convolver.state_from now reuses the precomputed state via ConvolverState.is_for_mask (old test compared mask to kernel shape → never matched); linear_func_operated_mapping_matrix_dict → cached_property (+ numba override .fget→.func companion fix); linear-func pair loop hoisted + upper-triangle mirrored in both sparse inversions. New tests incl. random non-symmetric noise map vs brute force (1e-12).
PyAutoGalaxy (1131 passed): numpy fast path in operated_mapping_matrix_override → one convolved_mapping_matrix_via_real_space_np_from call. Blurring-mask slim ordering verified identical to the ConvolverState frame (pure translation, 8+ configs) and pinned by test.

Measurements (this machine, OMP=1, loaded ±20-30%)

Trap found: the pixelization_numba breakdown harness re-evaluates one fixed instance, so the cross-eval memo hits and the MGE step reads ~0.003-0.01 s on both sides — blind to this change. Re-measured with AUTOARRAY_NUMBA_OPERATED_MEMO=0 on fresh FitImaging per call (a perturbed-instance memo-on variant agrees — that is the real-modelling regime):

MGE-60 operated matrix	before	after	ratio
hst	1.112 s	0.299 s	3.7×
euclid	0.678 s	0.125 s	5.4×

Matrices bitwise identical (max abs diff 0.0). Breakdown steps: Mapped reconstruction ~3.9× faster, Model data + chi² ~3.9×, Blurred image ~1.9× (state reuse). Curvature F / Mapper triplets / solve unchanged (as expected). hst pin 27661.910133664103 PASSED before and after; euclid LL bit-identical before/after (dataset auto-simulated; historical 5860.175 is pre-mesh-split so not comparable).

Follow-up (not this PR): make the breakdown/runtime harness perturb the instance or disable the memo for that step — otherwise the phase-2/3 measurements will inherit the same blind spot.

[Jammy2211]

Library PRs created

• PyAutoArray: perf: reuse Convolver state, cache operated-matrix dict, hoist linear-func pair loop (#496 phase 1) #497 (commit 0ff036ac, 1237 tests passed)
• PyAutoGalaxy: perf: batch MGE linear-func PSF convolution on the numpy path (PyAutoArray#496 phase 1) PyAutoGalaxy#588 (commit 10f5047e, 1131 tests passed)

Both labelled pending-release, independently mergeable. API changes: none — internal only. Workspace impact: none (workspace mentions of operated_mapping_matrix_override read an identical value). Shipped over Heart RED (release validation FAILED (stage integrate) — the unrelated rectangular_mge pin drift) on human authorisation. Smoke suites (autogalaxy / autolens / autolens_test) running against the worktree source; results to follow.

[Jammy2211]

Smoke tests (worktree source, PYAUTO_DISABLE_JAX=1 profile → numpy path exercised)

Workspace	Passed	Failed	Total
autogalaxy_workspace (scripts + notebooks)	16	0	16
autolens_workspace (scripts + notebooks)	37	0	37
autolens_workspace_test	21	3	24

The 3 autolens_workspace_test failures are pre-existing ENV: jax pin mismatches (interferometer/jax_likelihood/{rectangular,mge}.py, multi_dataset/jax_likelihood/mge.py) — the same run against canonical main fails the same 3 with bit-identical ACTUAL values, so the branch introduces zero new failures. multi_galaxy/features/pixelization/cpu_fast_modeling.py (the production route) PASS.

Source identity verified in the smoke env (ConvolverState.is_for_mask present, batched MGE convolve present; both absent on main).

Next: /prm once CI on #497 / PyAutoGalaxy#588 is green — merge stays human.

[Jammy2211]

Shipped

• PyAutoArray perf: reuse Convolver state, cache operated-matrix dict, hoist linear-func pair loop (#496 phase 1) #497 merged (86e2944a): Convolver state reuse, cached operated-matrix dict (+ numba .func fix), linear-func pair loop hoisted/mirrored.
• PyAutoGalaxy #588 merged (d55f3ab3): batched MGE PSF convolution on the numpy path.

MGE-60 operated matrix: hst 1.11 → 0.30 s, euclid 0.68 → 0.125 s; bitwise identical output, hst pin unchanged, smoke 0 new failures. Epic numba-cpu-likelihood phase 1 complete; follow-up filed: profiling harness memo-blindness (draft/feature/autolens_profiling/numba_breakdown_harness_memo_blind.md). Phase 3 (numba fnnls solver restoration) not yet filed.