[Jammy2211]

Summary

TransformerNUFFT and the NUFFT precision-operator builder now use finufft (CPU C++ NUFFT, NumPy API) on the NumPy path and keep nufftax for the JAX path. A use_jax=False interferometer run no longer imports nufftax or starts JAX at all.

Why: every eager NumPy NUFFT (the dirty image in apply_sparse_operator, the precision operator, visibilities_from for fits and visualisation) went through nufftax, so a NumPy/numba run silently started JAX and paid JAX's eager cost. On a collaborator's ALMA dataset (5.5M visibilities, 125×125 mask, eps 1e-12; RAL EPYC 7763, 8 threads):

Step	nufftax (before)	finufft (after)	Agreement
Precision operator	8.7 s	0.37 s	1.5e-14 rel
Dirty image	9.0 s	0.27 s	2.0e-14 rel
Forward visibilities_from	10.4 s	0.26 s	4.7e-13 rel

End to end on RAL with this branch, that collaborator's preprocessing took 2.4 s (was ~20 s and a ~44 GB unchunked peak).

Note: an earlier draft of #654 attributed a ≥16-worker Nautilus segfault to JAX. That was wrong: it persists with JAX blocked from import. The real cause (numba first imported inside forked workers) is fixed in the companion PyAutoFit PR. This PR stands on speed, memory and JAX-free NumPy runs.

Closes #654 (together with the PyAutoFit PR).

API Changes

• New core dependency: finufft>=2.4.0 (platform-marked to wheels: Linux/Windows x86_64, macOS arm64; Intel macOS capped <2.5). Elsewhere (e.g. Linux aarch64) the NumPy path falls back to nufftax with a one-time warning.
• psf_precision_operator_from(use_jax=...) / nufft_precision_operator_from(method="nufft", use_jax=...): use_jax now selects the backend (False = finufft default, True = nufftax). It was documented as ignored.
• PYAUTO_DISABLE_JAX no longer demotes the default NUFFT builder to the brute force (finufft is not JAX).
• xp=jnp paths are unchanged (jit / vmap / grad keep using nufftax).
  See full details below.

Test Plan

• pytest test_autoarray: 2191 passed, 4 xfailed (also with OMP_NUM_THREADS unset, i.e. multithreaded)
• Smoke (pyauto-heart smoke, worktree root): 165/165 pass across autofit, autogalaxy, autolens, autolens_workspace_test, euclid pipeline, HowToLens; autocti envs failed arcticpy install before running (unrelated)
• Emulated no-JAX CI leg (jax family blocked via a meta-path finder): 1810 passed, 362 skipped, 0 failed
• finufft vs nufftax parity for image_from / visibilities_from / transform_mapping_matrix (chunked/unchunked, ± primary beam); operator vs NumPy brute force and vs use_jax=True
• Subprocess guard: the whole NumPy route (operator, apply_sparse_operator, forward, adjoint, mapping-matrix transform) never imports nufftax or starts XLA
• libgomp fork regression test (see Decisions)
• RAL end to end on real ALMA data (source_pix[0] → source_pix[1] 28 cores → mass[1])

Full API Changes (for automation & release notes)

Added

• autoarray.operators.nufft_cpu — nufft2d1(x, y, c, n_modes, eps, isign), nufft2d2(x, y, f, eps, isign) (nufftax signature and orientation, finufft-backed, nufftax fallback), backend_available(), nufft_cpu_exception().
• _type1_real_grid_from(..., use_jax=False) backend switch (internal).

Changed

• TransformerNUFFT NumPy branches (visibilities_from, image_from, transform_mapping_matrix) use finufft; nufftax loads only on JAX branches. __init__ no longer imports nufftax.
• Interferometer.psf_precision_operator_from / apply_sparse_operator and nufft_precision_operator_from(method="nufft"): use_jax selects finufft vs nufftax. The dirty image in apply_sparse_operator always uses the NumPy path.
• pyproject.toml: finufft in base dependencies and dev.

Decisions taken (please review)

• Type-1 (nufft2d1) runs on one thread so the dirty image and precision operator are bitwise reproducible: a multithreaded type-1 spread sums thread subgrids in a non-deterministic order (13 bitwise-equality tests failed with OMP threads > 1; first CI run red on py3.13). Cost: ~0.9 s instead of ~0.3 s once per job on 5.5M visibilities (still ~10x faster than nufftax). The likelihood never calls type-1; type-2 (visibilities_from, deterministic) keeps finufft's threads.
• libgomp fork pin: finufft wheels bundle libgomp; a forked child that runs a multithreaded finufft call after its parent did deadlocks. nufft_cpu records the first process that ran finufft and passes nthreads=1 in processes forked from it. Users pinning OMP_NUM_THREADS=1 are unaffected.
• Platform markers + fallback instead of an unconditional dependency, so pip install autoarray never tries to build finufft from source.
• _sparse_parity_setup test now uses method="numpy": its curvature matrix (cond ~4e9) amplifies the on-spec 4e-13 NUFFT error past the test's 1e-4 tolerance; on realistic data both NUFFTs match the brute force to ~2e-14.

Migration

• None for users. Scripts that passed use_jax=True to the operator builder keep nufftax; dropping it gives the faster finufft build.

Generated by the PyAutoLabs agent workflow.

🤖 Generated with Claude Code

https://claude.ai/code/session_013C7Zd6ibK9J4V8CFEUzrJU