perf: build nufft_precision_operator_from as a type-1 NUFFT

[Jammy2211]

Overview

nufft_precision_operator_from builds the W~ preload by brute force — one cosine per (real-space offset, visibility) pair over the mask's bounding extent, O(N_pix·K) — and every w-tilde arm pays it once per dataset before a single likelihood is evaluated. That array is exactly Re of a type-1 (adjoint) NUFFT of the weights 1/σ² on the doubled extent grid, and should be built as one. Measured 289× on wall clock at alma (2101.5 s → 7.278 s) and 111× on CPU-seconds, with the log evidence bit-identical to the brute force at eps=1e-12. It removes a scaling wall rather than a constant, and it costs the library nothing new — nufftax is already a hard dependency of TransformerNUFFT.

Follow-up from autolens_profiling #229 (phase 3 of the retired epic numba-interferometer-revisit, merged as autolens_profiling PR #234; ledger PyAutoMind/complete/archive/epics/numba_interferometer_likelihood_revisit.md). Working: autolens_profiling/results/notes/numba_interferometer_verdict.md section 7 ("The preload") plus results/breakdown/interferometer/preload_breakdown_{sma,alma,alma_high}_v2026.8.17.1.json.

Measured saving

Wall seconds (CPU-seconds in brackets), single-threaded pinning where it takes:

instrument	K	N_pix	current numpy builder	type-1 NUFFT (eps=1e-12)
sma	190	3 852	0.1643 s (0.164)	0.0210 s (0.031)
alma	1 000 000	15 380	2101.5 s (2097.1)	7.278 s (18.84)
alma_high	5 000 000	61 572	refused (hours)	22.15 s (89.99)

The wall-clock ratio is not a single-core ratio: XLA's CPU runtime keeps its own intra-op pool that --xla_cpu_multi_thread_eigen=false does not close, so the NUFFT builder ran at 2.59× cores at alma and 4.06× at alma_high. Quote the CPU-seconds column (111×) before claiming a speed-up over a single-threaded loop.

Stacked on

This task is stacked on perf: apply the interferometer sparse operator with rfft2/irfft2 — #538, branch feature/interferometer-apply-operator-rfft2. Both tasks edit inversion_interferometer_util.py, in disjoint functions; to avoid two branches racing on one file, feature/interferometer-preload-nufft-type1 is cut from branch 538 and its PR is opened with base feature/interferometer-apply-operator-rfft2 (GitHub retargets it to main when #538 merges). /prm order is #538 first, then this one — this task's /prm must not run before #538 is merged.

Plan

• Add a type-1 NUFFT builder nufft_precision_operator_via_nufft_from, ported from the pinned reference implementation in the profiling pack, with its derivation and the eight-mapping search condensed into the docstring.
• Give the dispatcher nufft_precision_operator_from a method selector ("nufft" | "numpy" | "jax", default "nufft") that passes eps and chunk_size through; the two brute-force builders stay untouched as the test references.
• Fall back to "numpy" loudly (never silently) in exactly two cases: JAX disabled via the existing kill switch, and nufftax not importable. Anything else raises.
• Plumb method / eps / nufft_chunk_size through the interferometer dataset, defaulting to the transformer's own settings when it is a TransformerNUFFT, and rewrite the stale "N_vis·N_pix crossover" warning — the preload is now seconds, so the warning is about the DFT transformer, not the preload.
• Pin the mapping structurally, not just numerically: the parity rule, the exactly-zero padding row/column with non-zero neighbours, chunked == one-shot, and a control in which a wrong index permutation must fail the pin. Six of the eight candidate mappings are wrong by 21 % of the peak while being structurally plausible.
• Record in the docstring why eps=1e-12 (it is the only value that also holds the array pin, and it saturates fp64), why chunking is mandatory (~15 GB one-shot at K=5e6), and the CPU-seconds caveat above.

Detailed implementation plan

Work Classification

Library

Affected Repositories

• PyAutoArray (primary)

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoArray	main (47a00e8c)	clean

Suggested branch: feature/interferometer-preload-nufft-type1 (cut from feature/interferometer-apply-operator-rfft2, PR base the same until #538 merges)

Worktree root: ~/Code/PyAutoLabs-wt/interferometer-preload-nufft-type1/

Implementation Steps

1. New nufft_precision_operator_via_nufft_from(noise_map_real, uv_wavelengths, shape_masked_pixels_2d, grid_radians_2d, *, eps=1e-12, chunk_size=None) -> np.ndarray in autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py, ported from autolens_profiling/scripts/misc/numba_interferometer/preload.py::nufft_preload_from: Δ_rad from the grid (square-pixel guard — port _pixel_scale_radians_from), x = 2π·u·Δ, y = 2π·v·Δ, w = σ⁻², chunked nufftax.nufft2d1(-x, y, w, (2Nx, 2Ny), eps, +1) accumulating Re f in float64, ifftshift, zero row Ny and column Nx, contiguous float64. Do not route this through image_from — that path applies a half-pixel shift and a row flip, and neither belongs here.
2. Load nufftax the way transformer.py does — reuse _load_nufftax() / nufftax_exception() by import, do not duplicate them.
3. Dispatcher nufft_precision_operator_from (line 80) gains method: str = "nufft" with values "nufft" | "numpy" | "jax", passing eps and chunk_size through. use_jax=True (existing kwarg, kept for compatibility) maps to method="jax". nufft_precision_operator_via_np_from / _via_jax_from stay untouched as reference builders.
4. Two loud, logged fallbacks to "numpy", never silent: disable_jax() is set (PYAUTO_DISABLE_JAX=1, the test-mode kill switch already honoured at dataset.py:318), or nufftax is not importable (logger.warning naming the O(N_pix·K) cost). Anything else raises.
5. autoarray/dataset/interferometer/dataset.py: psf_precision_operator_from (~line 398) and apply_sparse_operator (~line 229) accept method="nufft", eps=None, nufft_chunk_size=None; when None they take the transformer's own eps / chunk_size if it is a TransformerNUFFT, else 1e-12 / None.
6. Rewrite the "N_vis·N_pix ≳ 1e7 crossover" warning text at dataset.py:351-376: the preload is now seconds; the warning is about the DFT transformer, not the preload.

Tests

test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py and test_autoarray/dataset/interferometer/test_dataset.py, porting the pins from the profiling pack's test_parity.py verbatim in spirit:

1. NUFFT vs _via_np_from on the shared 7×7 / K=5 fixture and a seeded 16×16 mask with K=300 (so the pin is not trivial), mixed tolerance assert_allclose(rtol=1e-10, atol=1e-10 * P[0, 0]), with a comment saying why mixed — a type-1 NUFFT bounds its error against Σ_k|c_k|, i.e. against the peak, so the preload's near-zero entries carry no relative guarantee.
2. Padding row Ny / column Nx exactly zero, and their neighbours non-zero (so the pin is not vacuous).
3. P[i, j] == P[-i, -j].
4. Chunked (chunk_size=64, K > 64) equals one-shot.
5. Control: negating uv[:, 0] must fail the pin (pytest.raises(AssertionError)).
6. Dispatcher: default returns the NUFFT array; method="numpy" returns the brute force; use_jax=True still routes to the JAX brute force; PYAUTO_DISABLE_JAX=1 falls back to numpy (extend the existing _disable_jax_overrides_an_explicit_use_jax test).
7. End-to-end: the existing test_interferometer.py sparse-vs-mapping comparisons run through the new default unchanged.

All under CPU JAX x64 (already enabled by autonerves on import); nufftax-gated tests use the conftest skip (pytest_collection_modifyitems, keep __dft__ in names where the fixture is DFT-based).

Verification

• Pins 1–6 pass; test_dataset.py green.
• pytest test_autoarray/inversion/inversion/interferometer test_autoarray/dataset/interferometer test_autoarray/operators -q green, plus the full test_autoarray/inversion tree once.
• The alma-scale timing is not re-measured here (no such dataset in PyAutoArray); the PR cites the phase-3 JSON numbers above.

Key Files

• autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py — dispatcher nufft_precision_operator_from (line 80), reference builders _via_np_from (249) / _via_jax_from (368), new _via_nufft_from.
• autoarray/dataset/interferometer/dataset.py — apply_sparse_operator (~229), psf_precision_operator_from (~398), crossover warning (~351-376), JAX kill switch (~318).
• autoarray/operators/transformer.py — _load_nufftax() / nufftax_exception() and the _x / _y convention (~334-335).
• autolens_profiling/scripts/misc/numba_interferometer/preload.py — nufft_preload_from, the pinned reference implementation.

Related

• Epic ledger (retired 2026-09-08): PyAutoMind/complete/archive/epics/numba_interferometer_likelihood_revisit.md
• Phase 3: research: interferometer preload on CPU — builder timings + type-1 NUFFT construction autolens_profiling#229, merged via research: interferometer preload on CPU — builder timings + type-1 NUFFT construction autolens_profiling#234
• Stacked on: perf: apply the interferometer sparse operator with rfft2/irfft2 #538
• Chunking precedent: PyAutoArray#330

Original Prompt

Click to expand starting prompt

Build nufft_precision_operator_from as a type-1 NUFFT — 35 minutes to 7 seconds

Type: feature
Target: autoarray
Repos:

• PyAutoArray
  Themes:
• interferometer
• nufft
• likelihood-profiling
  Difficulty: medium
  Autonomy: supervised
  Priority: high
  Epic: numba-interferometer-revisit
  Filed: 2026-09-08

Follow-up derived from autolens_profiling#229 (phase 3 of numba-interferometer-revisit)
— not a user request. Every number below is measured, and the working is in
autolens_profiling/results/notes/numba_interferometer_verdict.md section 7 ("The
preload") plus
results/breakdown/interferometer/preload_breakdown_{sma,alma,alma_high}_v2026.8.17.1.json.

What

nufft_precision_operator_from in
autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py builds the
W~ preload by brute force: one cosine per (real-space offset, visibility) pair over the
mask's bounding extent, O(N_pix·K). Every w-tilde arm pays it once per dataset before a
single likelihood is evaluated.

That array is exactly Re of a type-1 (adjoint) NUFFT of the weights 1/σ² on the
doubled extent grid, and should be built as one:

x, y = 2π·u·Δ_rad, 2π·v·Δ_rad          # the transformer's own _x/_y (transformer.py:334-335)
f = nufftax.nufft2d1(−x, y, w, n_modes=(2Nx, 2Ny), eps, isign=+1)
P = ifftshift(Re f)
P[Ny, :] = 0;  P[:, Nx] = 0            # the FFT padding row/column the brute force never fills

Do not route this through image_from — that path applies a half-pixel shift and a row
flip, and neither belongs here.

Reference implementation, already written and pinned:
autolens_profiling/scripts/misc/numba_interferometer/preload.py::nufft_preload_from.

Why

Wall seconds (CPU-seconds in brackets), single-threaded pinning where it takes:

instrument	K	N_pix	current numpy builder	type-1 NUFFT (eps=1e-12)
sma	190	3 852	0.1643 s (0.164)	0.0210 s (0.031)
alma	1 000 000	15 380	2101.5 s (2097.1)	7.278 s (18.84)
alma_high	5 000 000	61 572	refused (hours)	22.15 s (89.99)

289× on wall clock at alma, 111× on CPU-seconds. It removes a scaling wall rather than a
constant — the brute force is linear in K, the NUFFT is K·nspread² + M log M — so the gap
widens with instrument size (predicted 1590× at alma_high). And it is not an approximation:
at eps=1e-12 the log evidence at alma is bit-identical to the brute force's
(-12050103.936303042, rel diff 0).

It costs the library nothing new: nufftax is already a hard dependency of
TransformerNUFFT, the mapping uses that transformer's own convention, and the same call
runs on GPU.

Requirements

• Keep nufft_precision_operator_via_np_from as the test reference. Pin the new path
  against it at a small K with the mixed tolerance
  allclose(rtol=1e-10, atol=1e-10·P[0,0]) — a type-1 NUFFT bounds its error against
  Σ_k|c_k|, i.e. against the peak, so the preload's near-zero entries (five orders below
  it, at the fp64 noise floor) carry no relative guarantee and a relative-only test there
  measures round-off, not the builder.
• Carry the structural pins, not just a smoke test. The exactly-zero padding row/column
  (with non-zero neighbours, so the pin is not vacuous), P[i,j] == P[−i,−j], and a control
  in which a wrong index permutation must fail the parity rule. Six of the eight candidate
  mappings (axis swap × sign of x × sign of y) are wrong by 21 % of the peak while being
  structurally plausible; the mapping is convention-bound and must be pinned as such.
• Default eps=1e-12. Every value tested (1e-6, 1e-9, 1e-12) holds the rtol=1e-6
  log-evidence pin by at least five orders of margin, so the dial is not load-bearing for
  correctness — but 1e-12 is the only one that also holds the array pin at sma, it
  saturates fp64 (1e-14 moves max|Δ| from 1.7e-17 only to 1.4e-17), and it costs
  seconds. A preload is reused by every likelihood call in a fit; error budget spent here is
  spent for the whole run.
• Visibility chunking is mandatory, not optional. The spreader's gather buffer is
  K·nspread² complex128; alma_high's 5 M visibilities need ~15 GB in one shot and were
  OOM-killed in the profiling run before chunking was added. Chunk at the transformer's
  existing chunk_size (PyAutoArray#330) and pin chunked == one-shot.
• This applies to every backend. The JAX chunked builder becomes redundant, and the
  dirty image already goes through the same adjoint — so the NUFFT path is the one builder,
  not a GPU-only alternative. It also gives the GPU path a preload it does not currently
  have.

Caveat to record

The wall-clock ratio is not a single-core ratio. XLA's CPU runtime keeps its own intra-op
pool that --xla_cpu_multi_thread_eigen=false does not close, so the NUFFT builder ran at
2.59× cores at alma and 4.06× at alma_high. Quote the CPU-seconds column (111×) before
claiming a speed-up over a single-threaded loop.

Side finding (not this task, but worth knowing)

The recovered numba brute-force builder is the only one of four that fails the alma array
pin, at 2.195e-11 of the peak against a 1e-11 bound. It is summation round-off, not an
algebraic disagreement: it accumulates all 10⁶ visibilities into one running scalar per
offset while the NumPy reference sums them in chunk_k = 2048 blocks, so the two diverge
with K (2e-15 at K = 190, 2e-11 at K = 10⁶, on 11 of 78 400 entries). The NUFFT
agrees with the reference a thousand times better than that numba builder does
(1.9e-14 vs 2.2e-11). A 10⁶-term naive accumulation is worth knowing about wherever
else one lives.

[Jammy2211]

Library PR Created

• PR: perf: build nufft_precision_operator_from as a type-1 NUFFT #541 — stacked on perf: apply the interferometer sparse operator with rfft2/irfft2 #540 (base feature/interferometer-apply-operator-rfft2; GitHub retargets it to main when perf: apply the interferometer sparse operator with rfft2/irfft2 #540 merges, so /prm must run perf: apply the interferometer sparse operator with rfft2/irfft2 #540 first).
• Commit: 87205219 on feature/interferometer-preload-nufft-type1.
• Tests: pytest test_autoarray/inversion test_autoarray/dataset test_autoarray/operators -q → 685 passed, 0 failed (114 s). black --check clean on all four changed files.
• Heart: YELLOW (score 70), two reasons, both organism-scope — "workspace validation not passing (cloud#34099198772)" and "release validation incomplete: no rehearsal for current source". Neither names PyAutoArray or a library test.

Downstream workspace impact — (iii) no migration needed

Every pre-existing keyword is kept (chunk_k, show_progress, show_memory, use_jax), so no workspace call breaks. Grep across autofit_workspace, autogalaxy_workspace, autolens_workspace, *_workspace_test, autolens_workspace_developer, HowTo*, autocti_workspace, euclid_strong_lens_modeling_pipeline for psf_precision_operator_from / apply_sparse_operator( / nufft_precision_operator_from(:

• autogalaxy_workspace/scripts/interferometer/features/pixelization/many_visibilities_preparation.py — apply_sparse_operator(use_jax=True, chunk_k=2048, ...) and psf_precision_operator_from(use_jax=True, chunk_k=2048, ...)
• autogalaxy_workspace/scripts/interferometer/features/pixelization/{fit,modeling}.py — apply_sparse_operator(use_jax=True, show_progress=True)
• prose-only mentions in interferometer/features/{operated_light_profile,extra_galaxies,linear_light_profiles}/modeling.py (plus the mirrored notebooks)

All still run, unchanged. Note that use_jax=True routes to the JAX brute-force builder, so those scripts keep the old runtime until the kwarg is dropped.

Workspace follow-up (not done here — no workspace edits in this PR)

autogalaxy_workspace/scripts/interferometer/features/pixelization/many_visibilities_preparation.py (and its notebook) is the script whose whole premise — "the nufft_precision_operator computation takes minutes to hours, so prepare it separately and cache it to .npy" — is now stale: with the default method="nufft" it is seconds. Its prose and its use_jax=True calls should be updated in a workspace task. The cached .npy files themselves stay valid: the array is unchanged to 4.4e-14 of peak.

Next: /prm — merge #540, then this PR.

[Jammy2211]

Follow-up pushed to #541 (61f07fd): use_jax is now only honoured when a brute-force method is selected, so the existing workspace apply_sparse_operator(use_jax=True) calls land on the type-1 NUFFT rather than the JAX brute force — no workspace edit needed; method="jax" remains the explicit way to request the brute force.

[Jammy2211]

Shipped

PR #541 — MERGED, merge commit 9bd76799, head 61f07fd8. Stacked on #540 (merged first, 7a4cb700); GitHub retargeted this PR to main.

nufft_precision_operator_from now builds the interferometer sparse-operator preload as a type-1 (adjoint) NUFFT of the inverse-variance weights on the doubled-extent grid, with method="nufft" the new default and eps=1e-12. The numpy and JAX brute-force builders are untouched and stay as the reference implementations the NUFFT is pinned against; chunk_k still chunks those. Visibility chunking for the NUFFT is taken from the transformer's own chunk_size when not given — a memory ceiling, not an optimisation (the spreader's gather buffer is ~15 GB one-shot at K = 5e6).

A second commit (61f07fd8) removed the last demotion path: use_jax=True no longer routes a caller back to the brute force. Existing workspace calls that pass use_jax=True land on the NUFFT path with no edit; opting back into a reference builder now takes an explicit method="numpy" or method="jax".

Pins

• Parity against nufft_precision_operator_via_np_from, mixed tolerance: max deviation 4.4e-14 of peak.
• Chunked (chunk_size=64) == one-shot: 5.3e-14.
• Wrong-sign control (uv[:, 0] negated) fails the pin at 0.187 of peak — the pin discriminates.
• Padding row Ny / column Nx exactly zero with non-zero neighbours; P[i, j] == P[-i, -j].

Measured

alma: 2101 s -> 7.3 s wall (289x), 111x in CPU-seconds (the brute force is single-threaded, the spreader is not). alma_high: 22 s, where the brute force refused.

CI green on head 61f07fd8: unittest (3.12), unittest (3.13), unittest-nojax.

pending-release: PyAutoArray@https://github.com/PyAutoLabs/PyAutoArray/pull/541 — merged but not released; the change reaches users at the next autoarray release.