From 157e210c030769df5cd149804a76fbefe0b537d6 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Sat, 24 Oct 2026 16:31:20 +0100 Subject: [PATCH 1/2] feat: finufft as the CPU NUFFT backend; nufftax only on the JAX path (#654) TransformerNUFFT (forward, adjoint, mapping-matrix transform) and the NUFFT precision-operator / sparse-terms builders use finufft on the NumPy path, so a use_jax=False interferometer run never imports nufftax or starts JAX. use_jax now selects the operator builder backend. finufft is a platform-marked core dependency with a nufftax fallback; forked children are pinned to one finufft thread to avoid a libgomp fork deadlock. 25-35x faster than nufftax on CPU, agreeing to ~1e-14. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_013C7Zd6ibK9J4V8CFEUzrJU --- autoarray/dataset/interferometer/dataset.py | 33 +- .../inversion_interferometer_util.py | 115 +++-- autoarray/operators/nufft_cpu.py | 250 +++++++++++ autoarray/operators/transformer.py | 88 ++-- pyproject.toml | 18 +- test_autoarray/conftest.py | 27 +- .../dataset/interferometer/test_dataset.py | 28 +- .../interferometer/test_interferometer.py | 11 +- .../test_inversion_interferometer_util.py | 107 ++++- test_autoarray/operators/test_nufft_cpu.py | 406 ++++++++++++++++++ test_autoarray/operators/test_transformer.py | 9 +- 11 files changed, 965 insertions(+), 127 deletions(-) create mode 100644 autoarray/operators/nufft_cpu.py create mode 100644 test_autoarray/operators/test_nufft_cpu.py diff --git a/autoarray/dataset/interferometer/dataset.py b/autoarray/dataset/interferometer/dataset.py index 287d7b3d8..133e8ad34 100644 --- a/autoarray/dataset/interferometer/dataset.py +++ b/autoarray/dataset/interferometer/dataset.py @@ -551,8 +551,11 @@ def apply_sparse_operator( the inversion module. The backend the operator is *applied* with follows the inversion's `xp` (NumPy/scipy for - `xp=np`, JAX for `xp=jnp`); `use_jax` below only selects which brute-force builder computes - the preload, and never puts JAX on a NumPy fit's application path. + `xp=np`, JAX for `xp=jnp`); `use_jax` below only selects the backend that *builds* the + preload, and never puts JAX on a NumPy fit's application path. With the default + `use_jax=False` the whole call -- precision operator and dirty image, both on `finufft` -- + never imports JAX, so a NumPy fit does not start XLA before a multiprocessing search forks + its workers. The default builder (`method="nufft"`) computes the precision operator as a type-1 NUFFT, so it costs `O(N_vis * nspread^2 + M log M)` for `M = 4 * Ny * Nx` — seconds even at a million @@ -589,11 +592,11 @@ def apply_sparse_operator( show_memory If `True`, memory usage statistics are printed while computing the NUFFT precision matrix. use_jax - Only honoured when a brute-force builder is selected: `method="numpy"` with - `use_jax=True` runs the JAX brute force (equivalent to `method="jax"`). Under the - default `method="nufft"` it is ignored, because the NUFFT already runs on JAX -- - so an existing `use_jax=True` call keeps the fast path rather than being demoted - to the `O(N_vis * N_pix)` brute force. + Selects the backend that builds the precision operator. Under the default + `method="nufft"`, `False` runs the type-1 NUFFT on `finufft` (CPU, no JAX) and + `True` on `nufftax` (JAX); both are fast paths. Under `method="numpy"`, `True` runs + the JAX brute force (equivalent to `method="jax"`). The dirty image is always formed + on the NumPy path (`finufft`), whatever this says. `PYAUTO_DISABLE_JAX=1` overrides this to `False`. That variable is a harness-level switch, not a preference: it is the documented way to force the @@ -602,7 +605,8 @@ def apply_sparse_operator( explicit `use_jax=True` in a script -- which is the right thing for a script demonstrating the production path to say -- must therefore not defeat it, or the harness pays 2.3-3.2 s of compile for a backend it asked to disable. (The - same switch demotes the `"nufft"` builder to the NumPy brute force inside + `"nufft"` builder then runs on `finufft`, which is not JAX; only if `finufft` is + missing does the switch demote it to the NumPy brute force inside `nufft_precision_operator_from`.) pool_noise_map If `True`, pool the real and imaginary noise sigma of every visibility in @@ -653,9 +657,9 @@ def apply_sparse_operator( is `False`. """ - # `use_jax` now only selects between the two brute forces (the `"nufft"` builder - # runs on JAX whatever it says), so this clears it before it can upgrade - # `method="numpy"` to the JAX brute force under the kill switch. + # `use_jax` selects the builder's backend (`finufft` vs `nufftax` under `"nufft"`, the + # NumPy vs JAX brute force under `"numpy"`), so this clears it before it can put the + # build on JAX under the kill switch. if disable_jax(): use_jax = False @@ -992,9 +996,10 @@ def psf_precision_operator_from( show_memory If `True`, memory usage statistics are printed during computation. use_jax - Only honoured when a brute-force builder is selected: `method="numpy"` with - `use_jax=True` runs the JAX brute force (equivalent to `method="jax"`). It is - ignored under the default `method="nufft"`, which already runs on JAX. + Selects the backend that builds the operator. Under the default `method="nufft"`, + `False` runs the type-1 NUFFT on `finufft` (CPU, never imports JAX) and `True` on + `nufftax` (JAX). Under `method="numpy"`, `True` runs the JAX brute force + (equivalent to `method="jax"`). method Which builder computes the operator: `"nufft"` (default), `"numpy"` or `"jax"`. eps diff --git a/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py b/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py index 6539c1c89..a848e98aa 100644 --- a/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py +++ b/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py @@ -10,6 +10,7 @@ from autoarray import exc +from autoarray.operators import nufft_cpu from autoarray.operators.transformer import _load_nufftax, nufftax_exception try: @@ -474,25 +475,31 @@ def nufft_precision_operator_from( - `"nufft"` (default) -- `nufft_precision_operator_via_nufft_from`, the type-1 (adjoint) NUFFT. `O(K*nspread^2 + M log M)` for `M = 4*Ny*Nx`, i.e. seconds - where the brute-force builders take minutes to hours. Needs `nufftax`. + where the brute-force builders take minutes to hours. Runs `finufft` (CPU, no + JAX) by default and `nufftax` (JAX) with `use_jax=True`. - `"numpy"` -- `nufft_precision_operator_via_np_from`, the brute-force `O(N_pix*K)` reference builder. Kept as the reference the NUFFT builder is pinned against, and used as the fallback below. - `"jax"` -- `nufft_precision_operator_via_jax_from`, the same brute force on JAX. `method="jax"` is the explicit way to ask for it. - `use_jax` is only honoured when a brute-force method is selected: it upgrades - `method="numpy"` to `method="jax"` and is otherwise ignored. Under the default - `method="nufft"` it does nothing, because the NUFFT already runs on JAX -- an - existing `use_jax=True` caller therefore keeps the fast path rather than being - demoted to the `O(N_pix*K)` brute force. + `use_jax` picks the backend of the chosen method: under the default + `method="nufft"`, `use_jax=False` runs the type-1 NUFFT on `finufft` (CPU, never + imports JAX, so a NumPy fit does not start XLA before a multiprocessing search + forks its workers) and `use_jax=True` runs it on `nufftax` (JAX); under + `method="numpy"`, `use_jax=True` upgrades to `method="jax"`. Either way an + existing `use_jax=True` caller keeps a NUFFT fast path rather than being demoted + to the `O(N_pix*K)` brute force. Two fallbacks to `"numpy"` are taken, both logged loudly (never silently), because they cost `O(N_pix*K)` where the NUFFT is `O(K*nspread^2 + M log M)`: - 1. `disable_jax()` is true (`PYAUTO_DISABLE_JAX=1`, the test-mode kill switch). - Both `"nufft"` and `"jax"` run on JAX, so both are demoted. - 2. `nufftax` is not importable, so `"nufft"` cannot run. + 1. `disable_jax()` is true (`PYAUTO_DISABLE_JAX=1`, the test-mode kill switch) + and the method would run on JAX: `"jax"`, or `"nufft"` with `use_jax=True` + (or with `finufft` missing, where the NumPy NUFFT falls back to `nufftax`). + `"nufft"` on `finufft` is not JAX, so it is not demoted. + 2. No NUFFT backend for the requested route is importable: `nufftax` for + `use_jax=True`, or neither `finufft` nor `nufftax` for `use_jax=False`. Any other `method` raises `ValueError`. @@ -519,9 +526,9 @@ def nufft_precision_operator_from( chunk_k The visibility chunk size of the two brute-force builders. use_jax - Only honoured when a brute-force method is selected: `method="numpy"` with - `use_jax=True` runs the JAX brute force (equivalent to `method="jax"`). It is - ignored under the default `method="nufft"`, which already runs on JAX. + Under `method="nufft"`, `False` (default) runs the NUFFT on `finufft` and + `True` on `nufftax` (JAX). Under `method="numpy"`, `True` runs the JAX brute + force (equivalent to `method="jax"`). """ if method == "numpy" and use_jax: method = "jax" @@ -533,23 +540,32 @@ def nufft_precision_operator_from( "brute-force reference builders)." ) - if method in ("nufft", "jax") and disable_jax(): + # The NumPy NUFFT runs on JAX only when finufft is missing and it falls back to nufftax. + nufft_on_jax = use_jax or nufft_cpu._load_finufft() is None + + if (method == "jax" or (method == "nufft" and nufft_on_jax)) and disable_jax(): logger.warning( f"INTERFEROMETER - `PYAUTO_DISABLE_JAX=1` is set, so the NUFFT precision " - f"operator cannot be built with method={method!r} (both the NUFFT and the " - f"JAX brute force run on JAX). Falling back to the NumPy brute force, which " + f"operator cannot be built with method={method!r} on JAX (use_jax={use_jax}, " + f"or `finufft` is not installed). Falling back to the NumPy brute force, which " f"is O(N_pix * K) rather than O(K * nspread^2 + M log M) and can take " f"minutes to hours on a real dataset." ) method = "numpy" - if method == "nufft" and _load_nufftax() is None: + # nufftax is only looked for on the JAX route; the NumPy route checks for a backend + # without importing nufftax (which would start JAX). + if method == "nufft" and ( + (use_jax and _load_nufftax() is None) + or (not use_jax and not nufft_cpu.backend_available()) + ): logger.warning( - "INTERFEROMETER - `nufftax` is not installed, so the NUFFT precision " - "operator cannot be built with the type-1 NUFFT. Falling back to the NumPy " - "brute force, which is O(N_pix * K) rather than O(K * nspread^2 + M log M) " - "and can take minutes to hours on a real dataset. Install it via " - "`pip install nufftax`." + f"INTERFEROMETER - No NUFFT backend is installed for use_jax={use_jax} " + f"(`nufftax` for use_jax=True, `finufft` or `nufftax` for use_jax=False), so " + f"the NUFFT precision operator cannot be built with the type-1 NUFFT. Falling " + f"back to the NumPy brute force, which is O(N_pix * K) rather than " + f"O(K * nspread^2 + M log M) and can take minutes to hours on a real dataset. " + f"Install it via `pip install finufft` (or `pip install nufftax` for JAX)." ) method = "numpy" @@ -561,6 +577,7 @@ def nufft_precision_operator_from( grid_radians_2d=grid_radians_2d, eps=eps, chunk_size=chunk_size, + use_jax=use_jax, ) if method == "jax": @@ -666,6 +683,7 @@ def _type1_real_grid_from( centred: bool = False, zero_nyquist: bool = False, out: Optional[np.ndarray] = None, + use_jax: bool = False, ) -> np.ndarray: """ The real part of a **type-1 (adjoint) NUFFT** of per-visibility `values` onto a regular @@ -713,18 +731,23 @@ def _type1_real_grid_from( which is returned; nothing of the grid's size is allocated beyond `nufft2d1`'s own output and its real part per chunk. This is how `sparse_terms_from_chunks` accumulates the fine grids across a stream. If `None` a new array is returned. + use_jax + If `False` (default) the transform runs on `finufft` via `nufft_cpu` (CPU, never imports + JAX); if `True` it runs on `nufftax` (JAX). The two agree to ~1e-14 relative and return + the same NumPy array. Returns ------- np.ndarray `[N_y, N_x]` float64. """ - nufftax = _load_nufftax() + if use_jax: + nufftax = _load_nufftax() - if nufftax is None: - nufftax_exception() + if nufftax is None: + nufftax_exception() - import jax.numpy as jnp + import jax.numpy as jnp uv_wavelengths = np.asarray(uv_wavelengths, dtype=np.float64) @@ -761,14 +784,24 @@ def _type1_real_grid_from( for k0 in range(0, total_visibilities, chunk_size): k1 = min(total_visibilities, k0 + chunk_size) - f = nufftax.nufft2d1( - jnp.asarray(-x[k0:k1]), - jnp.asarray(y[k0:k1]), - jnp.asarray(values[k0:k1], dtype=jnp.complex128), - n_modes, - eps, - 1, - ) + if use_jax: + f = nufftax.nufft2d1( + jnp.asarray(-x[k0:k1]), + jnp.asarray(y[k0:k1]), + jnp.asarray(values[k0:k1], dtype=jnp.complex128), + n_modes, + eps, + 1, + ) + else: + f = nufft_cpu.nufft2d1( + -x[k0:k1], + y[k0:k1], + np.asarray(values[k0:k1], dtype=np.complex128), + n_modes, + eps, + 1, + ) if real_modes is not None: real_modes += np.asarray(np.real(f), dtype=np.float64) @@ -803,6 +836,7 @@ def nufft_precision_operator_via_nufft_from( *, eps: float = 1.0e-12, chunk_size: Optional[int] = None, + use_jax: bool = False, ) -> np.ndarray: """ Returns the `W~` precision operator built as the real part of a **type-1 (adjoint) @@ -833,7 +867,8 @@ def nufft_precision_operator_via_nufft_from( P[i, j] = Re sum_k w_k exp(i(-j * x_k + i * y_k)) i.e. the real part of a type-1 NUFFT of the weights onto the `(2Ny, 2Nx)` mode grid. - `nufftax.nufft2d1(x, y, c, n_modes=(N1, N2), eps, isign)` returns + `nufftax.nufft2d1(x, y, c, n_modes=(N1, N2), eps, isign)` (and `nufft_cpu.nufft2d1`, the + `finufft` wrapper that mirrors it) returns `f[m2, m1] = sum_k c_k exp(isign * i(m1 * x_k + m2 * y_k))` on the **centred** mode grid, shape `(N2, N1)`, so the mapping is @@ -902,6 +937,9 @@ def nufft_precision_operator_via_nufft_from( The requested NUFFT precision. chunk_size Cap on the visibilities passed to `nufft2d1` in one call, or `None` for one shot. + use_jax + If `False` (default) the NUFFT runs on `finufft` (CPU, never imports JAX); if `True` + on `nufftax` (JAX). Both return the same NumPy array to ~1e-14 relative. Returns ------- @@ -927,6 +965,7 @@ def nufft_precision_operator_via_nufft_from( chunk_size=chunk_size, centred=False, zero_nyquist=True, + use_jax=use_jax, ) @@ -2543,7 +2582,7 @@ def fine_grid_shapes_from( def _log_fine_grid_memory(shape_operator_fine, shape_dirty_image_fine, oversample): """ Log one line estimating the memory the `oversample` fine grids need: the two float64 grids - held for the whole stream, plus the largest per-call `nufftax.nufft2d1` working set (its 2x + held for the whole stream, plus the largest per-call `nufft2d1` working set (its 2x upsampled complex128 grid, its complex128 output and that output's float64 real part). """ held = 8 * ( @@ -2665,7 +2704,7 @@ def sparse_terms_from_chunks( grids need the NUFFT: `oversample` with a transformer other than `TransformerNUFFT` raises. Memory and cost. At a 400 x 400 native grid and `q = 8`, `precision_operator_fine` is 8000^2 - float64 (512 MB) and `dirty_image_fine` 6400^2 (328 MB), and `nufftax`'s 2x upsampled complex + float64 (512 MB) and `dirty_image_fine` 6400^2 (328 MB), and the NUFFT's 2x upsampled complex work grid adds ~4 GB during each call: plan for 5-6 GB peak (`q = 4` is ~1.5 GB; ~112 pixels at `q = 8` is under 0.5 GB). One info line logs the estimate for the actual shapes. The FFT part of each fine NUFFT costs the same whatever the chunk length (seconds per chunk at @@ -2709,6 +2748,8 @@ def sparse_terms_from_chunks( method, eps, chunk_size, chunk_k, use_jax, show_progress Passed to `nufft_precision_operator_from` for each chunk (`chunk_size` is the NUFFT builder's own inner chunk, a memory ceiling within a chunk). `eps=None` is `1e-12`. + `use_jax` also picks the backend of the `oversample` fine grids (`finufft` for `False`, + `nufftax` for `True`); the dirty image and beam always use the transformer's NumPy path. phase_centre The `(y, x)` phase-centre shift in arcseconds applied to every chunk's visibilities before forming the dirty image (see "Phase-centre shift" above). `None` applies no @@ -2889,6 +2930,7 @@ def sparse_terms_from_chunks( centred=False, zero_nyquist=False, out=precision_operator_fine, + use_jax=use_jax, ) _type1_real_grid_from( @@ -2904,6 +2946,7 @@ def sparse_terms_from_chunks( centred=True, zero_nyquist=False, out=dirty_image_fine, + use_jax=use_jax, ) dirty_beam_native = np.asarray( diff --git a/autoarray/operators/nufft_cpu.py b/autoarray/operators/nufft_cpu.py new file mode 100644 index 000000000..af5776ccb --- /dev/null +++ b/autoarray/operators/nufft_cpu.py @@ -0,0 +1,250 @@ +""" +The CPU (NumPy) NUFFT backend of `TransformerNUFFT` and the NUFFT precision-operator builders. + +Every eager NumPy NUFFT in PyAutoArray goes through the two functions below, which wrap `finufft` +(https://github.com/flatironinstitute/finufft), the CPU C++ library that `nufftax` ports to JAX. +They mirror `nufftax.nufft2d1` / `nufftax.nufft2d2` exactly -- same arguments, same output +orientation -- so a NumPy branch swaps one call for the other and nothing else changes. + +The point of the split is that the NumPy path never imports JAX. Importing `nufftax` alone +initialises XLA (its spreader builds JAX arrays at import), and a parent process that has started +JAX's threads before a multiprocessing search forks its workers can crash those workers when numba +compiles in them. The array namespace therefore picks the backend: `xp=np` (eager NumPy) comes +here and runs `finufft`; `xp=jnp` / `use_jax=True` stays on `nufftax`, which `jax.jit`, `jax.grad` +and `jax.vmap` need. `finufft` is also 25-35x faster than `nufftax` on CPU, and the two agree to +~1e-14 relative. + +Orientation: `finufft` returns (and takes) the mode grid in `n_modes` order, `(n_modes[0], +n_modes[1])`, with `x` paired to the first axis; `nufftax` uses `(n_modes[1], n_modes[0])` with +`x` paired to the last axis. The wrappers transpose the last two axes to give `nufftax`'s layout. + +`finufft` has no wheel for every platform (notably Linux aarch64). Where it is not importable the +functions fall back to `nufftax` with a one-time warning, since that puts JAX back on the NumPy +path; where neither is importable they raise `nufft_cpu_exception`. +""" + +import importlib.util +import logging +import os +import numpy as np +from typing import Tuple + +logger = logging.getLogger(__name__) + +# finufft is a compiled extension with no JAX dependency, but it is still only imported on first +# use so that sessions which never build an interferometer transformer do not pay for it. +_finufft = None +_finufft_loaded = False +_fallback_warned = False + +# The pid of the process in which finufft first ran, i.e. whose OpenMP (libgomp) thread pool it +# started. The value is inherited by forked children, which is how they recognise themselves. +_openmp_pid = None + + +def _load_finufft(): + global _finufft, _finufft_loaded + if _finufft_loaded: + return _finufft + _finufft_loaded = True + try: + import finufft + except ImportError: + return None + _finufft = finufft + return _finufft + + +def _threads_kwargs() -> dict: + """ + The `nthreads` keyword for a finufft call: none (finufft's default, `OMP_NUM_THREADS` or all + cores) in the process that first ran finufft, `nthreads=1` in a process forked from it. + + GNU OpenMP (libgomp, which the finufft wheels bundle) is not fork-safe: once a parent has run a + multi-threaded OpenMP region, a forked child that opens another multi-threaded region + deadlocks. This is exactly the multiprocessing-search pattern -- the parent builds the + precision operator and dirty image with finufft, then Nautilus forks workers whose likelihood + calls `TransformerNUFFT.visibilities_from` -- so forked children are pinned to one thread + (which is what a worker sharing the machine with its siblings wants anyway). A child whose + parent never ran finufft keeps finufft's default. + """ + global _openmp_pid + + pid = os.getpid() + + if _openmp_pid is None: + _openmp_pid = pid + + if pid == _openmp_pid: + return {} + + return {"nthreads": 1} + + +def backend_available() -> bool: + """ + Whether any NUFFT backend can run the NumPy path: `finufft`, or `nufftax` as the fallback. + + `nufftax` is only located (`importlib.util.find_spec`), never imported, because importing it + starts JAX -- which is exactly what a NumPy user checking for a backend must not do. + """ + if _load_finufft() is not None: + return True + return importlib.util.find_spec("nufftax") is not None + + +def nufft_cpu_exception(): + raise ModuleNotFoundError( + "\n--------------------\n" + "You are attempting to perform interferometer analysis with `TransformerNUFFT` or the " + "NUFFT precision operator.\n\n" + "However, neither of its NUFFT backends is installed: finufft " + "(https://github.com/flatironinstitute/finufft), used on the NumPy path, nor nufftax " + "(https://github.com/GragasLab/nufftax), used on the JAX path.\n\n" + "Install finufft via the command `pip install finufft`.\n\n" + "If your platform has no finufft wheel, install nufftax via the command " + "`pip install nufftax` (it requires JAX), or use " + "`transformer_class=TransformerDFT` instead.\n\n" + "----------------------" + ) + + +def _nufftax_fallback(): + """ + Returns `nufftax` for the NumPy path when `finufft` is not importable, warning once that this + starts JAX, or raises `nufft_cpu_exception` when `nufftax` is not importable either. + """ + global _fallback_warned + + # Imported here: transformer.py imports this module at its top. + from autoarray.operators.transformer import _load_nufftax + + nufftax = _load_nufftax() + + if nufftax is None: + nufft_cpu_exception() + + if not _fallback_warned: + _fallback_warned = True + logger.warning( + "INTERFEROMETER - `finufft` is not installed (it has no wheel for some platforms, " + "e.g. Linux aarch64), so the NumPy NUFFT path falls back to `nufftax`, which starts " + "JAX in this process. A multiprocessing search (e.g. Nautilus with " + "`number_of_cores > 1`) that forks numba workers after JAX has started can then " + "crash with a segmentation fault. Install `finufft` via `pip install finufft` to " + "keep the NumPy path JAX-free." + ) + + return nufftax + + +def nufft2d1( + x: np.ndarray, + y: np.ndarray, + c: np.ndarray, + n_modes: Tuple[int, int], + eps: float = 1.0e-12, + isign: int = 1, +) -> np.ndarray: + """ + Type-1 (adjoint) 2D NUFFT on the CPU, a drop-in for `nufftax.nufft2d1`: + + f[..., m2, m1] = sum_k c[..., k] exp(isign * i(m1 * x_k + m2 * y_k)) + + on the centred mode grid (mode 0 at index `N // 2`), returned with shape + `(..., n_modes[1], n_modes[0])` -- `nufftax`'s orientation, not `finufft`'s. + + Runs `finufft.nufft2d1`, whose thread count is left to `finufft`: it uses OpenMP and honours + `OMP_NUM_THREADS`, so pinning `OMP_NUM_THREADS=1` per process (e.g. one thread per Nautilus + worker) pins this transform too. The one exception is a process forked from one that already + ran finufft, which is pinned to one thread because libgomp deadlocks otherwise (see + `_threads_kwargs`). Falls back to `nufftax` (with a one-time warning) when + `finufft` is not importable. + + Parameters + ---------- + x, y + `[K]` non-uniform frequencies in `[-pi, pi)`; `x` pairs with `n_modes[0]`. + c + `[K]` (or stacked `[n_trans, K]`) complex strengths. + n_modes + `(N1, N2)`, the mode counts along the `x` and `y` axes. + eps + The requested NUFFT precision. + isign + The sign of the exponent, `+1` or `-1`. + + Returns + ------- + np.ndarray + `[..., N2, N1]` complex128. + """ + finufft = _load_finufft() + + if finufft is None: + return np.asarray(_nufftax_fallback().nufft2d1(x, y, c, n_modes, eps, isign)) + + f = finufft.nufft2d1( + np.ascontiguousarray(x, dtype=np.float64), + np.ascontiguousarray(y, dtype=np.float64), + np.ascontiguousarray(c, dtype=np.complex128), + tuple(int(n) for n in n_modes), + eps=float(eps), + isign=int(isign), + **_threads_kwargs(), + ) + + return np.swapaxes(f, -1, -2) + + +def nufft2d2( + x: np.ndarray, + y: np.ndarray, + f: np.ndarray, + eps: float = 1.0e-12, + isign: int = -1, +) -> np.ndarray: + """ + Type-2 (forward) 2D NUFFT on the CPU, a drop-in for `nufftax.nufft2d2`: + + c[..., k] = sum_{m1, m2} f[..., m2, m1] exp(isign * i(m1 * x_k + m2 * y_k)) + + for a mode grid `f` in `nufftax`'s orientation, shape `(..., N2, N1)`, centred (mode 0 at + index `N // 2`). A stacked `f` of shape `(n_trans, N2, N1)` is transformed in one call and + returns `(n_trans, K)`, as `TransformerNUFFT.transform_mapping_matrix` relies on. + + Runs `finufft.nufft2d2`, whose thread count is left to `finufft` (OpenMP, honouring + `OMP_NUM_THREADS`; one thread in a forked child, see `_threads_kwargs`). Falls back to `nufftax` (with a one-time warning) when `finufft` is not + importable. + + Parameters + ---------- + x, y + `[K]` non-uniform frequencies in `[-pi, pi)`; `x` pairs with the last axis of `f`. + f + `[N2, N1]` (or stacked `[n_trans, N2, N1]`) complex mode grid. + eps + The requested NUFFT precision. + isign + The sign of the exponent, `+1` or `-1`. + + Returns + ------- + np.ndarray + `[K]` (or `[n_trans, K]`) complex128. + """ + finufft = _load_finufft() + + if finufft is None: + return np.asarray(_nufftax_fallback().nufft2d2(x, y, f, eps, isign)) + + f = np.asarray(f, dtype=np.complex128) + + return finufft.nufft2d2( + np.ascontiguousarray(x, dtype=np.float64), + np.ascontiguousarray(y, dtype=np.float64), + np.ascontiguousarray(np.swapaxes(f, -1, -2)), + eps=float(eps), + isign=int(isign), + **_threads_kwargs(), + ) diff --git a/autoarray/operators/transformer.py b/autoarray/operators/transformer.py index 9bfefdf01..230e9252d 100644 --- a/autoarray/operators/transformer.py +++ b/autoarray/operators/transformer.py @@ -10,13 +10,16 @@ from autoarray.structures.visibilities import Visibilities from autoarray.structures.arrays import array_2d_util +from autoarray.operators import nufft_cpu from autoarray.operators import transformer_util from autoarray.dataset.interferometer.primary_beam import primary_beam_slim_from -# nufftax pulls in jax at import (~0.7s), which sessions that never touch an -# interferometer transformer should not pay for — deferred to _load_nufftax(), -# called from TransformerNUFFT's entry points (not only __init__, because -# unpickled instances in multiprocessing workers never re-run __init__). +# nufftax pulls in jax at import (~0.7s) and starts XLA, which sessions that +# never touch an interferometer transformer should not pay for, and which the +# NumPy path must never trigger (it runs finufft via `nufft_cpu` instead) — +# deferred to _load_nufftax(), called from the JAX branches of +# TransformerNUFFT's entry points (not from __init__, because unpickled +# instances in multiprocessing workers never re-run __init__). _nufftax = None _nufftax_loaded = False @@ -94,6 +97,13 @@ def batcher(args, dims, **kwargs): ) +def _require_nufftax(): + """Returns `nufftax` for a JAX branch, raising `nufftax_exception` if it is not installed.""" + if _load_nufftax() is None: + nufftax_exception() + return _nufftax + + def nufftax_exception(): raise ModuleNotFoundError( "\n--------------------\n" @@ -325,13 +335,25 @@ def __init__( **kwargs, ): """ - JAX-native Non-Uniform FFT for image -> visibilities, backed by `nufftax`. - - This is the default `TransformerNUFFT` in PyAutoArray. It uses the - `nufftax` library (https://github.com/GragasLab/nufftax), a pure-JAX - NUFFT implementation that supports `jax.jit`, `jax.grad`, and - `jax.vmap`. Note that nufftax requires JAX; on platforms with no JAX - wheels (notably Intel macOS) use `TransformerDFT` instead. + Non-Uniform FFT for image -> visibilities, backed by `finufft` on the + NumPy path and `nufftax` on the JAX path. + + This is the default `TransformerNUFFT` in PyAutoArray. The array + namespace of each call picks the backend: + + - `xp=np` (eager NumPy) runs `finufft` + (https://github.com/flatironinstitute/finufft), the CPU C++ NUFFT, + via `autoarray.operators.nufft_cpu`. This path never imports JAX, so + a NumPy fit (e.g. a numba Nautilus search with forked workers) does + not start XLA in its parent process. `finufft` uses OpenMP and + honours `OMP_NUM_THREADS`. + - `xp=jnp` runs `nufftax` (https://github.com/GragasLab/nufftax), a + pure-JAX port of `finufft` that supports `jax.jit`, `jax.grad` and + `jax.vmap`. `nufftax` is imported only on this path. + + The two agree to ~1e-14 relative. Where `finufft` has no wheel (e.g. + Linux aarch64) the NumPy path falls back to `nufftax` with a warning; + where neither is installed use `TransformerDFT` instead. Convention recipe (matches `TransformerDFT` to ~1e-13 relative across odd/even/non-square image sizes): @@ -345,8 +367,8 @@ def __init__( visibilities = nufftax.nufft2d2(x, y, image_flipped, eps, -1) * shift The `shift` factor is the half-pixel correction between autoarray's - grid centre at index `(N - 1) / 2` and nufftax's mode-0 at index - `N // 2`; nufftax does not apply it internally. + grid centre at index `(N - 1) / 2` and the NUFFT's mode-0 at index + `N // 2`; neither backend applies it internally. Parameters ---------- @@ -356,14 +378,14 @@ def __init__( real_space_mask The 2D mask defining the real-space image grid. eps - Requested NUFFT precision passed to nufftax. Defaults to `1e-12` + Requested NUFFT precision passed to the backend. Defaults to `1e-12` (effectively machine precision); relax to `1e-9` or `1e-6` for faster execution if marginal accuracy is acceptable. chunk_size If set to a positive integer, the forward and adjoint NUFFT calls split the visibility axis into chunks of this size and iterate (via ``jax.lax.scan`` on the JAX path, a Python loop - on the numpy path). This caps the nufftax gather-buffer + on the numpy path). This caps the NUFFT gather-buffer allocation (~``2 * chunk_size * nspread^2 * dtype_size``) at the cost of per-chunk overhead. Required for visibility counts above ~5M on a 40-80 GB GPU. If ``None`` (default), a single @@ -390,8 +412,10 @@ def __init__( """ from astropy import units - if _load_nufftax() is None: - nufftax_exception() + # Checks that a backend exists without importing nufftax (and so JAX): + # NumPy users must not start XLA just by building a transformer. + if not nufft_cpu.backend_available(): + nufft_cpu.nufft_cpu_exception() if chunk_size is not None and chunk_size <= 0: raise ValueError( @@ -431,16 +455,17 @@ def _forward_native(self, image_native_2d, xp=np): When ``self.chunk_size`` is set, the visibility axis is processed in fixed-size chunks via ``jax.lax.scan`` (JAX path) or a Python loop - (numpy path) — caps the nufftax gather-buffer allocation per call. + (numpy path) — caps the NUFFT gather-buffer allocation per call. The + numpy path runs `finufft` (`nufft_cpu`), the JAX path `nufftax`. """ - _load_nufftax() - K = int(self._x.shape[0]) if xp.__name__.startswith("jax"): import jax import jax.numpy as jnp + _require_nufftax() + img = jnp.asarray(image_native_2d)[::-1, :].astype(jnp.complex128) x_all = jnp.asarray(self._x) y_all = jnp.asarray(self._y) @@ -471,7 +496,7 @@ def body(carry, i): img = image_native_2d[::-1, :].astype(np.complex128) if self.chunk_size is None or self.chunk_size >= K: - out = _nufftax.nufft2d2(self._x, self._y, img, self.eps, -1) * self._shift + out = nufft_cpu.nufft2d2(self._x, self._y, img, self.eps, -1) * self._shift return np.array(np.asarray(out)) cs = int(self.chunk_size) @@ -479,7 +504,7 @@ def body(carry, i): for k0 in range(0, K, cs): k1 = min(k0 + cs, K) vis = ( - _nufftax.nufft2d2(self._x[k0:k1], self._y[k0:k1], img, self.eps, -1) + nufft_cpu.nufft2d2(self._x[k0:k1], self._y[k0:k1], img, self.eps, -1) * self._shift[k0:k1] ) parts.append(np.asarray(vis)) @@ -522,7 +547,8 @@ def image_from( """ Adjoint NUFFT: visibilities -> real-space (dirty) image. - Implemented as `nufftax.nufft2d1` with `conj(shift)` applied to the + Implemented as a type-1 `nufft2d1` (`finufft` via `nufft_cpu` for + `xp=np`, `nufftax` for `xp=jnp`) with `conj(shift)` applied to the visibilities and a final row-flip to return to autoarray's native orientation. The real part is taken to discard imaginary residue. @@ -536,8 +562,6 @@ def image_from( un-beamed transform, so the dirty image, dirty beam and the w-tilde quantities built from it do not depend on `primary_beam_slim`. """ - _load_nufftax() - n_y, n_x = self.real_space_mask.shape_native n_modes = (n_x, n_y) # nufftax wants (n1, n2) = (N_x, N_y) K = int(self._x.shape[0]) @@ -546,6 +570,8 @@ def image_from( import jax import jax.numpy as jnp + _require_nufftax() + x_all = jnp.asarray(self._x) y_all = jnp.asarray(self._y) shift_conj_all = jnp.asarray(np.conj(self._shift)) @@ -583,7 +609,7 @@ def body(f_accum, i): c_all = visibilities.array * np.conj(self._shift) if self.chunk_size is None or self.chunk_size >= K: - f = _nufftax.nufft2d1(self._x, self._y, c_all, n_modes, self.eps, +1) + f = nufft_cpu.nufft2d1(self._x, self._y, c_all, n_modes, self.eps, +1) image = np.array(np.asarray(f)[::-1, :].real) return Array2D(values=image, mask=self.real_space_mask) @@ -591,7 +617,7 @@ def body(f_accum, i): f_total = np.zeros((n_y, n_x), dtype=np.complex128) for k0 in range(0, K, cs): k1 = min(k0 + cs, K) - f_chunk = _nufftax.nufft2d1( + f_chunk = nufft_cpu.nufft2d1( self._x[k0:k1], self._y[k0:k1], c_all[k0:k1], n_modes, self.eps, +1 ) f_total = f_total + np.asarray(f_chunk) @@ -611,9 +637,9 @@ def transform_mapping_matrix(self, mapping_matrix, xp=np): If the transformer has a `primary_beam_slim` every column is multiplied by it on the slim grid before the scatter, giving `F P M`. - """ - _load_nufftax() + The numpy path runs `finufft` (`nufft_cpu`), the JAX path `nufftax`. + """ mapping_matrix = self._mapping_matrix_beamed_from(mapping_matrix, xp=xp) n_src = mapping_matrix.shape[1] @@ -623,6 +649,8 @@ def transform_mapping_matrix(self, mapping_matrix, xp=np): if xp.__name__.startswith("jax"): import jax.numpy as jnp + _require_nufftax() + # Real scatter, then cast: GPU complex128 scatter floor (autolens_profiling#308). mm_T = jnp.asarray(mapping_matrix).T source_images = jnp.zeros((n_src, n_y, n_x), dtype=mm_T.dtype) @@ -647,7 +675,7 @@ def transform_mapping_matrix(self, mapping_matrix, xp=np): source_images[np.arange(n_src)[:, None], rows[None, :], cols[None, :]] = mm_T flipped = source_images[:, ::-1, :].astype(np.complex128) vis_batched = ( - _nufftax.nufft2d2(self._x, self._y, flipped, self.eps, -1) + nufft_cpu.nufft2d2(self._x, self._y, flipped, self.eps, -1) * self._shift[None, :] ) return np.array(np.asarray(vis_batched).T) diff --git a/pyproject.toml b/pyproject.toml index 601bc489f..b78a6725b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,6 +44,15 @@ dependencies = [ "astropy>=5.0", "decorator>=4.0.0", "dill>=0.3.1.1", + # finufft is the CPU NUFFT of TransformerNUFFT and the precision-operator + # builder on the NumPy path, which must never import JAX (a JAX-threaded + # parent crashes the numba workers a multiprocessing search forks). The + # markers limit it to platforms with wheels (Linux / Windows x86_64, macOS + # arm64; Intel macOS is capped to 2.4.x, the last series with an x86_64 + # macOS wheel, so pip does not try to build 2.5's sdist); elsewhere (e.g. + # Linux aarch64) the NumPy path falls back to nufftax with a warning. + 'finufft>=2.4.0; platform_machine == "AMD64" or (platform_machine == "x86_64" and sys_platform != "darwin") or (sys_platform == "darwin" and platform_machine == "arm64")', + 'finufft>=2.4.0,<2.5; sys_platform == "darwin" and platform_machine == "x86_64"', "matplotlib>=3.7.0", "scipy<=1.17.1", "scikit-image<=0.26.0", @@ -87,7 +96,14 @@ optional = [ "tfp-nightly==0.26.0.dev20260713" ] test = ["pytest"] -dev = ["pytest", "black", "numba", "nufftax>=0.6.1,<0.7.0"] +dev = [ + "pytest", + "black", + "numba", + "nufftax>=0.6.1,<0.7.0", + 'finufft>=2.4.0; platform_machine == "AMD64" or (platform_machine == "x86_64" and sys_platform != "darwin") or (sys_platform == "darwin" and platform_machine == "arm64")', + 'finufft>=2.4.0,<2.5; sys_platform == "darwin" and platform_machine == "x86_64"', +] [tool.pytest.ini_options] testpaths = ["test_autoarray"] diff --git a/test_autoarray/conftest.py b/test_autoarray/conftest.py index b0c8ffcae..72df9afcc 100644 --- a/test_autoarray/conftest.py +++ b/test_autoarray/conftest.py @@ -283,25 +283,24 @@ def make_acs_quadrant(): def pytest_collection_modifyitems(config, items): - """Skip interferometer tests that need the default ``nufftax`` backend when - it is not installed. - - The default ``Interferometer`` transformer is the JAX-native - ``TransformerNUFFT``, backed by the optional ``nufftax`` dependency. When - that backend is absent, any test that builds a default interferometer or - transformer raises ``ModuleNotFoundError`` at construction. Skip exactly - those tests while keeping the explicit ``DFT`` transformer tests, which - have no ``nufftax`` dependency. + """Skip interferometer tests that need a ``TransformerNUFFT`` backend when + none is installed. + + The default ``Interferometer`` transformer is ``TransformerNUFFT``, backed + by ``finufft`` on the NumPy path (a base dependency on platforms with + wheels) and the optional ``nufftax`` on the JAX path. When neither is + importable, any test that builds a default interferometer or transformer + raises ``ModuleNotFoundError`` at construction. Skip exactly those tests + while keeping the explicit ``DFT`` transformer tests, which need neither. + Tests of the JAX path skip themselves via ``pytest.importorskip``. """ - try: - import nufftax # noqa: F401 + from autoarray.operators import nufft_cpu + if nufft_cpu.backend_available(): return - except ImportError: - pass skip_nufftax = pytest.mark.skip( - reason="nufftax (the default JAX interferometer backend) is not installed" + reason="no TransformerNUFFT backend (finufft or nufftax) is installed" ) for item in items: name = item.nodeid.rsplit("::", 1)[-1] diff --git a/test_autoarray/dataset/interferometer/test_dataset.py b/test_autoarray/dataset/interferometer/test_dataset.py index a7f003449..c12d0e799 100644 --- a/test_autoarray/dataset/interferometer/test_dataset.py +++ b/test_autoarray/dataset/interferometer/test_dataset.py @@ -351,13 +351,18 @@ def dataset(): assert recorded == [False, True, True, True, True] - # The default builder is the type-1 NUFFT, which runs on JAX too (nufftax is a JAX library), - # so the same kill switch has to demote it to the NumPy brute force rather than to the JAX - # one. Nothing above tests that: `use_jax` only ever selected between the two brute forces. + # Under `use_jax=True` the type-1 NUFFT builder runs on nufftax (JAX), so the same kill switch + # has to demote it to the NumPy brute force rather than to the JAX one. The default + # `use_jax=False` builder runs on finufft, which is not JAX, so the switch leaves it alone. monkeypatch.setattr(aa.Interferometer, "psf_precision_operator_from", original) monkeypatch.setenv("PYAUTO_DISABLE_JAX", "1") - operator_under_kill_switch = np.asarray(dataset().psf_precision_operator_from()) + operator_under_kill_switch = np.asarray( + dataset().psf_precision_operator_from(use_jax=True) + ) + operator_default_under_kill_switch = np.asarray( + dataset().psf_precision_operator_from() + ) monkeypatch.delenv("PYAUTO_DISABLE_JAX", raising=False) operator_via_numpy = np.asarray( @@ -366,6 +371,13 @@ def dataset(): np.testing.assert_array_equal(operator_under_kill_switch, operator_via_numpy) + pytest.importorskip("finufft") + + np.testing.assert_array_equal( + operator_default_under_kill_switch, + np.asarray(dataset().psf_precision_operator_from()), + ) + def _interferometer_for_precision_operator(mask_2d_7x7, transformer_class): n_visibilities = 5 @@ -466,14 +478,14 @@ def test__apply_sparse_operator__method_and_nufft_kwargs_reach_the_builder(mask_ # The operator only keeps `Khat`, so the plumbing is checked through it: routing to the brute # force and to the NUFFT must give the same operator to the pin's tolerance. np.testing.assert_allclose( - np.asarray(dataset_default.sparse_operator.Khat), - np.asarray(dataset_via_numpy.sparse_operator.Khat), + np.asarray(dataset_default.sparse_operator.khat_np), + np.asarray(dataset_via_numpy.sparse_operator.khat_np), rtol=1.0e-10, atol=1.0e-10 * np.abs(operator_via_numpy[0, 0]), ) np.testing.assert_allclose( - np.asarray(dataset_chunked.sparse_operator.Khat), - np.asarray(dataset_default.sparse_operator.Khat), + np.asarray(dataset_chunked.sparse_operator.khat_np), + np.asarray(dataset_default.sparse_operator.khat_np), rtol=1.0e-10, atol=1.0e-10 * np.abs(operator_via_numpy[0, 0]), ) diff --git a/test_autoarray/inversion/inversion/interferometer/test_interferometer.py b/test_autoarray/inversion/inversion/interferometer/test_interferometer.py index 37d84fa7a..947959888 100644 --- a/test_autoarray/inversion/inversion/interferometer/test_interferometer.py +++ b/test_autoarray/inversion/inversion/interferometer/test_interferometer.py @@ -487,7 +487,16 @@ def _sparse_parity_setup(seed=0): transformer_class=aa.TransformerDFT, ) - return mask, rng, dataset, dataset.apply_sparse_operator(use_jax=False) + # The brute-force builder (25 pixel-visibility pairs here) keeps the operator exact: the + # curvature matrices below have a condition number of ~4e9, which would amplify even the + # NUFFT builder's ~1e-12 (`eps`) error past the reconstruction tolerance, and these tests + # pin the sparse algebra, not the builder (pinned in test_inversion_interferometer_util). + return ( + mask, + rng, + dataset, + dataset.apply_sparse_operator(method="numpy", use_jax=False), + ) def _mapper_from(mask, pixels, shape, regularization=None): diff --git a/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py b/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py index 0f2673c34..423ba9b13 100644 --- a/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py +++ b/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py @@ -562,6 +562,11 @@ def test__nufft_precision_operator_from__method_routes_to_each_builder(): **inputs ) ) + operator_via_nufft_jax = np.asarray( + aa.util.inversion_interferometer.nufft_precision_operator_via_nufft_from( + **inputs, use_jax=True + ) + ) # `"numpy"` and `"jax"` are pure routing, so they return the brute-force arrays unchanged. np.testing.assert_array_equal( @@ -581,8 +586,7 @@ def test__nufft_precision_operator_from__method_routes_to_each_builder(): operator_via_jax, ) - # `use_jax=True` only upgrades a brute-force method, so `method="numpy"` with it set is - # the JAX brute force. + # `use_jax=True` upgrades `method="numpy"` to the JAX brute force. np.testing.assert_array_equal( np.asarray( aa.util.inversion_interferometer.nufft_precision_operator_from( @@ -592,23 +596,23 @@ def test__nufft_precision_operator_from__method_routes_to_each_builder(): operator_via_jax, ) - # Under the default `method="nufft"` it is ignored -- the NUFFT already runs on JAX, so - # honouring it there would demote every existing `use_jax=True` caller (the workspace - # `apply_sparse_operator(use_jax=True)` calls) from seconds to the O(N_pix * K) brute force. + # Under the default `method="nufft"` it selects the nufftax NUFFT -- still a fast path, so + # every existing `use_jax=True` caller (the workspace `apply_sparse_operator(use_jax=True)` + # calls) keeps seconds rather than being demoted to the O(N_pix * K) brute force. operator_use_jax = np.asarray( aa.util.inversion_interferometer.nufft_precision_operator_from( use_jax=True, **inputs ) ) - np.testing.assert_array_equal(operator_use_jax, operator_via_nufft) + np.testing.assert_array_equal(operator_use_jax, operator_via_nufft_jax) _assert_matches_brute_force(operator_use_jax, operator_via_np) # The control: the NUFFT array is not the JAX brute-force array, so the assertion above is # testing the routing and not two builders that happen to agree bitwise. assert not np.array_equal(operator_use_jax, operator_via_jax) - # The default is the NUFFT builder, and it agrees with the brute force. + # The default is the NUFFT builder on finufft, and it agrees with the brute force. operator_default = np.asarray( aa.util.inversion_interferometer.nufft_precision_operator_from(**inputs) ) @@ -635,15 +639,13 @@ def test__nufft_precision_operator_from__disable_jax_falls_back_to_the_numpy_bui aa.util.inversion_interferometer.nufft_precision_operator_via_np_from(**inputs) ) - # `PYAUTO_DISABLE_JAX=1` is the harness-level kill switch. Both the default NUFFT builder and - # the `"jax"` brute force run on JAX, so both must fall back -- and loudly, because the NumPy - # brute force is O(N_pix * K) where the NUFFT is O(K * nspread^2 + M log M). `use_jax=True` - # falls back either way: ignored under the default, and demoted again when it upgrades - # `method="numpy"` to the JAX brute force. + # `PYAUTO_DISABLE_JAX=1` is the harness-level kill switch. Every route that runs on JAX -- + # the `"jax"` brute force and the nufftax NUFFT (`use_jax=True`) -- must fall back, and + # loudly, because the NumPy brute force is O(N_pix * K) where the NUFFT is + # O(K * nspread^2 + M log M). monkeypatch.setenv("PYAUTO_DISABLE_JAX", "1") for kwargs in ( - {}, {"method": "jax"}, {"use_jax": True}, {"method": "numpy", "use_jax": True}, @@ -660,11 +662,44 @@ def test__nufft_precision_operator_from__disable_jax_falls_back_to_the_numpy_bui np.testing.assert_array_equal(operator, operator_via_np) assert "PYAUTO_DISABLE_JAX" in caplog.text + # The default NUFFT builder runs on finufft, which is not JAX, so the switch leaves it alone. + # Only when finufft is missing (and the NumPy NUFFT would fall back to nufftax) is it demoted. + if aa.util.inversion_interferometer.nufft_cpu._load_finufft() is not None: + caplog.clear() + + with caplog.at_level("WARNING"): + operator = np.asarray( + aa.util.inversion_interferometer.nufft_precision_operator_from(**inputs) + ) + + np.testing.assert_array_equal( + operator, + np.asarray( + aa.util.inversion_interferometer.nufft_precision_operator_via_nufft_from( + **inputs + ) + ), + ) + assert "PYAUTO_DISABLE_JAX" not in caplog.text + + monkeypatch.setattr( + aa.util.inversion_interferometer.nufft_cpu, "_load_finufft", lambda: None + ) + caplog.clear() + + with caplog.at_level("WARNING"): + operator = np.asarray( + aa.util.inversion_interferometer.nufft_precision_operator_from(**inputs) + ) + + np.testing.assert_array_equal(operator, operator_via_np) + assert "PYAUTO_DISABLE_JAX" in caplog.text + + monkeypatch.undo() + # Without the variable the default is the NUFFT builder again, so the fallback is the switch's # doing and not a permanent demotion. - monkeypatch.delenv("PYAUTO_DISABLE_JAX", raising=False) - - pytest.importorskip("nufftax") + pytest.importorskip("finufft") np.testing.assert_array_equal( np.asarray( @@ -687,13 +722,15 @@ def test__nufft_precision_operator_from__nufftax_absent_falls_back_to_the_numpy_ aa.util.inversion_interferometer.nufft_precision_operator_via_np_from(**inputs) ) - # `nufftax` is an optional dependency, so the default builder has to survive its absence -- - # loudly, naming the O(N_pix * K) cost the caller now pays, rather than silently. + # `nufftax` is an optional dependency, so the `use_jax=True` NUFFT route has to survive its + # absence -- loudly, naming the O(N_pix * K) cost the caller now pays, rather than silently. monkeypatch.setattr(aa.util.inversion_interferometer, "_load_nufftax", lambda: None) with caplog.at_level("WARNING"): operator = np.asarray( - aa.util.inversion_interferometer.nufft_precision_operator_from(**inputs) + aa.util.inversion_interferometer.nufft_precision_operator_from( + use_jax=True, **inputs + ) ) np.testing.assert_array_equal(operator, operator_via_np) @@ -702,9 +739,36 @@ def test__nufft_precision_operator_from__nufftax_absent_falls_back_to_the_numpy_ # The builder itself raises rather than falling back: only the dispatcher chooses. with pytest.raises(ModuleNotFoundError): aa.util.inversion_interferometer.nufft_precision_operator_via_nufft_from( - **inputs + **inputs, use_jax=True ) + # The default route runs on finufft and does not need nufftax at all. + if aa.util.inversion_interferometer.nufft_cpu._load_finufft() is not None: + np.testing.assert_array_equal( + np.asarray( + aa.util.inversion_interferometer.nufft_precision_operator_from(**inputs) + ), + np.asarray( + aa.util.inversion_interferometer.nufft_precision_operator_via_nufft_from( + **inputs + ) + ), + ) + + # With no NUFFT backend at all the default route falls back to the brute force too. + monkeypatch.setattr( + aa.util.inversion_interferometer.nufft_cpu, "backend_available", lambda: False + ) + caplog.clear() + + with caplog.at_level("WARNING"): + operator = np.asarray( + aa.util.inversion_interferometer.nufft_precision_operator_from(**inputs) + ) + + np.testing.assert_array_equal(operator, operator_via_np) + assert "finufft" in caplog.text + def _numpy_backend_fixtures(): """ @@ -2069,9 +2133,10 @@ def test__nufft_precision_operator_via_nufft__type1_core_refactor_is_bit_for_bit for inputs in (_nufft_pin_inputs_7x7(), _nufft_pin_inputs_16x16()): for chunk_size in (None, 64): + # The reference is the nufftax loop, so it pins the `use_jax=True` route. np.testing.assert_array_equal( aa.util.inversion_interferometer.nufft_precision_operator_via_nufft_from( - **inputs, chunk_size=chunk_size + **inputs, chunk_size=chunk_size, use_jax=True ), _nufft_precision_operator_via_nufft_reference_from( **inputs, chunk_size=chunk_size diff --git a/test_autoarray/operators/test_nufft_cpu.py b/test_autoarray/operators/test_nufft_cpu.py new file mode 100644 index 000000000..3e6c85b56 --- /dev/null +++ b/test_autoarray/operators/test_nufft_cpu.py @@ -0,0 +1,406 @@ +""" +The CPU (finufft) NUFFT backend of the NumPy path: its orientation against a direct sum, its parity +with nufftax through `TransformerNUFFT` and the precision-operator builder, its fallback to nufftax, +and a fresh-process guard that the whole NumPy interferometer route never starts JAX. +""" + +import os +import subprocess +import sys +import textwrap + +import numpy as np +import pytest + +import autoarray as aa +from autoarray.operators import nufft_cpu + + +def _relative_max_difference(a, b): + a = np.asarray(a) + b = np.asarray(b) + return np.max(np.abs(a - b)) / np.max(np.abs(b)) + + +def _direct_type1(x, y, c, n_modes, isign): + """`f[m2, m1] = sum_k c_k exp(isign * i(m1 * x_k + m2 * y_k))` on the centred mode grid.""" + n1, n2 = n_modes + m1 = np.arange(n1) - n1 // 2 + m2 = np.arange(n2) - n2 // 2 + phase = np.exp( + isign + * 1j + * (m1[None, None, :] * x[:, None, None] + m2[None, :, None] * y[:, None, None]) + ) + return np.sum(c[:, None, None] * phase, axis=0) + + +def test__nufft2d1_and_nufft2d2__match_direct_sums_in_nufftax_orientation(): + pytest.importorskip("finufft") + + rng = np.random.default_rng(seed=0) + x = rng.uniform(-np.pi, np.pi, size=30) + y = rng.uniform(-np.pi, np.pi, size=30) + c = rng.normal(size=30) + 1j * rng.normal(size=30) + n_modes = (6, 5) # non-square, odd and even, so a transposed result cannot pass + + f = nufft_cpu.nufft2d1(x, y, c, n_modes, 1.0e-12, 1) + + assert f.shape == (5, 6) + assert _relative_max_difference(f, _direct_type1(x, y, c, n_modes, +1)) < 1.0e-10 + + grid = rng.normal(size=(5, 6)) + 1j * rng.normal(size=(5, 6)) + + vis = nufft_cpu.nufft2d2(x, y, grid, 1.0e-12, -1) + + m1 = np.arange(6) - 3 + m2 = np.arange(5) - 2 + expected = np.array( + [ + np.sum(grid * np.exp(-1j * (m1[None, :] * xk + m2[:, None] * yk))) + for xk, yk in zip(x, y) + ] + ) + + assert vis.shape == (30,) + assert _relative_max_difference(vis, expected) < 1.0e-10 + + # A stacked grid transforms in one call, one row per grid. + stacked = nufft_cpu.nufft2d2(x, y, np.stack([grid, 2.0 * grid]), 1.0e-12, -1) + + assert stacked.shape == (2, 30) + assert _relative_max_difference(stacked[1], 2.0 * expected) < 1.0e-10 + + +def _transformer_setup(seed=3): + rng = np.random.default_rng(seed=seed) + + mask = aa.Mask2D.circular(shape_native=(10, 11), pixel_scales=0.2, radius=0.9) + uv_wavelengths = rng.normal(size=(53, 2)) * 2.0e5 + + image = aa.Array2D(values=rng.normal(size=mask.pixels_in_mask), mask=mask) + visibilities = aa.Visibilities(visibilities=rng.normal(size=(53, 2))) + mapping_matrix = rng.normal(size=(mask.pixels_in_mask, 4)) + primary_beam = aa.Array2D( + values=rng.uniform(0.2, 1.0, size=mask.pixels_in_mask), mask=mask + ) + + return mask, uv_wavelengths, image, visibilities, mapping_matrix, primary_beam + + +def _transforms_from(transformer, image, visibilities, mapping_matrix): + return ( + np.asarray(transformer.visibilities_from(image=image).array), + np.asarray(transformer.image_from(visibilities=visibilities).array), + np.asarray(transformer.transform_mapping_matrix(mapping_matrix=mapping_matrix)), + ) + + +def test__transformer_nufft__numpy_path_finufft_matches_nufftax(monkeypatch): + """ + Every NumPy-path transform of `TransformerNUFFT` on finufft matches the same transform on + nufftax (reached through the fallback), chunked and unchunked, with and without a primary + beam -- and matches the JAX path, which is nufftax by construction. + """ + pytest.importorskip("finufft") + pytest.importorskip("nufftax") + import jax.numpy as jnp + + mask, uv_wavelengths, image, visibilities, mapping_matrix, primary_beam = ( + _transformer_setup() + ) + + for chunk_size in (None, 16): + for beam in (None, primary_beam): + transformer = aa.TransformerNUFFT( + uv_wavelengths=uv_wavelengths, + real_space_mask=mask, + chunk_size=chunk_size, + primary_beam=beam, + ) + + finufft_transforms = _transforms_from( + transformer, image, visibilities, mapping_matrix + ) + + with monkeypatch.context() as m: + m.setattr(nufft_cpu, "_load_finufft", lambda: None) + m.setattr(nufft_cpu, "_fallback_warned", True) + nufftax_transforms = _transforms_from( + transformer, image, visibilities, mapping_matrix + ) + + for finufft_transform, nufftax_transform in zip( + finufft_transforms, nufftax_transforms + ): + assert ( + _relative_max_difference(finufft_transform, nufftax_transform) + < 1.0e-10 + ) + + jax_vis = transformer.visibilities_from(image=image, xp=jnp) + + assert ( + _relative_max_difference( + finufft_transforms[0], np.asarray(jax_vis.array) + ) + < 1.0e-10 + ) + + +def _precision_operator_inputs(seed=5): + rng = np.random.default_rng(seed=seed) + + mask = aa.Mask2D.circular(shape_native=(12, 12), pixel_scales=0.2, radius=0.9) + uv_wavelengths = rng.normal(size=(200, 2)) * 2.0e5 + noise_map_real = rng.uniform(0.5, 2.0, size=200) + + return dict( + noise_map_real=noise_map_real, + uv_wavelengths=uv_wavelengths, + shape_masked_pixels_2d=mask.shape_native_masked_pixels, + grid_radians_2d=mask.derive_grid.all_false.in_radians.native.array, + ) + + +def _assert_matches_mixed(operator, reference): + np.testing.assert_allclose( + operator, reference, rtol=1.0e-10, atol=1.0e-10 * np.abs(reference[0, 0]) + ) + + +def test__nufft_precision_operator_from__default_finufft_matches_brute_force_and_nufftax(): + pytest.importorskip("finufft") + + inputs = _precision_operator_inputs() + + operator_default = np.asarray( + aa.util.inversion_interferometer.nufft_precision_operator_from(**inputs) + ) + operator_via_np = np.asarray( + aa.util.inversion_interferometer.nufft_precision_operator_from( + method="numpy", **inputs + ) + ) + + _assert_matches_mixed(operator_default, operator_via_np) + + # Chunking only splits the linear sum. + _assert_matches_mixed( + np.asarray( + aa.util.inversion_interferometer.nufft_precision_operator_from( + chunk_size=37, **inputs + ) + ), + operator_via_np, + ) + + pytest.importorskip("nufftax") + + operator_use_jax = np.asarray( + aa.util.inversion_interferometer.nufft_precision_operator_from( + use_jax=True, **inputs + ) + ) + + _assert_matches_mixed(operator_default, operator_use_jax) + + +def test__nufft_cpu__falls_back_to_nufftax_with_one_warning_when_finufft_is_missing( + monkeypatch, caplog +): + pytest.importorskip("nufftax") + + rng = np.random.default_rng(seed=1) + x = rng.uniform(-np.pi, np.pi, size=20) + y = rng.uniform(-np.pi, np.pi, size=20) + c = rng.normal(size=20) + 0j + + monkeypatch.setattr(nufft_cpu, "_load_finufft", lambda: None) + monkeypatch.setattr(nufft_cpu, "_fallback_warned", False) + + with caplog.at_level("WARNING"): + f = nufft_cpu.nufft2d1(x, y, c, (6, 5), 1.0e-12, 1) + nufft_cpu.nufft2d2(x, y, f, 1.0e-12, -1) + + assert isinstance(f, np.ndarray) + assert f.shape == (5, 6) + assert _relative_max_difference(f, _direct_type1(x, y, c, (6, 5), +1)) < 1.0e-10 + + warnings = [r for r in caplog.records if "finufft" in r.getMessage()] + + assert len(warnings) == 1 + assert "JAX" in warnings[0].getMessage() + assert "segmentation fault" in warnings[0].getMessage() + + +def test__nufft_cpu__no_backend_raises(monkeypatch): + monkeypatch.setattr(nufft_cpu, "_load_finufft", lambda: None) + monkeypatch.setattr(nufft_cpu, "backend_available", lambda: False) + + from autoarray.operators import transformer + + monkeypatch.setattr(transformer, "_load_nufftax", lambda: None) + + with pytest.raises(ModuleNotFoundError) as exc: + aa.TransformerNUFFT( + uv_wavelengths=np.array([[0.2, 1.0], [0.5, 1.1]]), + real_space_mask=aa.Mask2D.all_false(shape_native=(5, 5), pixel_scales=0.1), + ) + + assert "pip install finufft" in str(exc.value) + + with pytest.raises(ModuleNotFoundError): + nufft_cpu.nufft2d1(np.zeros(2), np.zeros(2), np.ones(2) + 0j, (4, 4)) + + +_GUARD_SCRIPT = textwrap.dedent( + """ + import sys + + import numpy as np + + import autoarray as aa + + rng = np.random.default_rng(seed=0) + n_vis = 300 + + mask = aa.Mask2D.circular(shape_native=(16, 16), pixel_scales=0.2, radius=1.2) + dataset = aa.Interferometer( + data=aa.Visibilities(visibilities=rng.normal(size=(n_vis, 2))), + noise_map=aa.VisibilitiesNoiseMap(visibilities=np.ones((n_vis, 2))), + uv_wavelengths=rng.normal(size=(n_vis, 2)) * 2.0e5, + real_space_mask=mask, + transformer_class=aa.TransformerNUFFT, + ) + + operator = dataset.psf_precision_operator_from() + dataset_sparse = dataset.apply_sparse_operator(nufft_precision_operator=operator) + dataset.apply_sparse_operator() + + transformer = dataset.transformer + image = aa.Array2D(values=rng.normal(size=mask.pixels_in_mask), mask=mask) + transformer.visibilities_from(image=image) + transformer.image_from(visibilities=dataset.data) + transformer.transform_mapping_matrix( + mapping_matrix=rng.normal(size=(mask.pixels_in_mask, 3)) + ) + + assert dataset_sparse.sparse_operator is not None + + print("nufftax_imported", "nufftax" in sys.modules) + + backends = None + if "jax" in sys.modules: + from jax._src import xla_bridge + + backends = len(xla_bridge._backends) + print("jax_imported", "jax" in sys.modules) + print("xla_backends", backends) + """ +) + + +def test__numpy_interferometer_route__never_imports_nufftax_or_starts_jax(): + """ + The whole NumPy interferometer route -- precision operator, sparse operator (dirty image), + forward, adjoint and mapping-matrix transforms -- run in a fresh process must neither import + nufftax nor start JAX's XLA backend. A parent that has started JAX before a multiprocessing + search forks its numba workers crashes them (PyAutoArray#654), so this is the property the + finufft backend exists for. + """ + pytest.importorskip("finufft") + + env = dict(os.environ) + env.setdefault("NUMBA_CACHE_DIR", "/tmp/numba_cache") + env.setdefault("MPLCONFIGDIR", "/tmp/matplotlib") + env.pop("PYAUTO_DISABLE_JAX", None) + + result = subprocess.run( + [sys.executable, "-c", _GUARD_SCRIPT], + capture_output=True, + text=True, + env=env, + timeout=600, + ) + + assert result.returncode == 0, result.stderr[-4000:] + + lines = dict( + line.split(" ", 1) + for line in result.stdout.splitlines() + if line.startswith(("nufftax_imported", "jax_imported", "xla_backends")) + ) + + assert lines["nufftax_imported"] == "False" + assert lines["xla_backends"] in ("None", "0"), result.stdout + + +_FORK_SCRIPT = textwrap.dedent( + """ + import multiprocessing + + import numpy as np + + from autoarray.operators import nufft_cpu + + + def transform(seed): + rng = np.random.default_rng(seed) + x = rng.uniform(-np.pi, np.pi, size=20000) + y = rng.uniform(-np.pi, np.pi, size=20000) + f = nufft_cpu.nufft2d1(x, y, rng.normal(size=20000) + 0j, (64, 64)) + return float(np.abs(nufft_cpu.nufft2d2(x, y, f)).sum()) + + + if __name__ == "__main__": + transform(0) + with multiprocessing.get_context("fork").Pool(2) as pool: + pool.map(transform, range(4)) + print("forked_ok") + """ +) + + +def test__nufft_cpu__forked_children_after_a_multithreaded_parent_do_not_deadlock( + tmp_path, +): + """ + libgomp deadlocks a forked child that opens a multi-threaded OpenMP region after its parent + ran one -- the Nautilus pattern of building the operator in the parent and transforming in + forked workers. `nufft_cpu` pins forked children to one thread; without that this hangs. + """ + pytest.importorskip("finufft") + + import multiprocessing + + if "fork" not in multiprocessing.get_all_start_methods(): + pytest.skip("fork start method unavailable on this platform") + + script = tmp_path / "fork_script.py" + script.write_text(_FORK_SCRIPT) + + env = dict(os.environ) + env["OMP_NUM_THREADS"] = "2" + env.setdefault("NUMBA_CACHE_DIR", "/tmp/numba_cache") + env.setdefault("MPLCONFIGDIR", "/tmp/matplotlib") + + # Own session, so a deadlocked run's forked workers are killed with it rather than orphaned. + process = subprocess.Popen( + [sys.executable, str(script)], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + env=env, + start_new_session=True, + ) + + try: + stdout, stderr = process.communicate(timeout=120) + except subprocess.TimeoutExpired: + os.killpg(process.pid, 9) + process.communicate() + pytest.fail("forked nufft_cpu workers deadlocked after a multi-threaded parent") + + assert process.returncode == 0, stderr[-4000:] + assert "forked_ok" in stdout diff --git a/test_autoarray/operators/test_transformer.py b/test_autoarray/operators/test_transformer.py index 67ea899e8..a18464035 100644 --- a/test_autoarray/operators/test_transformer.py +++ b/test_autoarray/operators/test_transformer.py @@ -142,10 +142,13 @@ def test__nufft__transform_mapping_matrix__ones_mapping_matrix__first_element_ma def test__nufft__transform_mapping_matrix__real_scatter_matches_complex_scatter_exactly(): """Scattering the real mapping matrix and casting afterwards (autolens_profiling#308) must be bit-identical to the previous - cast-then-scatter formula, for NumPy and under ``jax.jit``.""" + cast-then-scatter formula, for NumPy (on ``finufft``) and under ``jax.jit`` + (on ``nufftax``).""" + pytest.importorskip("nufftax") import jax import jax.numpy as jnp + from autoarray.operators import nufft_cpu from autoarray.operators import transformer as transformer_module rng = np.random.default_rng(seed=4) @@ -170,7 +173,7 @@ def test__nufft__transform_mapping_matrix__real_scatter_matches_complex_scatter_ ) expected_numpy = np.array( np.asarray( - nufftax.nufft2d2( + nufft_cpu.nufft2d2( transformer._x, transformer._y, source_images[:, ::-1, :], @@ -278,6 +281,7 @@ def test__nufft__chunk_size__image_from_numpy_matches_unchunked(): def test__nufft__chunk_size__jax_paths_match_unchunked(): + pytest.importorskip("nufftax") import jax import jax.numpy as jnp @@ -324,6 +328,7 @@ def test__nufft__chunk_size__jax_jit_traces_with_scan(): without unrolling the chunk loop. We exercise this by JIT-ing and confirming the compiled HLO graph has bounded size (no per-chunk re-emission).""" + pytest.importorskip("nufftax") import jax import jax.numpy as jnp From 3f3c8619f4dc91d746c3379c3b5bdcebc53127ea Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Sat, 24 Oct 2026 17:16:00 +0100 Subject: [PATCH 2/2] fix: run the type-1 finufft transform on one thread so its products are reproducible A multithreaded type-1 spread sums thread subgrids in a non-deterministic order, so two identical precision-operator / dirty-image builds differed in the last bits (13 bitwise-equality tests failed with OMP threads > 1; CI py3.13 red). These are one-off precomputations; one thread is ~0.9 s for 5.5M visibilities, still ~10x faster than nufftax, and the likelihood never uses this transform. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_013C7Zd6ibK9J4V8CFEUzrJU --- autoarray/operators/nufft_cpu.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/autoarray/operators/nufft_cpu.py b/autoarray/operators/nufft_cpu.py index af5776ccb..3b54ba518 100644 --- a/autoarray/operators/nufft_cpu.py +++ b/autoarray/operators/nufft_cpu.py @@ -154,12 +154,13 @@ def nufft2d1( on the centred mode grid (mode 0 at index `N // 2`), returned with shape `(..., n_modes[1], n_modes[0])` -- `nufftax`'s orientation, not `finufft`'s. - Runs `finufft.nufft2d1`, whose thread count is left to `finufft`: it uses OpenMP and honours - `OMP_NUM_THREADS`, so pinning `OMP_NUM_THREADS=1` per process (e.g. one thread per Nautilus - worker) pins this transform too. The one exception is a process forked from one that already - ran finufft, which is pinned to one thread because libgomp deadlocks otherwise (see - `_threads_kwargs`). Falls back to `nufftax` (with a one-time warning) when - `finufft` is not importable. + Runs `finufft.nufft2d1` on **one thread**. A multithreaded type-1 spread adds each thread's + subgrid into the output in a non-deterministic order, so two identical calls can differ in + the last bits; the products built from this transform (the dirty image and the NUFFT + precision operator) are one-off precomputations that must be bitwise reproducible, and one + thread is still ~10x faster than `nufftax` on the CPU (~0.9 s for 5.5M visibilities). It also + keeps this transform out of the libgomp fork deadlock `_threads_kwargs` guards against. + Falls back to `nufftax` (with a one-time warning) when `finufft` is not importable. Parameters ---------- @@ -191,7 +192,7 @@ def nufft2d1( tuple(int(n) for n in n_modes), eps=float(eps), isign=int(isign), - **_threads_kwargs(), + nthreads=1, ) return np.swapaxes(f, -1, -2)