From 87205219fda6c3e00b7246b61c7dc56f757c5a99 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Tue, 8 Sep 2026 13:11:09 -0400 Subject: [PATCH 1/2] perf: build nufft_precision_operator_from as a type-1 NUFFT (#539) The sparse-operator preload is the real part of a type-1 (adjoint) NUFFT of the inverse-variance weights 1/sigma^2 evaluated on the doubled-extent grid: `nufft2d1(-x, y, w, (2Nx, 2Ny), eps, +1)`, followed by an `ifftshift` and zeroing of the Nyquist row `Ny` and column `Nx`. That replaces the brute-force O(N_pix * K) cosine accumulation with an O(K + N_pix log N_pix) transform. - New `nufft_precision_operator_via_nufft_from`, and `method="nufft"` (the new default) on the `nufft_precision_operator_from` dispatcher, with `eps=1e-12` and mandatory visibility chunking taken from the transformer. - The numpy and JAX brute-force builders are kept untouched as the reference implementations (`method="numpy"` / `"jax"`; `use_jax=True` still routes to the JAX one). - Loud, logged fallback to `method="numpy"` when `PYAUTO_DISABLE_JAX=1` is set or nufftax is not importable; never silent, anything else raises. - `dataset.py`: `apply_sparse_operator` and `psf_precision_operator_from` gain `method`, `eps` and `nufft_chunk_size`, defaulting to the transformer's own `eps` / `chunk_size` when it is a `TransformerNUFFT`. Pins: mixed-tolerance parity against the numpy builder at 4.4e-14 of peak at K=300; chunked == one-shot at 5.3e-14; the wrong-sign control fails at 0.187 of peak. Measured in autolens_profiling#229: alma 2101 s -> 7.3 s wall (289x; 111x in CPU-seconds), alma_high 22 s. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_018hLF3ZAcz5MmaSJBEcLkvF --- autoarray/dataset/interferometer/dataset.py | 134 +++++-- .../inversion_interferometer_util.py | 312 ++++++++++++++- .../dataset/interferometer/test_dataset.py | 127 ++++++ .../test_inversion_interferometer_util.py | 360 +++++++++++++++++- 4 files changed, 883 insertions(+), 50 deletions(-) diff --git a/autoarray/dataset/interferometer/dataset.py b/autoarray/dataset/interferometer/dataset.py index 4b9759ee2..33b5035f4 100644 --- a/autoarray/dataset/interferometer/dataset.py +++ b/autoarray/dataset/interferometer/dataset.py @@ -230,6 +230,9 @@ def apply_sparse_operator( self, nufft_precision_operator=None, batch_size: int = 128, + method: str = "nufft", + eps: Optional[float] = None, + nufft_chunk_size: Optional[int] = None, chunk_k: int = 2048, show_progress: bool = False, show_memory: bool = False, @@ -246,24 +249,12 @@ def apply_sparse_operator( and is used automatically by `FitInterferometer` when performing pixelized reconstructions via the inversion module. - Computing the NUFFT precision matrix from scratch can be very slow (runtime scales with both - the number of visibilities and the real-space mask resolution — potentially hours for large - datasets). The result can be cached to disk and reloaded to avoid recomputation. - - Both `TransformerDFT` and `TransformerNUFFT` are supported here and agree to ~3e-13 relative. - Which is faster is governed by the product `N_vis * N_pix`, because the DFT setup cost scales - as `O(N_vis * N_pix)` whereas the NUFFT scales as `O((N_vis + N_pix) log N)` on top of a fixed - ~2s overhead: - - - below `N_vis * N_pix ~ 1e7`, `TransformerDFT` is faster (measured 0.2-0.7x the NUFFT time) - - above it, `TransformerNUFFT` is faster (1.2-1.9x at 1e7-1e8) - - above `~1e8`, the DFT's `O(N_vis * N_pix)` allocation makes it infeasible rather than merely - slow — extrapolating a measured 446 MB at N_vis=4e3 / N_pix=1e4 gives ~109 GB at 1M - visibilities, whereas the NUFFT path allocates nothing beyond its working buffers. - - As a rule of thumb at a typical 64x64 mask the crossover sits near 5,000 visibilities, but on a - coarser mask the DFT stays ahead well beyond that — it is the product that matters, not the - visibility count alone. + 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 + visibilities. The brute-force builders (`method="numpy"` / `"jax"`, and the `use_jax` kwarg) + are `O(N_vis * N_pix)` and can take minutes to hours; they are kept as the reference the + NUFFT builder is pinned against. Either way the result can be cached to disk and reloaded + via `nufft_precision_operator=`. Parameters ---------- @@ -274,9 +265,20 @@ def apply_sparse_operator( batch_size The number of real-space pixels processed per batch when building the sparse operator. Reducing this lowers peak memory usage at the cost of speed. + method + Which builder computes the precision operator: `"nufft"` (default, the type-1 NUFFT), + `"numpy"` or `"jax"` (the brute-force reference builders). + eps + The requested NUFFT precision of the `"nufft"` builder. `None` takes the transformer's + own `eps` when it is a `TransformerNUFFT`, else `1e-12`. + nufft_chunk_size + The visibility chunk size of the `"nufft"` builder, a memory ceiling rather than an + optimisation. `None` takes the transformer's own `chunk_size` when it is a + `TransformerNUFFT`, else no chunking. chunk_k - The number of visibilities processed per chunk when computing the NUFFT precision matrix - inside `psf_precision_operator_from()`. Reducing this lowers peak memory usage. + The number of visibilities processed per chunk by the brute-force builders when computing + the NUFFT precision matrix inside `psf_precision_operator_from()`. Reducing this lowers + peak memory usage. show_progress If `True`, a progress bar is displayed while computing the NUFFT precision matrix. show_memory @@ -350,25 +352,41 @@ def apply_sparse_operator( if nufft_precision_operator is None: - logger.info( - "INTERFEROMETER - Computing NUFFT Precision Operator; runtime scales with visibility count and mask resolution, CPU run times may exceed hours." - ) + logger.info("INTERFEROMETER - Computing NUFFT Precision Operator.") n_vis = self.uv_wavelengths.shape[0] n_pix = self.real_space_mask.pixels_in_mask - if n_vis * n_pix < 10**7 and not isinstance( - self.transformer, TransformerDFT - ): + if method != "nufft" or use_jax: logger.info( - f"INTERFEROMETER - This dataset is small for the NUFFT setup path " - f"(N_vis x N_pix = {n_vis * n_pix:.1e}, below the ~1e7 crossover). " - f"`TransformerDFT` computes this operator faster below that point; " - f"`TransformerNUFFT` wins above it, and above ~1e8 it is the only " - f"option that fits in memory. Both are supported here." + f"INTERFEROMETER - The precision operator is being built by a brute-force " + f"builder, which is O(N_vis x N_pix) = O({n_vis * n_pix:.1e}) and can take " + f"minutes to hours. The default `method='nufft'` builds the same array as a " + f"type-1 NUFFT in seconds." ) + if isinstance(self.transformer, TransformerDFT): + # This is about the transformer, not the precision operator: the operator is + # built by `nufft_precision_operator_from` either way and does not go through + # the transformer at all. The DFT transformer allocates O(N_vis x N_pix) for + # every subsequent transform, which is what becomes infeasible at scale -- + # extrapolating a measured 446 MB at N_vis=4e3 / N_pix=1e4 gives ~109 GB at a + # million visibilities, whereas `TransformerNUFFT` allocates nothing beyond its + # working buffers. Below `N_vis x N_pix ~ 1e7` the DFT is the faster transform + # (0.2-0.7x the NUFFT time) and there is nothing to warn about. + if n_vis * n_pix > 10**7: + logger.info( + f"INTERFEROMETER - This dataset uses `TransformerDFT` at " + f"N_vis x N_pix = {n_vis * n_pix:.1e}, above the ~1e7 crossover where " + f"`TransformerNUFFT` transforms faster (1.2-1.9x at 1e7-1e8); above ~1e8 " + f"the DFT's O(N_vis x N_pix) allocation makes it infeasible rather than " + f"merely slow. The two agree to ~3e-13 relative." + ) + nufft_precision_operator = self.psf_precision_operator_from( + method=method, + eps=eps, + nufft_chunk_size=nufft_chunk_size, chunk_k=chunk_k, show_progress=show_progress, show_memory=show_memory, @@ -401,6 +419,9 @@ def psf_precision_operator_from( show_progress: bool = False, show_memory: bool = False, use_jax: bool = False, + method: str = "nufft", + eps: Optional[float] = None, + nufft_chunk_size: Optional[int] = None, ): """ Compute the NUFFT precision matrix for this interferometer dataset. @@ -409,23 +430,34 @@ def psf_precision_operator_from( baseline, weighted by the noise map. It is the core precomputed quantity required for efficient pixelized source reconstruction via the sparse linear algebra formalism. - This computation can be very slow for large datasets (runtime scales with the number of - visibilities multiplied by the number of unmasked real-space pixels). For datasets with - tens of thousands of visibilities and high-resolution masks, computation can take hours - on a CPU. The result should be saved to disk and reloaded rather than recomputed on each - run. Use `apply_sparse_operator(nufft_precision_operator=...)` to attach a cached result. + The default builder (`method="nufft"`) computes this as a type-1 (adjoint) NUFFT, which is + `O(N_vis * nspread^2 + M log M)` for `M = 4 * Ny * Nx` — seconds even at a million + visibilities. The brute-force builders (`method="numpy"` / `"jax"`, and the `use_jax` + kwarg) are `O(N_vis * N_pix)` and can take minutes to hours on a CPU for a + high-resolution mask; they are kept as the reference the NUFFT builder is pinned against. + The result can still be saved to disk and reloaded rather than recomputed on each run — + use `apply_sparse_operator(nufft_precision_operator=...)` to attach a cached result. Parameters ---------- chunk_k - The number of visibilities processed per chunk. Reducing this lowers peak memory - usage during computation at the cost of speed. + The number of visibilities processed per chunk by the brute-force builders. Reducing + this lowers peak memory usage during computation at the cost of speed. show_progress - If `True`, a progress bar is shown during computation. + If `True`, a progress bar is shown during computation by the NumPy brute force. show_memory If `True`, memory usage statistics are printed during computation. use_jax - If `True`, JAX is used to accelerate the computation. + If `True`, the JAX brute-force builder is used (equivalent to `method="jax"`). + method + Which builder computes the operator: `"nufft"` (default), `"numpy"` or `"jax"`. + eps + The requested NUFFT precision of the `"nufft"` builder. `None` takes the transformer's + own `eps` when it is a `TransformerNUFFT`, else `1e-12`. + nufft_chunk_size + The visibility chunk size of the `"nufft"` builder, a memory ceiling rather than an + optimisation. `None` takes the transformer's own `chunk_size` when it is a + `TransformerNUFFT`, else no chunking. Returns ------- @@ -433,11 +465,29 @@ def psf_precision_operator_from( The NUFFT precision matrix of shape (total_pixels, total_pixels) where total_pixels is the number of unmasked real-space pixels. """ + transformer = self.transformer + + # The NUFFT builder and `TransformerNUFFT` spread the same visibilities onto a mode grid + # with the same library, so a dataset that has already chosen an accuracy and a memory + # ceiling for its transformer should not have to repeat them here. + if eps is None: + eps = ( + transformer.eps + if isinstance(transformer, TransformerNUFFT) + else 1.0e-12 + ) + + if nufft_chunk_size is None and isinstance(transformer, TransformerNUFFT): + nufft_chunk_size = transformer.chunk_size + return inversion_interferometer_util.nufft_precision_operator_from( noise_map_real=self.noise_map.array.real, uv_wavelengths=self.uv_wavelengths, - shape_masked_pixels_2d=self.transformer.grid.mask.shape_native_masked_pixels, - grid_radians_2d=self.transformer.grid.mask.derive_grid.all_false.in_radians.native.array, + shape_masked_pixels_2d=transformer.grid.mask.shape_native_masked_pixels, + grid_radians_2d=transformer.grid.mask.derive_grid.all_false.in_radians.native.array, + method=method, + eps=eps, + chunk_size=nufft_chunk_size, chunk_k=chunk_k, show_memory=show_memory, show_progress=show_progress, diff --git a/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py b/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py index c19a7401b..d147c8c10 100644 --- a/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py +++ b/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py @@ -3,6 +3,25 @@ import numpy as np import time from pathlib import Path +from typing import Optional + +from autoarray.operators.transformer import _load_nufftax, nufftax_exception + +try: + from autonerves.test_mode import disable_jax +except ImportError: + # Mirrors the fallback in `autoarray/dataset/interferometer/dataset.py`: + # `disable_jax()` arrives in the autonerves release that closes + # PyAutoNerves#159, and a `--no-deps` install, an editable checkout or a + # hand-built virtualenv can all put an older autonerves on the path + # regardless of the floor in `pyproject.toml`. Degrade to the predicate's + # own one-line body rather than fail at module load; delete the fallback + # when the floor names a release carrying the predicate. + import os + + def disable_jax(): + return os.environ.get("PYAUTO_DISABLE_JAX") == "1" + logger = logging.getLogger(__name__) @@ -83,6 +102,9 @@ def nufft_precision_operator_from( shape_masked_pixels_2d, grid_radians_2d: np.ndarray, *, + method: str = "nufft", + eps: float = 1.0e-12, + chunk_size: Optional[int] = None, chunk_k: int = 2048, show_progress: bool = False, show_memory: bool = False, @@ -209,8 +231,25 @@ def nufft_precision_operator_from( Notes ----- - - If use_jax=True, the JAX implementation is used (requires JAX installed). - - If use_jax=False, the NumPy implementation is used. + Three builders compute the same array; `method` selects which: + + - `"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`. + - `"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. `use_jax=True` is kept for backwards compatibility and maps to this. + + 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. + + Any other `method` raises `ValueError`. Parameters ---------- @@ -225,8 +264,57 @@ def nufft_precision_operator_from( grid_radians_2d The 2D (y,x) grid of coordinates in radians corresponding to real-space mask within which the image that is Fourier transformed is computed. + method + Which builder computes the operator: `"nufft"` (default), `"numpy"` or `"jax"`. + eps + The requested NUFFT precision, used by `method="nufft"` only. + chunk_size + The visibility chunk size of the NUFFT builder (a memory ceiling, not an + optimisation), used by `method="nufft"` only. `None` is one shot. + chunk_k + The visibility chunk size of the two brute-force builders. """ if use_jax: + method = "jax" + + if method not in ("nufft", "numpy", "jax"): + raise ValueError( + f"nufft_precision_operator_from: unknown method {method!r}. " + 'Use "nufft" (the default type-1 NUFFT builder), "numpy" or "jax" (the ' + "brute-force reference builders)." + ) + + if method in ("nufft", "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"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: + 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`." + ) + method = "numpy" + + if method == "nufft": + return nufft_precision_operator_via_nufft_from( + noise_map_real=noise_map_real, + uv_wavelengths=uv_wavelengths, + shape_masked_pixels_2d=shape_masked_pixels_2d, + grid_radians_2d=grid_radians_2d, + eps=eps, + chunk_size=chunk_size, + ) + + if method == "jax": return nufft_precision_operator_via_jax_from( noise_map_real=noise_map_real, uv_wavelengths=uv_wavelengths, @@ -246,6 +334,226 @@ def nufft_precision_operator_from( ) +def _pixel_scale_radians_from(grid_radians_2d: np.ndarray) -> float: + """ + Returns the pixel scale in radians, `delta_rad`, read off the radian grid as an + adjacent-pixel difference. + + It is taken from the grid rather than from `mask.pixel_scales` because the grid is what + the brute-force builders difference to get their `dx` / `dy`: deriving it any other way + would let a unit or half-pixel convention drift in between the two implementations that + have to agree exactly. + + Square pixels are asserted rather than handled. Every interferometer preset is square, + the `[2Ny, 2Nx]` offset grid has a single mode spacing per axis by construction, and a + rectangular-pixel dataset would silently produce a *plausible* wrong operator. + + Parameters + ---------- + grid_radians_2d + The 2D (y,x) native grid of coordinates in radians, shape `[ny, nx, 2]`. + """ + grid_radians_2d = np.asarray(grid_radians_2d, dtype=np.float64) + + if grid_radians_2d.ndim != 3 or grid_radians_2d.shape[-1] != 2: + raise ValueError( + f"grid_radians_2d must be [ny, nx, 2] native; got {grid_radians_2d.shape}." + ) + + n_y, n_x = grid_radians_2d.shape[:2] + + if n_y < 2 or n_x < 2: + raise ValueError( + "grid_radians_2d must be at least 2x2 for the pixel scale to be read off as an " + f"adjacent-pixel difference; got {(n_y, n_x)}." + ) + + # Native y decreases down the rows, x increases along the columns. + delta_y = float(grid_radians_2d[0, 0, 0] - grid_radians_2d[1, 0, 0]) + delta_x = float(grid_radians_2d[0, 1, 1] - grid_radians_2d[0, 0, 1]) + + if not np.isclose(delta_y, delta_x, rtol=1.0e-12, atol=0.0): + raise ValueError( + "The NUFFT precision operator requires square pixels: the radian grid's row " + f"spacing {delta_y!r} and column spacing {delta_x!r} differ." + ) + + return delta_x + + +def nufft_precision_operator_via_nufft_from( + noise_map_real: np.ndarray, + uv_wavelengths: np.ndarray, + shape_masked_pixels_2d, + grid_radians_2d: np.ndarray, + *, + eps: float = 1.0e-12, + chunk_size: Optional[int] = None, +) -> np.ndarray: + """ + Returns the `W~` precision operator built as the real part of a **type-1 (adjoint) + NUFFT**. + + Same signature family and same return value as the brute-force builders + `nufft_precision_operator_via_np_from` / `..._via_jax_from`, but `O(K * nspread^2 + + M log M)` instead of `O(N_pix * K)`, where `M = 4 * Ny * Nx` is the doubled offset + grid. At ALMA scale (`K = 1e6`, `N_pix = 15380`) that is the difference between + ~35 minutes and ~7 seconds. + + The construction + ---------------- + The brute-force builders compute, over the mask's bounding extent + `(Ny, Nx) = shape_native_masked_pixels`: + + P[i, j] = sum_k w_k cos(2 pi (dx * u_k + dy * v_k)), w_k = 1 / sigma_k^2 + + with `dx = -j * delta_rad` and `dy = +i * delta_rad` on autoarray's radian grid + (native `y` *decreases* down the rows, `x` increases along the columns), the four + quadrants filled from the four corners so that offset `0` sits at `[0, 0]` and + negative offsets sit at negative indices (wraparound / FFT ordering), with the middle + row `Ny` and column `Nx` left zero as padding. + + Writing `x_k = 2 pi u_k delta_rad` and `y_k = 2 pi v_k delta_rad` -- the transformer's + own scaled frequencies -- that is exactly + + 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 + `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 + + f = nufft2d1(-x, y, w, n_modes=(2Nx, 2Ny), eps, isign=+1) + P = ifftshift(Re f); P[Ny, :] = 0; P[:, Nx] = 0 + + Why this mapping and not one of the other seven + ----------------------------------------------- + Pinned empirically against `nufft_precision_operator_via_np_from`. Of the eight + candidates (axis swap x sign of `x` x sign of `y`) exactly two agree with the brute + force -- `(-x, +y)` above and `(+x, -y)` -- at `max|delta| = 1.7e-17`, i.e. `8.7e-14` + of the peak `P[0, 0]`. The other six are wrong by `2.1e-1` of the peak, so the + identification is not marginal: the discrimination is thirteen orders of magnitude. + + The two survivors are the *same* construction: `w` is real, so `f(-x, +y)` and + `f(+x, -y)` are complex conjugates and their real parts are identical. `(-x, +y)` is + kept because it reads off the formula above term by term. That degeneracy is also why + `P[i, j] == P[-i, -j]` (cosine evenness) holds -- pinned separately. + + `ifftshift` vs `fftshift` is likewise not a choice here: both axes have even length + `2N`, and for even `N` the two shifts are the same permutation. The canonical + `ifftshift` (centred -> wraparound) is used because that is the direction the transform + actually goes. + + The padding row / column is at index `Ny` / `Nx`, not `Ny - 1` / `Nx - 1`: after + `ifftshift`, index `Ny` carries mode `-Ny`, the Nyquist mode, which the brute force + never evaluates (its quadrants span offsets `-(Ny - 1) ... Ny - 1`). The NUFFT *does* + return a value there, so it is zeroed explicitly. + + Accuracy + -------- + `eps` is the NUFFT's requested precision and its error bound is **peak-scaled**, not + elementwise-relative: a type-1 NUFFT bounds `max|delta|` against `sum_k |c_k|`, so the + near-zero entries of `P` -- five orders below its peak -- carry no relative accuracy + guarantee at all. `eps = 1e-12` saturates fp64 at every instrument profiled: the + measured `max|delta| = 1.7e-17` is already the round-off floor (`eps = 1e-14` only + reaches `1.4e-17`), yet the worst *elementwise* relative error is `6.0e-10` on 32 of + 19600 entries. A pin against this builder must therefore be **mixed** -- + `rtol = 1e-10` with `atol = 1e-10 * P[0, 0]` -- never `rtol` with `atol = 0`. + + Chunking is mandatory at scale + ------------------------------ + `chunk_size` is not an optimisation, it is a memory ceiling -- the same one + `TransformerNUFFT` carries as its own `chunk_size`. The spreader's gather buffer is + `K * nspread^2` complex128; at `eps = 1e-12`, `nspread ~ 14`, so `K = 1e6` needs ~3 GB + and `K = 5e6` needs ~15 GB, which is where an unchunked call on a 15 GB machine is + killed by the OOM reaper rather than returning slowly. Use the instrument's own + transformer chunk size. The transform is linear in the weights, so the chunks' + transforms are summed and the result is the same array (to summation order). + + A caveat on the timings quoted above: they are **CPU** seconds. The saving is in the + algorithm, not the backend -- the brute force is `O(N_pix * K)` whatever it runs on. + + Parameters + ---------- + noise_map_real + `[K]` real noise-map values of the interferometer data; `w = 1 / sigma^2`. + uv_wavelengths + `[K, 2]` `(u, v)` baselines in wavelengths. + shape_masked_pixels_2d + `(Ny, Nx)`, the mask's bounding extent (`mask.shape_native_masked_pixels`). + grid_radians_2d + `[ny, nx, 2]` native `(y, x)` grid in radians. Only its pixel spacing is used, so + the full native grid and the extent sub-grid give the same answer. + eps + The requested NUFFT precision. + chunk_size + Cap on the visibilities passed to `nufft2d1` in one call, or `None` for one shot. + + Returns + ------- + np.ndarray + `[2Ny, 2Nx]` float64, wraparound-ordered, with the padding row / column zero. + """ + nufftax = _load_nufftax() + + if nufftax is None: + nufftax_exception() + + import jax.numpy as jnp + + noise_map_real = np.asarray(noise_map_real, dtype=np.float64) + uv_wavelengths = np.asarray(uv_wavelengths, dtype=np.float64) + grid_radians_2d = np.asarray(grid_radians_2d, dtype=np.float64) + + y_shape, x_shape = (int(s) for s in shape_masked_pixels_2d) + + pixel_scale_radians = _pixel_scale_radians_from(grid_radians_2d) + + # The transformer's own scaled frequencies. + x = 2.0 * np.pi * uv_wavelengths[:, 0] * pixel_scale_radians + y = 2.0 * np.pi * uv_wavelengths[:, 1] * pixel_scale_radians + + w = 1.0 / (noise_map_real**2) + + n_modes = (2 * x_shape, 2 * y_shape) + total_visibilities = int(x.shape[0]) + + if chunk_size is None or chunk_size >= total_visibilities: + chunk_size = total_visibilities + + if chunk_size <= 0: + raise ValueError( + f"chunk_size must be a positive integer or None, got {chunk_size}." + ) + + # Only Re(f) is ever used, so each chunk's real part is accumulated in float64 and the + # complex block is released before the next one is spread. + real_modes = np.zeros((2 * y_shape, 2 * x_shape), dtype=np.float64) + + 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(w[k0:k1], dtype=jnp.complex128), + n_modes, + eps, + 1, + ) + + real_modes += np.asarray(np.real(f), dtype=np.float64) + + del f + + nufft_precision_operator = np.ascontiguousarray(np.fft.ifftshift(real_modes)) + + nufft_precision_operator[y_shape, :] = 0.0 + nufft_precision_operator[:, x_shape] = 0.0 + + return nufft_precision_operator + + def nufft_precision_operator_via_np_from( noise_map_real: np.ndarray, uv_wavelengths: np.ndarray, diff --git a/test_autoarray/dataset/interferometer/test_dataset.py b/test_autoarray/dataset/interferometer/test_dataset.py index b7e1e4766..2b1d33662 100644 --- a/test_autoarray/dataset/interferometer/test_dataset.py +++ b/test_autoarray/dataset/interferometer/test_dataset.py @@ -349,3 +349,130 @@ def dataset(): dataset().apply_sparse_operator(use_jax=True) 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. + 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()) + + monkeypatch.delenv("PYAUTO_DISABLE_JAX", raising=False) + operator_via_numpy = np.asarray( + dataset().psf_precision_operator_from(method="numpy") + ) + + np.testing.assert_array_equal(operator_under_kill_switch, operator_via_numpy) + + +def _interferometer_for_precision_operator(mask_2d_7x7, transformer_class): + n_visibilities = 5 + rng = np.random.default_rng(seed=0) + + return aa.Interferometer( + data=aa.Visibilities( + visibilities=rng.normal(size=(n_visibilities, 2)).astype(np.float64) + ), + noise_map=aa.VisibilitiesNoiseMap( + visibilities=np.ones((n_visibilities, 2), dtype=np.float64) + ), + uv_wavelengths=rng.normal(size=(n_visibilities, 2)).astype(np.float64), + real_space_mask=mask_2d_7x7, + transformer_class=transformer_class, + ) + + +def test__psf_precision_operator_from__nufft_default_matches_the_numpy_brute_force( + mask_2d_7x7, +): + dataset = _interferometer_for_precision_operator( + mask_2d_7x7, transformer.TransformerDFT + ) + + operator_via_numpy = np.asarray(dataset.psf_precision_operator_from(method="numpy")) + operator_default = np.asarray(dataset.psf_precision_operator_from()) + + # Mixed tolerance: a type-1 NUFFT bounds its error against the sum of the weights, so the + # near-zero entries carry no relative accuracy guarantee and need the absolute floor. + np.testing.assert_allclose( + operator_default, + operator_via_numpy, + rtol=1.0e-10, + atol=1.0e-10 * np.abs(operator_via_numpy[0, 0]), + ) + + # `nufft_chunk_size` is a memory ceiling, not an approximation. + np.testing.assert_allclose( + np.asarray(dataset.psf_precision_operator_from(nufft_chunk_size=2)), + operator_default, + rtol=1.0e-10, + atol=1.0e-10 * np.abs(operator_via_numpy[0, 0]), + ) + + +def test__psf_precision_operator_from__eps_and_chunk_size_default_to_the_transformers( + mask_2d_7x7, monkeypatch +): + recorded = [] + + original = aa.util.inversion_interferometer.nufft_precision_operator_from + + def spy(*args, eps, chunk_size, **kwargs): + recorded.append((eps, chunk_size)) + return original(*args, eps=eps, chunk_size=chunk_size, **kwargs) + + monkeypatch.setattr( + aa.util.inversion_interferometer, "nufft_precision_operator_from", spy + ) + + # A `TransformerNUFFT` has already chosen an accuracy and a memory ceiling; the precision + # operator spreads the same visibilities with the same library, so it inherits them. + dataset_nufft = _interferometer_for_precision_operator( + mask_2d_7x7, transformer.TransformerNUFFT + ) + dataset_nufft.transformer.eps = 1.0e-9 + dataset_nufft.transformer.chunk_size = 3 + + dataset_nufft.psf_precision_operator_from() + + assert recorded[-1] == (1.0e-9, 3) + + # A `TransformerDFT` has neither, so the builder's own defaults are used. + _interferometer_for_precision_operator( + mask_2d_7x7, transformer.TransformerDFT + ).psf_precision_operator_from() + + assert recorded[-1] == (1.0e-12, None) + + # An explicit value always wins over both. + dataset_nufft.psf_precision_operator_from(eps=1.0e-11, nufft_chunk_size=4) + + assert recorded[-1] == (1.0e-11, 4) + + +def test__apply_sparse_operator__method_and_nufft_kwargs_reach_the_builder(mask_2d_7x7): + dataset = _interferometer_for_precision_operator( + mask_2d_7x7, transformer.TransformerDFT + ) + + operator_via_numpy = np.asarray(dataset.psf_precision_operator_from(method="numpy")) + + dataset_via_numpy = dataset.apply_sparse_operator(method="numpy") + dataset_default = dataset.apply_sparse_operator() + dataset_chunked = dataset.apply_sparse_operator(nufft_chunk_size=2, eps=1.0e-12) + + # 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), + 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), + rtol=1.0e-10, + atol=1.0e-10 * np.abs(operator_via_numpy[0, 0]), + ) 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 600a6250b..fe125b568 100644 --- a/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py +++ b/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py @@ -68,21 +68,38 @@ def test__data_vector_via_transformed_mapping_matrix_from(): assert (data_vector_complex_via_blurred == data_vector_via_transformed).all() -def _dataset_from(mask, n_visibilities, seed): +def _dataset_from(mask, n_visibilities, seed, uv_scale=1.0, vary_noise=False): """ Returns a small `TransformerDFT` interferometer dataset on the input mask, with `n_visibilities` - seeded random visibilities and unit noise, alongside the random generator used to build it. + seeded random visibilities, alongside the random generator used to build it. + + `uv_scale` multiplies the `(u, v)` baselines. At the default `1.0` the seeded normal baselines + are of order a wavelength, so `2 pi u delta_rad` is ~1e-4 and every phase in the precision + operator is near zero — fine for a shape or plumbing test, useless as a numerical pin. A scale + near `0.5 / delta_rad` (~1e5 for arcsecond pixels) puts the phases in `[-pi, pi]`, which is + where the operator's entries actually vary and where a wrong sign convention is visible. + + `vary_noise` gives each visibility its own sigma, so the weights `w = 1 / sigma^2` are not all + equal and a builder that dropped them would be caught. """ rng = np.random.default_rng(seed=seed) + if vary_noise: + sigma = rng.uniform(0.5, 1.5, size=n_visibilities).astype(np.float64) + noise_map = np.stack([sigma, sigma], axis=1) + else: + noise_map = np.ones((n_visibilities, 2), dtype=np.float64) + dataset = aa.Interferometer( data=aa.Visibilities( visibilities=rng.normal(size=(n_visibilities, 2)).astype(np.float64) ), - noise_map=aa.VisibilitiesNoiseMap( - visibilities=np.ones((n_visibilities, 2), dtype=np.float64) - ), - uv_wavelengths=rng.normal(size=(n_visibilities, 2)).astype(np.float64), + noise_map=aa.VisibilitiesNoiseMap(visibilities=noise_map), + uv_wavelengths=( + uv_scale * rng.uniform(-1.0, 1.0, size=(n_visibilities, 2)) + if uv_scale != 1.0 + else rng.normal(size=(n_visibilities, 2)) + ).astype(np.float64), real_space_mask=mask, transformer_class=aa.TransformerDFT, ) @@ -328,3 +345,334 @@ def test__interferometer_sparse_operator__apply_operator__rfft2_matches_complex_ rtol=1.0e-10, atol=1.0e-10 * np.abs(operated_via_complex_fft2).max(), ) + + +def _preload_inputs_from(dataset): + """ + Returns the four arguments `Interferometer.psf_precision_operator_from` passes to every NUFFT + precision operator builder, read off a dataset. + + The pins below compare builders, so they must be handed exactly the arrays the library hands + them; deriving any of the four differently here would let a convention drift in between the + implementations the pins exist to hold together. + """ + mask = dataset.transformer.grid.mask + + return { + "noise_map_real": np.asarray(dataset.noise_map.array.real, dtype=np.float64), + "uv_wavelengths": np.asarray(dataset.uv_wavelengths, dtype=np.float64), + "shape_masked_pixels_2d": mask.shape_native_masked_pixels, + "grid_radians_2d": np.asarray( + mask.derive_grid.all_false.in_radians.native.array, dtype=np.float64 + ), + } + + +def _nufft_pin_inputs_7x7(): + """ + The shared 7x7 / K=5 fixture the rest of this module's operator tests are built from. + """ + dataset, _ = _dataset_from(mask=_mask_7x7(), n_visibilities=5, seed=3) + + return _preload_inputs_from(dataset) + + +def _nufft_pin_inputs_16x16(): + """ + A seeded 16x16 circular mask with K=300, varying noise sigmas and baselines scaled so the + phases span `[-pi, pi]`. + + The 7x7 fixture alone is a weak pin: its five normal baselines give phases of order 1e-4, so + every entry of the operator sits within 1e-8 of the peak and a wrong sign convention would be + almost invisible. On this fixture the entries vary over the full range and the wrong-sign + control below is off by 18% of the peak. + """ + dataset, _ = _dataset_from( + mask=aa.Mask2D.circular(shape_native=(16, 16), pixel_scales=1.0, radius=6.0), + n_visibilities=300, + seed=7, + uv_scale=1.0e5, + vary_noise=True, + ) + + return _preload_inputs_from(dataset) + + +def _assert_matches_brute_force(operator, operator_via_np): + """ + Asserts that a NUFFT-built precision operator matches the NumPy brute-force builder. + + The tolerance is deliberately **mixed**. A type-1 NUFFT bounds its error against `sum_k |c_k|`, + i.e. peak-scaled, not elementwise-relative: the near-zero entries of the operator sit orders + below its peak and carry no relative accuracy guarantee at all, so a pure `rtol` pin on them + would be measuring fp64 round-off rather than the builder. `atol = 1e-10 * P[0, 0]` puts a + floor under exactly those entries while every entry carrying signal stays under a full + relative test. Never `rtol` with `atol = 0`. + """ + np.testing.assert_allclose( + operator, + operator_via_np, + rtol=1.0e-10, + atol=1.0e-10 * np.abs(operator_via_np[0, 0]), + ) + + +def test__nufft_precision_operator_via_nufft__matches_the_numpy_brute_force(): + pytest.importorskip("nufftax") + + for inputs in (_nufft_pin_inputs_7x7(), _nufft_pin_inputs_16x16()): + operator_via_np = np.asarray( + aa.util.inversion_interferometer.nufft_precision_operator_via_np_from( + **inputs + ) + ) + operator_via_nufft = np.asarray( + aa.util.inversion_interferometer.nufft_precision_operator_via_nufft_from( + **inputs + ) + ) + + assert operator_via_nufft.shape == operator_via_np.shape + assert operator_via_nufft.dtype == np.float64 + + _assert_matches_brute_force(operator_via_nufft, operator_via_np) + + +def test__nufft_precision_operator_via_nufft__nyquist_row_and_column_are_zero(): + pytest.importorskip("nufftax") + + for inputs in (_nufft_pin_inputs_7x7(), _nufft_pin_inputs_16x16()): + y_shape, x_shape = (int(s) for s in inputs["shape_masked_pixels_2d"]) + + operator = np.asarray( + aa.util.inversion_interferometer.nufft_precision_operator_via_nufft_from( + **inputs + ) + ) + + # After `ifftshift`, index `y_shape` / `x_shape` carries the Nyquist mode, which the brute + # force never evaluates (its quadrants span offsets -(N-1) ... N-1). The NUFFT does return + # a value there, so the builder must zero it explicitly -- exactly, not approximately. + assert (operator[y_shape, :] == 0.0).all() + assert (operator[:, x_shape] == 0.0).all() + + # The neighbours are not zero, so the assertion above is testing the padding and not an + # operator that came back empty. + assert np.abs(operator[y_shape - 1, :]).max() > 0.0 + assert np.abs(operator[y_shape + 1, :]).max() > 0.0 + assert np.abs(operator[:, x_shape - 1]).max() > 0.0 + assert np.abs(operator[:, x_shape + 1]).max() > 0.0 + + +def test__nufft_precision_operator_via_nufft__is_even_under_negated_offsets(): + pytest.importorskip("nufftax") + + for inputs in (_nufft_pin_inputs_7x7(), _nufft_pin_inputs_16x16()): + y_shape, x_shape = (int(s) for s in inputs["shape_masked_pixels_2d"]) + + operator = np.asarray( + aa.util.inversion_interferometer.nufft_precision_operator_via_nufft_from( + **inputs + ) + ) + + # `P[i, j] = sum_k w_k cos(...)` is even in the offset, so `P[i, j] == P[-i, -j]`. The + # wraparound ordering means the negated offset is a plain negative index. + for i in range(-(y_shape - 1), y_shape): + for j in range(-(x_shape - 1), x_shape): + assert operator[i, j] == pytest.approx(operator[-i, -j], rel=1.0e-12) + + +def test__nufft_precision_operator_via_nufft__chunked_matches_one_shot(): + pytest.importorskip("nufftax") + + inputs = _nufft_pin_inputs_16x16() + + operator_via_np = np.asarray( + aa.util.inversion_interferometer.nufft_precision_operator_via_np_from(**inputs) + ) + one_shot = np.asarray( + aa.util.inversion_interferometer.nufft_precision_operator_via_nufft_from( + **inputs + ) + ) + # K = 300, so a chunk size of 64 spreads five chunks and sums their transforms. Chunking is a + # memory ceiling, not an approximation -- the transform is linear in the weights -- so the two + # agree to summation order, well inside the pin the builder is held to. + chunked = np.asarray( + aa.util.inversion_interferometer.nufft_precision_operator_via_nufft_from( + **inputs, chunk_size=64 + ) + ) + + _assert_matches_brute_force(chunked, operator_via_np) + + np.testing.assert_allclose( + chunked, + one_shot, + rtol=1.0e-10, + atol=1.0e-10 * np.abs(operator_via_np[0, 0]), + ) + + +def test__nufft_precision_operator_via_nufft__negated_u_fails_the_pin(): + pytest.importorskip("nufftax") + + inputs = _nufft_pin_inputs_16x16() + + operator_via_np = np.asarray( + aa.util.inversion_interferometer.nufft_precision_operator_via_np_from(**inputs) + ) + + inputs_wrong_sign = dict(inputs) + inputs_wrong_sign["uv_wavelengths"] = inputs["uv_wavelengths"].copy() + inputs_wrong_sign["uv_wavelengths"][:, 0] *= -1.0 + + operator_wrong_sign = np.asarray( + aa.util.inversion_interferometer.nufft_precision_operator_via_nufft_from( + **inputs_wrong_sign + ) + ) + + # The control: of the eight candidate mappings (axis swap x sign of x x sign of y) only two + # agree with the brute force, and they are the same construction. Negating `u` selects one of + # the six wrong ones, which must fail the pin loudly -- if it passed, the pin would be + # measuring nothing about the sign convention. + with pytest.raises(AssertionError): + _assert_matches_brute_force(operator_wrong_sign, operator_via_np) + + +def test__nufft_precision_operator_from__method_routes_to_each_builder(): + pytest.importorskip("nufftax") + + inputs = _nufft_pin_inputs_16x16() + + operator_via_np = np.asarray( + aa.util.inversion_interferometer.nufft_precision_operator_via_np_from(**inputs) + ) + operator_via_jax = np.asarray( + aa.util.inversion_interferometer.nufft_precision_operator_via_jax_from(**inputs) + ) + operator_via_nufft = np.asarray( + aa.util.inversion_interferometer.nufft_precision_operator_via_nufft_from( + **inputs + ) + ) + + # `"numpy"` and `"jax"` are pure routing, so they return the brute-force arrays unchanged. + np.testing.assert_array_equal( + np.asarray( + aa.util.inversion_interferometer.nufft_precision_operator_from( + method="numpy", **inputs + ) + ), + operator_via_np, + ) + np.testing.assert_array_equal( + np.asarray( + aa.util.inversion_interferometer.nufft_precision_operator_from( + method="jax", **inputs + ) + ), + operator_via_jax, + ) + + # `use_jax=True` is kept for backwards compatibility and maps onto `method="jax"`. + np.testing.assert_array_equal( + np.asarray( + aa.util.inversion_interferometer.nufft_precision_operator_from( + use_jax=True, **inputs + ) + ), + operator_via_jax, + ) + + # The default is the NUFFT builder, and it agrees with the brute force. + operator_default = np.asarray( + aa.util.inversion_interferometer.nufft_precision_operator_from(**inputs) + ) + + np.testing.assert_array_equal(operator_default, operator_via_nufft) + _assert_matches_brute_force(operator_default, operator_via_np) + + +def test__nufft_precision_operator_from__unknown_method_raises(): + inputs = _nufft_pin_inputs_7x7() + + with pytest.raises(ValueError): + aa.util.inversion_interferometer.nufft_precision_operator_from( + method="type-1", **inputs + ) + + +def test__nufft_precision_operator_from__disable_jax_falls_back_to_the_numpy_builder( + monkeypatch, caplog +): + inputs = _nufft_pin_inputs_7x7() + + operator_via_np = np.asarray( + 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). + monkeypatch.setenv("PYAUTO_DISABLE_JAX", "1") + + for kwargs in ({}, {"method": "jax"}, {"use_jax": True}): + caplog.clear() + + with caplog.at_level("WARNING"): + operator = np.asarray( + aa.util.inversion_interferometer.nufft_precision_operator_from( + **kwargs, **inputs + ) + ) + + np.testing.assert_array_equal(operator, operator_via_np) + assert "PYAUTO_DISABLE_JAX" in caplog.text + + # 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") + + 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 + ) + ), + ) + + +def test__nufft_precision_operator_from__nufftax_absent_falls_back_to_the_numpy_builder( + monkeypatch, caplog +): + inputs = _nufft_pin_inputs_7x7() + + operator_via_np = np.asarray( + 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. + 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) + ) + + np.testing.assert_array_equal(operator, operator_via_np) + assert "nufftax" in caplog.text + + # 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 + ) From 61f07fd832454dab5209e5085f43c11364bf7959 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Tue, 8 Sep 2026 13:24:38 -0400 Subject: [PATCH 2/2] fix: use_jax no longer demotes the NUFFT preload to the JAX brute force (#539) Every existing workspace `apply_sparse_operator(use_jax=True)` call was mapped onto `method="jax"`, silently putting the production path on the brute force; `use_jax` is now only honoured when a brute-force method is already selected. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_018hLF3ZAcz5MmaSJBEcLkvF --- autoarray/dataset/interferometer/dataset.py | 33 ++++++++++++------- .../inversion_interferometer_util.py | 20 ++++++++--- .../test_inversion_interferometer_util.py | 32 +++++++++++++++--- 3 files changed, 65 insertions(+), 20 deletions(-) diff --git a/autoarray/dataset/interferometer/dataset.py b/autoarray/dataset/interferometer/dataset.py index 33b5035f4..b36a593dd 100644 --- a/autoarray/dataset/interferometer/dataset.py +++ b/autoarray/dataset/interferometer/dataset.py @@ -251,10 +251,10 @@ def apply_sparse_operator( 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 - visibilities. The brute-force builders (`method="numpy"` / `"jax"`, and the `use_jax` kwarg) - are `O(N_vis * N_pix)` and can take minutes to hours; they are kept as the reference the - NUFFT builder is pinned against. Either way the result can be cached to disk and reloaded - via `nufft_precision_operator=`. + visibilities. The brute-force builders (`method="numpy"` / `"jax"`) are `O(N_vis * N_pix)` + and can take minutes to hours; they are kept as the reference the NUFFT builder is pinned + against. Either way the result can be cached to disk and reloaded via + `nufft_precision_operator=`. Parameters ---------- @@ -284,7 +284,11 @@ def apply_sparse_operator( show_memory If `True`, memory usage statistics are printed while computing the NUFFT precision matrix. use_jax - If `True`, JAX is used to accelerate the NUFFT precision matrix computation. + 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. `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 @@ -292,7 +296,9 @@ def apply_sparse_operator( and the smoke profiles set it so a fast run does not pay a JIT compile. An 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 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_precision_operator_from`.) Precondition ------------ @@ -317,6 +323,9 @@ def apply_sparse_operator( If any visibility has unequal real and imaginary noise sigma. """ + # `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. if disable_jax(): use_jax = False @@ -357,7 +366,7 @@ def apply_sparse_operator( n_vis = self.uv_wavelengths.shape[0] n_pix = self.real_space_mask.pixels_in_mask - if method != "nufft" or use_jax: + if method != "nufft": logger.info( f"INTERFEROMETER - The precision operator is being built by a brute-force " f"builder, which is O(N_vis x N_pix) = O({n_vis * n_pix:.1e}) and can take " @@ -432,9 +441,9 @@ def psf_precision_operator_from( The default builder (`method="nufft"`) computes this as a type-1 (adjoint) NUFFT, which is `O(N_vis * nspread^2 + M log M)` for `M = 4 * Ny * Nx` — seconds even at a million - visibilities. The brute-force builders (`method="numpy"` / `"jax"`, and the `use_jax` - kwarg) are `O(N_vis * N_pix)` and can take minutes to hours on a CPU for a - high-resolution mask; they are kept as the reference the NUFFT builder is pinned against. + visibilities. The brute-force builders (`method="numpy"` / `"jax"`) are + `O(N_vis * N_pix)` and can take minutes to hours on a CPU for a high-resolution mask; + they are kept as the reference the NUFFT builder is pinned against. The result can still be saved to disk and reloaded rather than recomputed on each run — use `apply_sparse_operator(nufft_precision_operator=...)` to attach a cached result. @@ -448,7 +457,9 @@ def psf_precision_operator_from( show_memory If `True`, memory usage statistics are printed during computation. use_jax - If `True`, the JAX brute-force builder is used (equivalent to `method="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. 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 d147c8c10..c5e8f195d 100644 --- a/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py +++ b/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py @@ -122,15 +122,15 @@ def nufft_precision_operator_from( matrix construction without performing a NUFFT per source pixel. ------------------------------------------------------------------------------- - Backend behaviour + Backend behaviour (the two brute-force builders) ------------------------------------------------------------------------------- - - NumPy backend (use_jax=False, default): + - NumPy backend (`method="numpy"`): * CPU execution * Explicit Python loop over visibility chunks * Supports progress bars and optional memory reporting * Numerically closest to the original reference implementation - - JAX backend (use_jax=True): + - JAX backend (`method="jax"`, or `method="numpy"` with `use_jax=True`): * JIT-compilable and GPU/TPU capable * Uses fixed-size chunking and lax.fori_loop * No Python-side loops during execution @@ -240,7 +240,13 @@ def nufft_precision_operator_from( `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. `use_jax=True` is kept for backwards compatibility and maps to this. + 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. 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)`: @@ -273,8 +279,12 @@ def nufft_precision_operator_from( optimisation), used by `method="nufft"` only. `None` is one shot. 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. """ - if use_jax: + if method == "numpy" and use_jax: method = "jax" if method not in ("nufft", "numpy", "jax"): 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 fe125b568..26cba69a4 100644 --- a/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py +++ b/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py @@ -577,16 +577,33 @@ def test__nufft_precision_operator_from__method_routes_to_each_builder(): operator_via_jax, ) - # `use_jax=True` is kept for backwards compatibility and maps onto `method="jax"`. + # `use_jax=True` only upgrades a brute-force method, so `method="numpy"` with it set is + # the JAX brute force. np.testing.assert_array_equal( np.asarray( aa.util.inversion_interferometer.nufft_precision_operator_from( - use_jax=True, **inputs + method="numpy", use_jax=True, **inputs ) ), 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. + 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) + _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. operator_default = np.asarray( aa.util.inversion_interferometer.nufft_precision_operator_from(**inputs) @@ -616,10 +633,17 @@ def test__nufft_precision_operator_from__disable_jax_falls_back_to_the_numpy_bui # `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). + # 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. monkeypatch.setenv("PYAUTO_DISABLE_JAX", "1") - for kwargs in ({}, {"method": "jax"}, {"use_jax": True}): + for kwargs in ( + {}, + {"method": "jax"}, + {"use_jax": True}, + {"method": "numpy", "use_jax": True}, + ): caplog.clear() with caplog.at_level("WARNING"):