Skip to content

perf(numba): speed up the curvature matrix F at HST resolution (#505) - #506

Merged
Jammy2211 merged 3 commits into
mainfrom
feature/numba-hst-curvature-matrix-speedup
Aug 28, 2026
Merged

Jammy2211 merged 3 commits into
mainfrom
feature/numba-hst-curvature-matrix-speedup

Conversation

@Jammy2211

Copy link
Copy Markdown
Collaborator

Summary

Phase 1 of the numba CPU-path speed-up of the curvature matrix F at HST
resolution (#505). Two changes:

  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.
    They are now per-block helpers over one shared weighted copy, filling both
    triangles directly. Verified bit-identical (np.array_equal True on 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 runs 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.

Measured at OMP_NUM_THREADS=1, memo off, n_repeats 10, with the before column
re-measured back-to-back with the after run (this host carries 20-30 %
session-to-session variance):

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

~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. RTU is shown in
parentheses because it is GPU-only by the 2026-08-28 decision and was re-run
once for currency rather than paired.

F is no longer the dominant term at HST resolution: the largest remaining steps
on hst bilinear are the mapper x mapper sparse-operator block (0.277 s) and the
MGE operated mapping matrix (~0.224 s). That is the phase-2 candidate; no
phase-2 prompt is filed by this task since the goal is met.

API Changes

Internal to the numba imaging-inversion path. One numba-only helper
(inversion_imaging_numba_util.curvature_matrix_mirrored_from) is removed
because F now fills both triangles directly; the identically-named functions in
inversion_util and inversion_imaging_util are untouched, and the only
downstream reference in the workspaces uses that untouched non-numba namespace.
One numba util is added for the scatter half of the split block, and Convolver
gains a cached reversed_kernel property. convolver.py is purely additive
(48 lines added, 0 removed). No user-facing class, signature or default changes.
See full details below.

Test Plan

  • pytest test_autoarray/ in the task worktree: 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.

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

  • Nautilus pool run — no thread oversubscription. A single-thread win
    only counts if it survives the one-process-per-core pool, and an FFT
    spinning up its own threads would show as a pool regression. Ran
    autolens_workspace's imaging/features/pixelization/cpu_fast_modeling.py
    (first, non-SLaM fit) on an 8-core host against canonical main and
    against this branch via PYTHONPATH (autoarray.__file__ asserted both
    ways), number_of_cores 2 -> 8, with a 60-component linear MGE bulge added
    so the FFT'd block is non-empty:

    | | pool of 8 (median of 3) | serial (1 core) |
    |---|---|---|
    | 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), which is the
    number that would drop on hidden threads. Consistent with the code:
    `Convolver`'s FFT path uses `np.fft.rfft2` (no `workers=`, single-threaded)
    and `scipy.fft` appears only as `next_fast_len`, a shape calculation.
    

Profiling artifacts and the full method: autolens_profiling,
results/notes/numba_curvature_matrix_f_split.md.

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. Reasons quoted
verbatim from pyauto-heart readiness --json:

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
this override does not extend to any other task.

The RED is pre-existing on main and unrelated to this branch.
~/.pyauto-heart/validation_report.json (ts 2026-08-28T14:59:58Z):
validation_outcome: fail, release_ready: false, stages.integrate = fail
(run 33177898708), stages.rehearse = pass, totals 692 passed / 3 failed /
1 timeout. The failing scripts are autofit scripts/plot/nautilus_plotter.py,
autolens_test scripts/imaging/jax_likelihood/rectangular_mge.py,
.../rectangular_mge_rtu.py (the known f0ef8f2 pin drift) and, as the timeout,
autolens_test scripts/multi_dataset/jax_likelihood/delaunay.py. All live in
other repos and all are on the JAX likelihood path or in autofit; this
branch touches only the numba imaging-inversion path plus an additive-only
Convolver.

Full API Changes (for automation & release notes)

Removed

  • autoarray.util.inversion_imaging_numba.curvature_matrix_mirrored_from(curvature_matrix) —
    no longer needed; the numba F assembly now writes both triangles as it goes.
    Not a removal of inversion_util.curvature_matrix_mirrored_from or
    inversion_imaging_util.curvature_matrix_mirrored_from, which are unchanged
    and still used by the interferometer and non-numba imaging paths.

Added

  • autoarray.util.inversion_imaging_numba.curvature_matrix_off_diags_via_mapper_and_blurred_curvature_weights_from(data_to_pix_unique, data_weights, pix_lengths, pix_pixels, blurred_curvature_weights) —
    the sparse scatter half of the mapper x linear-func block, taking curvature
    weights that have already been correlated with the PSF. The dense
    curvature_matrix_off_diags_via_mapper_and_linear_func_curvature_vector_from
    is retained and is the reference the FFT path is asserted against in the tests.
  • autoarray.operators.convolver.Convolver.reversed_kernel — cached property
    returning a Convolver built on the kernel reversed along both axes,
    inheriting this convolver's use_fft policy. Turns a PSF correlation into a
    convolution the batched FFT path can run.

Changed Behaviour

  • The numba curvature matrix F is assembled from one shared
    operated_mapping_matrix / noise_map ** 2 copy via per-block helpers, and the
    mapper x linear-func block is computed by FFT convolution rather than a dense
    sliding window. F is bit-identical for change 1 and agrees to 3e-18 relative
    after change 2; pinned log-likelihoods are unchanged.

Migration

  • None required. No public class, signature or default changed; the one removed
    symbol has no reference outside the module it was removed from.

Generated by the PyAutoLabs agent workflow.

Jammy2211 and others added 3 commits August 28, 2026 15:31
…#505)

`_curvature_matrix_func_list_and_mapper` assembled all three blocks of F in a
single pass, so the breakdown harness could only time F as one step.

Split the two loops out into private helpers that write their block into the
`curvature_matrix` they are passed and return it:

- `_curvature_matrix_mapper_func_blocks_from` — the mapper x linear-func loop
- `_curvature_matrix_func_func_blocks_from`   — the linear-func x linear-func loop

`_curvature_matrix_mapper_diag` (the mapper x mapper block) already existed.
`_curvature_matrix_func_list_and_mapper` now composes the three in the same
order, and the `curvature_matrix` cached property is untouched (mirror + diag
add unchanged).

Pure code motion — no behaviour change, bit-identical output.
`pytest test_autoarray/inversion`: 392 passed.

Step 0 of #505, so autolens_profiling can time each block separately.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SqrSVGPrFcUB1vvDsoTw3n
…#505)

Three redundant passes over F in the numba sparse imaging inversion, none of
which changed a single value:

1. A global `curvature_matrix_mirrored_from` ran over the whole (P, P) matrix
   after assembly. Every block was already symmetric or had a known transpose:
   the mapper x mapper blocks are folded and mirrored inside
   `curvature_matrix_via_sparse_operator_from`, and the linear-func x
   linear-func blocks already wrote both triangles. The mirror is removed and
   the two off-diagonal block writers (mapper x mapper, mapper x linear-func)
   now place their transpose alongside the block, so F leaves assembly
   symmetric.
2. `_curvature_matrix_mapper_diag` wrapped the three `mapper.unique_mappings`
   arrays in `np.array(...)`, copying them on every evaluation. They are
   already contiguous ndarrays of the dtype the numba kernel wants (see
   `UniqueMappings.__init__`, which casts on construction) and every other
   caller in this module passes them through untouched.
3. `_curvature_matrix_mapper_func_blocks_from` formed
   `operated_mapping_matrix / noise_map ** 2` inside the mapper loop and then
   copied the result again with `np.array`. The weights do not depend on the
   mapper, so they are formed once per linear func ahead of the loop and passed
   straight to the kernel.

`curvature_matrix_mirrored_from` in `inversion_imaging_numba_util` had no
remaining caller (in this repo or downstream) and is deleted rather than left
as dead numba code.

Also corrects the `curvature_matrix` docstring: it claimed the property is
"not a cached property" and is overwritten in memory by the regularization
add, but it is decorated `@cached_property` and `curvature_reg_matrix` adds
out-of-place via `np.add`.

Bit-identical output verified by capturing `inversion.curvature_matrix` before
and after on the autolens_profiling breakdown fiducial: `np.array_equal` True
for both euclid and hst, with the figures of merit matching the pins exactly
(hst 27661.910133664103, euclid 6213.3068738858765).

`pytest test_autoarray/inversion`: 392 passed.

Step 1 of #505.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Fr6iJ5T1RDARWfxttCuGkK
…ix (#505)

The mapper x linear-func block of `F` was 70-85 % of F on every profiled cell
(0.953 s of 1.275 s at HST resolution) and did not scale with the source mesh,
only with image pixels. The cost was the dense sliding window inside
`curvature_matrix_off_diags_via_mapper_and_linear_func_curvature_vector_from`:
for every one of the 60 linear-func columns it expanded the noise-weighted
curvature weights onto the native grid and ran an `ny x nx x ky x kx` correlation
in numba.

That correlation is exactly a convolution with the PSF reversed along both axes,
so it now runs through the existing batched FFT convolver:

- `Convolver.reversed_kernel` (new, cached) is the same convolver with its kernel
  reversed, reusing the preloaded `ConvolverState` geometry rebuilt for the
  reversed kernel. Cached because a dataset's PSF outlives the `Inversion` that
  is rebuilt for every likelihood evaluation, so the reversed kernel's FFT
  geometry is built once per fit rather than once per evaluation.
- `InversionImagingSparseNumba._blurred_curvature_weights_from` correlates a
  linear func's curvature weights through it, once per linear func rather than
  once per (mapper, linear func) pair, with the weights zero outside the mask
  (no blurring mapping matrix) exactly as the sliding window had them.
- `curvature_matrix_off_diags_via_mapper_and_blurred_curvature_weights_from`
  (new) is the scatter half of the old kernel: the genuinely sparse, irregular
  accumulation onto source pixels, which stays in numba.

The old dense kernel and `convolve_with_kernel_native` are kept, unused by the
inversion, as the reference the new path is asserted against in the tests.

The numpy convolution path is `scipy.signal.convolve(..., mode="same")`, whose
`(k - 1) // 2` "same" offset matches the kernel's `k // 2` centre for both odd
and even kernel widths, and which spawns no thread pool (the profiling campaign
runs one process per core).

Tests:

- `test_inversion_imaging_util.py` asserts the FFT block equals the dense kernel
  (rel 1e-6) on a small masked grid, parametrized over the file's existing
  asymmetric, non-square `KERNELS_ODD`; a control run without the reversal fails
  on all four. It also asserts `reversed_kernel` is the reversed kernel and is
  cached.
- `test_curvature_matrix_func_list_blocks.py` gains an inversion-level test of
  `_curvature_matrix_mapper_func_blocks_from` against the dense kernel with an
  asymmetric PSF, which also pins that the block's transpose is written (the
  global mirror having been removed in the previous commit). A control run
  without the reversal fails.

`pytest test_autoarray`: 1296 passed (1290 + 6 new).

Seconds per evaluation, `OMP_NUM_THREADS=1 AUTOARRAY_NUMBA_OPERATED_MEMO=0`,
n_repeats 10 (autolens_profiling breakdown cells):

| cell                | eval        | F           | mapper x l-func |
|---------------------|-------------|-------------|-----------------|
| hst bilinear        | 1.738 -> 0.819 | 1.275 -> 0.387 | 0.953 -> 0.0629 |
| euclid bilinear     | 0.444 -> 0.273 | 0.284 -> 0.0889 | 0.198 -> 0.0224 |
| hst Delaunay-1250   | 1.455 -> 0.754 | 1.055 -> 0.241 | 0.899 -> 0.0612 |
| hst rectangular RTU | 8.501 -> 7.058 | 1.250 -> 0.335 | 0.929 -> 0.0630 |

i.e. 15x on the block, 3.3x on F and 2.1x on the whole HST evaluation. Pinned
log-likelihoods PASSED on all three pinned cells; F itself agrees with the
sliding-window result to 3e-18 relative, and the hst figure of merit moves from
27661.910133664103 to 27661.91013366411 (9e-18 relative, well inside the
harness rtol of 1e-4).

Step 2 of #505.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Fr6iJ5T1RDARWfxttCuGkK
@Jammy2211

Copy link
Copy Markdown
Collaborator Author

Linked workspace PR: PyAutoLabs/autolens_profiling#189 (profiling harness split + re-baselined breakdown artifacts + the pool-run notes).

Library-first merge gate: this PR (PyAutoArray#506) merges first; autolens_profiling#189 must not merge until it is MERGED, since the profiling repo's "after" columns were measured against this branch.

@Jammy2211
Jammy2211 merged commit 1b89404 into main Aug 28, 2026
3 checks passed
@Jammy2211
Jammy2211 deleted the feature/numba-hst-curvature-matrix-speedup branch August 28, 2026 21:16
@Jammy2211 Jammy2211 removed the pending-release PR queued for the next release build label Sep 4, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant