Work Classification

Library

Affected Repositories

• PyAutoArray (primary)
• PyAutoGalaxy

Branch Survey

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).

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).