Skip to content

feat: finufft as the CPU NUFFT backend (JAX-free NumPy interferometer runs) #655

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Jammy2211 merged 2 commits into from
Oct 24, 2026
Merged

feat: finufft as the CPU NUFFT backend (JAX-free NumPy interferometer runs) #655

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
157e210
feat: finufft as the CPU NUFFT backend; nufftax only on the JAX path …
Jammy2211 Oct 24, 2026
3f3c861
fix: run the type-1 finufft transform on one thread so its products a…
Jammy2211 Oct 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 19 additions & 14 deletions autoarray/dataset/interferometer/dataset.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down
115 changes: 79 additions & 36 deletions autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@

from autoarray import exc

from autoarray.operators import nufft_cpu
from autoarray.operators.transformer import _load_nufftax, nufftax_exception

try:
Expand Down Expand Up @@ -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`.

Expand All @@ -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"
Expand All @@ -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"

Expand All @@ -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":
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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
-------
Expand All @@ -927,6 +965,7 @@ def nufft_precision_operator_via_nufft_from(
chunk_size=chunk_size,
centred=False,
zero_nyquist=True,
use_jax=use_jax,
)


Expand Down Expand Up @@ -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 * (
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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(
Expand All @@ -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(
Expand Down
Loading
Loading