Skip to content

feat: numba CPU interferometer curvature path — direct_conv, geometry-gated (#543) - #545

Merged
Jammy2211 merged 1 commit into
mainfrom
feature/interferometer-numba-cpu-direct-conv
Sep 8, 2026
Merged

Jammy2211 merged 1 commit into
mainfrom
feature/interferometer-numba-cpu-direct-conv

Conversation

@Jammy2211

Copy link
Copy Markdown
Collaborator

Summary

Adds a numba CPU direct_conv curvature path for single-mapper interferometer
inversions, 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 the
    autolens_profiling investigation pack (autolens_profiling#226). Instead of
    transforming each source column to the uv-plane, curvature_direct_conv
    convolves 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; a prange-parallel form is compiled
    lazily on first use (direct_conv_parallel_kernel()) rather than decorated at
    import time, so the parallel compile is never paid by a serial run.
  • sparse.pyInversionInterferometerSparseNumba, a subclass of
    InversionInterferometerSparse that overrides curvature_matrix_diag only.
    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 in autoarray/config/general.yaml; 0 disables the
    numba path entirely.
  • inversion_interferometer_from gains _use_interferometer_numba(...), which
    routes to the numba class only when all of: xp is np; the gate is > 0;
    exactly one Mapper and no AbstractLinearObjFuncList; no over-sampling
    (sub_fraction == 1); mean non-zeros per source column at or below the gate;
    import numba succeeds.

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.

mesh measured crossover (nnz / source column)
Delaunay ≈ 60
rectangular ≈ 77

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 Settings property
docstring both say so.

Parity

Validated in test_autoarray/inversion/inversion/interferometer_numba/test_interferometer_numba.py
against the existing sparse NumPy path:

quantity agreement
curvature matrix F max abs diff 1.8e-15
data vector D exact
reconstruction 4.9e-17
parallel kernel vs serial exact
1 % perturbation control fails, as it must

Behaviour 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-object
function lists and no over-sampling, whose mapping operator is at or below the
gate, now goes through the numba direct_conv curvature path by default when
numba 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:

inversion:
  interferometer_numba_nnz_per_source_max: 0

or pass Settings(interferometer_numba_nnz_per_source_max=0).

The new key lives in PyAutoArray's own packaged general.yaml and the property
falls back to 60.0 on KeyError, so no workspace config needs to change.

Preconditions

InversionInterferometerSparseNumba raises InversionException on: a
non-NumPy xp, more than one mapper, any AbstractLinearObjFuncList, and
over-sampling. The factory checks the same conditions and simply falls through
to InversionInterferometerSparse instead — the routing is silent, the direct
construction 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.

  • New class InversionInterferometerSparseNumba and its kernel module (a new
    sub-package, not exported at the autoarray top level).
  • New optional Settings argument / property
    interferometer_numba_nnz_per_source_max.
  • New general.yaml key of the same name, with a packaged default.
  • Changed behaviour of inversion_interferometer_from: CPU single-mapper
    interferometer 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 -q inside the task worktree — 1501 passed.
  • black --check on all changed files.
  • Parity vs the sparse NumPy path, parallel-vs-serial equality, and the 1 %
    perturbation control, all in the new test file.
  • Reviewer: confirm the behaviour change is acceptable as a default, and
    that interferometer_numba_nnz_per_source_max: 0 is a sufficient escape
    hatch.

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 through aa.Inversion so the library's own routing
is 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=)InversionInterferometerSparse subclass overriding curvature_matrix_diag with the numba direct_conv kernel; raises InversionException on non-NumPy xp, multiple mappers, any AbstractLinearObjFuncList, 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 the prange variant.
  • 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 property Settings.interferometer_numba_nnz_per_source_max — default 60.0.
  • general.yaml key inversion: interferometer_numba_nnz_per_source_max: 60.0.

Changed Signature

  • autoarray.settings.Settings.__init__ — new optional keyword interferometer_numba_nnz_per_source_max (default None → packaged value). Backwards compatible.

Changed Behaviour

  • autoarray.inversion.inversion.factory.inversion_interferometer_from — when dataset.sparse_operator is not None and use_sparse_operator, a single-mapper, no-func-list, no-over-sampling inversion with xp is np whose mean non-zeros per source column is <= settings.interferometer_numba_nnz_per_source_max now returns InversionInterferometerSparseNumba instead of InversionInterferometerSparse. Numerically equivalent (F 1.8e-15, D exact, reconstruction 4.9e-17); the kernel and the run time change.

Migration

  • None required. To keep the old route: general.yamlinversion: interferometer_numba_nnz_per_source_max: 0, or Settings(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

…-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
@Jammy2211 Jammy2211 added the pending-release PR queued for the next release build label Sep 8, 2026
@Jammy2211
Jammy2211 merged commit 35aa681 into main Sep 8, 2026
3 checks passed
@Jammy2211
Jammy2211 deleted the feature/interferometer-numba-cpu-direct-conv branch September 8, 2026 19:03
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

pending-release PR queued for the next release build

Projects

None yet

Development

Successfully merging this pull request may close these issues.

feat: numba CPU interferometer curvature path — direct_conv, geometry-gated

1 participant