feat: numba CPU interferometer curvature path — direct_conv, geometry-gated (#543) - #545
Merged
Merged
Conversation
…-gated (#543) Adds `autoarray/inversion/inversion/interferometer_numba/`, a numba CPU `direct_conv` curvature path for single-mapper interferometer inversions. The kernel is ported from the autolens_profiling investigation pack (issue #226): a serial `@numba_util.jit()` `curvature_direct_conv`, plus a lazily-built `prange` variant compiled on first use rather than decorated at import time. `InversionInterferometerSparseNumba` subclasses `InversionInterferometerSparse` and overrides `curvature_matrix_diag` only. It raises `InversionException` on every configuration the kernel cannot represent (non-NumPy `xp`, more than one mapper, any `AbstractLinearObjFuncList`, over-sampling), so a direct construction fails loudly where the factory's routing merely falls through. `Settings.interferometer_numba_nnz_per_source_max` (default 60.0, read from `autoarray/config/general.yaml` with a KeyError fallback to the same value; 0 disables the path) is the geometry gate, in mean non-zeros per source column. `inversion_interferometer_from` routes a single-mapper, no-func-list, no-over-sampling, `xp=np` interferometer inversion at or below the gate to the numba class; everything else keeps the FFT route. Behaviour change: the default CPU interferometer route changes with this PR. Sparse-operator interferometer fits with `xp=np`, one mapper, no over-sampling and a sparse-enough mapping operator now run the numba curvature path by default when numba is installed. Validated against the sparse NumPy path: curvature matrix F max abs difference 1.8e-15, data vector D exact, reconstruction 4.9e-17; the parallel kernel is exactly equal to the serial one; a 1 % perturbation control test fails as it must. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018hLF3ZAcz5MmaSJBEcLkvF
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds a numba CPU
direct_convcurvature path for single-mapper interferometerinversions, and routes to it automatically when the mapping operator is sparse
enough for it to win.
Design. New package
autoarray/inversion/inversion/interferometer_numba/:inversion_interferometer_numba_util.py— the kernel, ported from theautolens_profiling investigation pack (autolens_profiling#226). Instead of
transforming each source column to the uv-plane,
curvature_direct_convconvolves each source column of the mapping operator over the real-space
extent rectangle, at a cost that scales with the column's non-zeros. Serial
form is a
@numba_util.jit()function; aprange-parallel form is compiledlazily on first use (
direct_conv_parallel_kernel()) rather than decorated atimport time, so the parallel compile is never paid by a serial run.
sparse.py—InversionInterferometerSparseNumba, a subclass ofInversionInterferometerSparsethat overridescurvature_matrix_diagonly.Everything else (data vector, regularization, solve, evidence) is inherited,
so the two paths cannot diverge outside the curvature matrix.
Settings.interferometer_numba_nnz_per_source_max— the geometry gate,default
60.0, packaged inautoarray/config/general.yaml;0disables thenumba path entirely.
inversion_interferometer_fromgains_use_interferometer_numba(...), whichroutes to the numba class only when all of:
xp is np; the gate is > 0;exactly one
Mapperand noAbstractLinearObjFuncList; no over-sampling(
sub_fraction == 1); mean non-zeros per source column at or below the gate;import numbasucceeds.The gate, and its measured crossovers
The numba kernel's cost scales with the mapping operator's density; the FFT
route's does not. They therefore cross at a roughly fixed number of non-zeros
per source column,
mapper.pix_sizes_for_sub_slim_index.sum() / mapper.params— which is exactly the gate quantity.
Well below the crossover the numba kernel is 2–7× faster than the JAX/FFT
route on CPU. Measurements and method: autolens_profiling#226 verdict, §2. The
default takes the conservative of the two (60.0).
This constant is machine-dependent — it is set by the ratio of scalar AXPY
throughput to FFT throughput on the CPU running the fit. A machine with a very
different cache hierarchy or FFT library will cross somewhere else, so
re-measure before tuning it. The config comment and the
Settingspropertydocstring both say so.
Parity
Validated in
test_autoarray/inversion/inversion/interferometer_numba/test_interferometer_numba.pyagainst the existing sparse NumPy path:
FDBehaviour change (read this)
The default CPU interferometer route changes with this PR. A sparse-operator
interferometer inversion with
xp=np, a single mapper, no linear-objectfunction lists and no over-sampling, whose mapping operator is at or below the
gate, now goes through the numba
direct_convcurvature path by default whennumba is installed — where before it always went through the FFT route.
The results are numerically equivalent (table above); what changes is which
kernel runs, and therefore run time and the first-call compile.
To restore the previous behaviour unconditionally, set in
general.yaml:or pass
Settings(interferometer_numba_nnz_per_source_max=0).The new key lives in PyAutoArray's own packaged
general.yamland the propertyfalls back to
60.0onKeyError, so no workspace config needs to change.Preconditions
InversionInterferometerSparseNumbaraisesInversionExceptionon: anon-NumPy
xp, more than one mapper, anyAbstractLinearObjFuncList, andover-sampling. The factory checks the same conditions and simply falls through
to
InversionInterferometerSparseinstead — the routing is silent, the directconstruction is loud, and the two share the meaning of the conditions so they
cannot drift.
API Changes
Additive only — nothing removed, renamed or re-signatured.
InversionInterferometerSparseNumbaand its kernel module (a newsub-package, not exported at the
autoarraytop level).Settingsargument / propertyinterferometer_numba_nnz_per_source_max.general.yamlkey of the same name, with a packaged default.inversion_interferometer_from: CPU single-mapperinterferometer inversions at or below the gate now route to the numba path.
See the behaviour-change section above.
See full details below.
Test Plan
pytest test_autoarray -qinside the task worktree — 1501 passed.black --checkon all changed files.perturbation control, all in the new test file.
that
interferometer_numba_nnz_per_source_max: 0is a sufficient escapehatch.
Stacked on
Branched from
feature/interferometer-sparse-operator-numpy-cpu-path(PyAutoArray#544, merged 2026-09-08). This branch therefore diffs cleanly
against
main— the diff is task 4's seven files only, no rebase required.Follow-on
The crossover numbers above were measured against the standalone investigation
pack, not the library dispatch. The in-situ re-measurement — running
delaunay_numba.py's arms throughaa.Inversionso the library's own routingis what is timed, and confirming the 60/77 crossover on that machine — is filed
as
PyAutoMind/draft/research/autolens_profiling/interferometer_numba_library_dispatch_insitu.md.Full API Changes (for automation & release notes)
Added
autoarray.inversion.inversion.interferometer_numba.sparse.InversionInterferometerSparseNumba(dataset, linear_obj_list, settings=, xp=, preloads=)—InversionInterferometerSparsesubclass overridingcurvature_matrix_diagwith the numbadirect_convkernel; raisesInversionExceptionon non-NumPyxp, multiple mappers, anyAbstractLinearObjFuncList, or over-sampling.autoarray.inversion.inversion.interferometer_numba.inversion_interferometer_numba_util.curvature_direct_conv(...)— the serial numba kernel.autoarray.inversion.inversion.interferometer_numba.inversion_interferometer_numba_util.direct_conv_parallel_kernel()— lazily compiles theprangevariant.autoarray.inversion.inversion.interferometer_numba.inversion_interferometer_numba_util.kernel_inputs_from(...),nnz_per_source_column_from(mapper).autoarray.settings.Settings(interferometer_numba_nnz_per_source_max=None)and the propertySettings.interferometer_numba_nnz_per_source_max— default60.0.general.yamlkeyinversion: interferometer_numba_nnz_per_source_max: 60.0.Changed Signature
autoarray.settings.Settings.__init__— new optional keywordinterferometer_numba_nnz_per_source_max(defaultNone→ packaged value). Backwards compatible.Changed Behaviour
autoarray.inversion.inversion.factory.inversion_interferometer_from— whendataset.sparse_operator is not None and use_sparse_operator, a single-mapper, no-func-list, no-over-sampling inversion withxp is npwhose mean non-zeros per source column is<= settings.interferometer_numba_nnz_per_source_maxnow returnsInversionInterferometerSparseNumbainstead ofInversionInterferometerSparse. Numerically equivalent (F 1.8e-15, D exact, reconstruction 4.9e-17); the kernel and the run time change.Migration
general.yaml→inversion: interferometer_numba_nnz_per_source_max: 0, orSettings(interferometer_numba_nnz_per_source_max=0).Closes #543
Generated by the PyAutoLabs agent workflow.
🤖 Generated with Claude Code
https://claude.ai/code/session_018hLF3ZAcz5MmaSJBEcLkvF