perf: build nufft_precision_operator_from as a type-1 NUFFT

[Jammy2211]

Stacked on #540; merge #540 first, GitHub retargets this PR to main.

Summary

nufft_precision_operator_from builds the interferometer sparse-operator preload — the compact (2Ny, 2Nx) array that depends only on the (dy, dx) offset between image pixels. Until now it was built by brute force: an O(N_pix · K) cosine accumulation over every image pixel × visibility pair, which is minutes to hours at real dataset sizes and is why the workspace ships a separate "prepare the preload overnight and cache it to .npy" example.

The same array is the real part of a type-1 (adjoint) NUFFT of the inverse-variance weights w = 1/σ² evaluated on the doubled-extent grid:

P = Re nufft2d1(-x, y, w, (2Nx, 2Ny), eps, +1)   with x = 2π·u·Δ, y = 2π·v·Δ
    → ifftshift → zero the Nyquist row Ny and column Nx → contiguous float64

which costs O(N_vis · nspread² + M log M) for M = 4·Ny·Nx. The sign convention (-x, +y, isign +1) is the one of eight candidate mappings that reproduces the brute force; the derivation and the search are recorded in the function's docstring.

method="nufft" is the new default. The numpy and JAX brute-force builders are untouched and stay as the reference implementations the NUFFT is pinned against.

The array is bit-for-bit the same object (to the pins below), so existing workspace .npy preload caches remain valid — nothing needs regenerating.

Measured (autolens_profiling#229, phase 3)

dataset	brute force	type-1 NUFFT	speed-up
alma	2101 s	7.3 s	289× wall (111× CPU-seconds)
alma_high	—	22 s	—

CPU-seconds caveat: the 289× is wall-clock. The brute-force builder is single-threaded, the NUFFT spreader is multi-threaded, so in CPU-seconds the honest figure is 111×; the wall-clock number is what a user experiences on a workstation.

Pins

• Mixed-tolerance parity against _via_np_from, seeded 16×16 mask at K=300: max deviation 4.4e-14 of peak (rtol=1e-10, atol=1e-10·P[0,0] — mixed because the off-peak entries are many orders of magnitude below the peak, so a pure rtol would pin round-off).
• Chunked (chunk_size=64, K > 64) == one-shot: 5.3e-14.
• Wrong-sign control (uv[:, 0] negated) fails the pin at 0.187 of peak — i.e. the pin has real discriminating power.
• Padding row Ny / column Nx exactly zero, neighbours non-zero; P[i, j] == P[-i, -j].
• End-to-end: the existing test_interferometer.py sparse-vs-mapping comparisons run through the new default unchanged.

eps and chunking

eps = 1e-12 saturates fp64 at every instrument profiled — the measured max|Δ| is already at the round-off floor, and it is the only value that holds the array pin. chunk_size is not an optimisation, it is a memory ceiling: the spreader's gather buffer is one-shot ~15 GB at K = 5e6, so chunking is mandatory at real scale. When not given it is taken from the transformer's own chunk_size if it is a TransformerNUFFT.

Fallback semantics

Two loud, logged fallbacks to method="numpy" — never silent:

• PYAUTO_DISABLE_JAX=1 (the test-mode kill switch already honoured in dataset.py), or
• nufftax is not importable.

Both emit a logger.warning naming the O(N_pix · K) cost being paid. Any other unusable method raises.

API Changes

Additive only — every existing keyword is kept with its existing meaning, so no caller breaks.

• nufft_precision_operator_from gains method="nufft" (new default), eps=1e-12 and chunk_size=None. use_jax is now only honoured when a brute-force method is selected (method="numpy" + use_jax=True = the JAX brute force); it is ignored under the default method="nufft", which already runs on JAX. method="jax" is the explicit way to ask for the JAX brute force. chunk_k still chunks the brute-force builders.
• Interferometer.apply_sparse_operator and Interferometer.psf_precision_operator_from gain method="nufft", eps=None and nufft_chunk_size=None (None = take the transformer's).
• New public builder nufft_precision_operator_via_nufft_from(...).
• Behaviour change: callers get the NUFFT builder instead of the numpy brute force — same array, ~100–300× faster — whether or not they pass use_jax=True. Existing workspace apply_sparse_operator(use_jax=True) calls therefore land on the fast path with no edit; opting back into a brute force now takes an explicit method="numpy" or method="jax".

See full details below.

Test Plan

• pytest test_autoarray/inversion test_autoarray/dataset test_autoarray/operators -q → 685 passed
• python -m black --check on the four changed files → clean
• /prm: merge perf: apply the interferometer sparse operator with rfft2/irfft2 #540 first, then confirm this PR retargeted to main and CI is green per leg

Full API Changes (for automation & release notes)

Added

• autoarray.inversion.inversion.interferometer.inversion_interferometer_util.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 — builds the preload as a type-1 NUFFT.

Changed Signature

• nufft_precision_operator_from(..., *, method="nufft", eps=1e-12, chunk_size=None, chunk_k=2048, show_progress=False, show_memory=False, use_jax=False) — method, eps, chunk_size added; all pre-existing keywords unchanged.
• Interferometer.apply_sparse_operator(nufft_precision_operator=None, batch_size=128, method="nufft", eps=None, nufft_chunk_size=None, chunk_k=2048, show_progress=False, show_memory=False, use_jax=False)
• Interferometer.psf_precision_operator_from(..., method="nufft", eps=None, nufft_chunk_size=None, ...)

Changed Behaviour

• The default preload builder is now the type-1 NUFFT rather than the numpy brute force. The array is identical to within 4.4e-14 of peak; cached .npy preloads stay valid.
• The "N_vis · N_pix ≳ 1e7 crossover" warning in dataset.py was rewritten: the preload is now seconds, so the warning is about the DFT transformer, not the preload.

Migration

• None required. To opt back into the reference builder: method="numpy" (or method="jax" for the JAX one).
• Workspace scripts that pass use_jax=True get the NUFFT path unchanged — the kwarg is now a no-op under the default method, so no workspace edit is needed to buy the speed-up.

Closes #539

Generated by the PyAutoLabs agent workflow.

🤖 Generated with Claude Code

https://claude.ai/code/session_018hLF3ZAcz5MmaSJBEcLkvF

[Jammy2211]

Workspace PR: PyAutoLabs/autogalaxy_workspace#238 — corrects the interferometer preload prose in autogalaxy_workspace ("minutes to hours" -> built as a type-1 NUFFT in seconds). It waits on the next autoarray release per the pending-release chain.