From c2469b9d0c0db3f6ad6b8a3f2fc4a03e8f37450d Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Tue, 8 Sep 2026 14:39:45 -0400 Subject: [PATCH] perf: NumPy/scipy application path for InterferometerSparseOperator (#542) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The operator now keeps the raw (2y, 2x) real preload as `nufft_precision_operator` and lazily builds both transforms of it: `Khat` (JAX `rfft2`, under `jax.ensure_compile_time_eval()`) and `khat_np` (`scipy.fft.rfft2`). An operator constructed and applied with `xp=np` therefore never imports JAX. All six public methods (`apply_operator`, `curvature_matrix_diag_from`, `curvature_matrix_off_diag_from`, `operated_matrix_slim_from`, `curvature_matrix_off_diag_func_list_from`, `curvature_matrix_func_list_from`) take an `xp` array module and branch: `jax.numpy` keeps the existing `lax.fori_loop` / `segment_sum` / `dynamic_update_slice` bodies unchanged, `numpy` (the default) runs scipy `rfft2`/`irfft2` with a `scipy.sparse` CSC projection and a plain Python block loop (no padding or masking, which exist only to keep traced shapes static). The diagonal NumPy body symmetrises like the JAX one so the two agree to the pin. `InversionInterferometerSparse` passes `xp=self._xp` to every operator call, so the application backend follows the inversion's `xp` — what the fit passes — rather than the dataset's `use_jax` kwarg, which after #541 only selects a brute-force *builder* for the preload. The `np.array(curvature_matrix)` workaround at sparse.py:171 (needed only because the operator used to return JAX arrays on a NumPy fit) is deleted. Delaunay's padded `col = -1` triplets are dropped explicitly when building the CSC matrix. Verified: parity NumPy vs JAX at <= 5e-16 relative on all six methods (7x7/K=5, 12x12/K=64 seeded, Delaunay duplicate-COO, and a `batch_size=4` case exercising the block loop); end-to-end inversion curvature parity 2.8e-14; a subprocess test asserts `"jax" not in sys.modules` after a full NumPy fit; NumPy `curvature_matrix_diag_from` is 3.8x faster than JAX-CPU on a 40x40 / S=400 probe. 1486 tests pass. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_018hLF3ZAcz5MmaSJBEcLkvF --- autoarray/dataset/interferometer/dataset.py | 4 + .../inversion_interferometer_util.py | 366 +++++++++++++-- .../inversion/interferometer/sparse.py | 10 +- .../interferometer/test_interferometer.py | 158 +++++++ .../test_inversion_interferometer_util.py | 425 ++++++++++++++++++ 5 files changed, 926 insertions(+), 37 deletions(-) diff --git a/autoarray/dataset/interferometer/dataset.py b/autoarray/dataset/interferometer/dataset.py index b36a593dd..35879a31f 100644 --- a/autoarray/dataset/interferometer/dataset.py +++ b/autoarray/dataset/interferometer/dataset.py @@ -249,6 +249,10 @@ def apply_sparse_operator( and is used automatically by `FitInterferometer` when performing pixelized reconstructions via 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. + 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"`) are `O(N_vis * N_pix)` diff --git a/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py b/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py index c5e8f195d..f447bef75 100644 --- a/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py +++ b/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py @@ -1,4 +1,5 @@ from dataclasses import dataclass +from functools import cached_property import logging import numpy as np import time @@ -875,14 +876,13 @@ class InterferometerSparseOperator: `DatasetException` when it is violated. """ + nufft_precision_operator: np.ndarray # (2y, 2x) real preload dirty_image: np.ndarray y_shape: int x_shape: int M: int batch_size: int - w_dtype: "jax.numpy.dtype" - Khat: "jax.Array" # (2y, x+1), rfft2 of the real preload - col_offsets: "jax.Array" # (batch_size,) int32 + w_dtype: np.dtype """ Cached FFT operator state for fast interferometer curvature-matrix assembly. @@ -917,8 +917,24 @@ class InterferometerSparseOperator: - dtype / precision (float32 vs float64) - `batch_size` + Backends + -------- + Every public method takes an `xp` array module and branches on it: `xp=jnp` runs the + JAX bodies (`lax.fori_loop`, `segment_sum`, `dynamic_update_slice`), `xp=np` (the + default) runs NumPy/scipy bodies (`scipy.fft`, `scipy.sparse`, a Python block loop). + The two agree to floating-point round-off, and the backend follows the *inversion's* + `xp` — the dataset's `use_jax` kwarg only selects a brute-force builder for the + preload, never the application path. + + The JAX-only state (`Khat`, `col_offsets`) is therefore built lazily, so an operator + constructed and applied under `xp=np` never imports JAX. + Parameters stored ----------------- + nufft_precision_operator + The raw real-valued (2y_shape, 2x_shape) preload the operator is built from, kept + (rather than discarded after the FFT) because the NumPy and numba CPU paths index + it directly and it is small. dirty_image Convenience field for associated dirty image data (not used directly in curvature assembly in this method). Stored as a NumPy array to match @@ -932,11 +948,71 @@ class InterferometerSparseOperator: Larger batch sizes improve throughput on GPU but increase memory usage. w_dtype Floating-point dtype for weights and accumulations (e.g. float64). + + Cached properties + ----------------- Khat - Real FFT of the curvature preload, shape (2y_shape, x_shape + 1), complex. - This is the frequency-domain representation of the W~ operator kernel. + Real FFT of the curvature preload, shape (2y_shape, x_shape + 1), complex, as a + JAX array. The frequency-domain representation of the W~ operator kernel used by + the JAX branch. Computed on first access, so it never imports JAX for a NumPy run. + khat_np + The same transform via `scipy.fft.rfft2`, used by the NumPy branch. + col_offsets + `(batch_size,)` int32 JAX column indices used to mask the final partial block of + the JAX branch's `fori_loop`. """ + @cached_property + def Khat(self): + """ + The `rfft2` of the preload as a JAX array, of shape (2y_shape, x_shape + 1). + + Lazy so that constructing the operator, and applying it with `xp=np`, never + imports JAX. `functools.cached_property` writes straight into the instance + `__dict__`, so it is compatible with `frozen=True` (which only blocks + `__setattr__`) and does not participate in the dataclass's fields, `__eq__` or + `__hash__`. + + `jax.ensure_compile_time_eval()` is what makes the laziness safe. Under a trace + (`jit`, and the `lax.fori_loop` the JAX curvature bodies run this transform + inside) **every** `jax.numpy` call is staged into the jaxpr, even one whose + operands are all concrete -- so a first access from inside the loop body would + cache a `DynamicJaxprTracer` here and every later call would fail with + `UnexpectedTracerError`. Inside the context the transform is evaluated eagerly + and comes out a concrete `jax.Array`, which is both correct to cache and a + compile-time constant for the trace that asked for it. + """ + import jax + import jax.numpy as jnp + + with jax.ensure_compile_time_eval(): + return jnp.fft.rfft2(self.nufft_precision_operator) + + @cached_property + def khat_np(self) -> np.ndarray: + """ + The `rfft2` of the preload as a NumPy array, of shape (2y_shape, x_shape + 1). + + `scipy.fft` rather than `numpy.fft`: it is the transform the prototype measured + (autolens_profiling #226) and scipy is already a hard dependency. + """ + import scipy.fft + + return scipy.fft.rfft2(self.nufft_precision_operator) + + @cached_property + def col_offsets(self): + """ + `(batch_size,)` int32 JAX column indices, used by the JAX branch to zero the + columns of the final partial block. Lazy, and evaluated eagerly, for the same + reason as `Khat`. + """ + import jax + import jax.numpy as jnp + + with jax.ensure_compile_time_eval(): + return jnp.arange(int(self.batch_size), dtype=jnp.int32) + @classmethod def from_nufft_precision_operator( cls, @@ -952,9 +1028,11 @@ def from_nufft_precision_operator( The curvature preload is assumed to be defined on a (2y, 2x) rectangular grid of pixel offsets, where y and x correspond to the *unmasked extent* - of the real-space grid. The preload is real, so it is transformed once with - a real FFT (`rfft2`) to obtain `Khat` of shape (2y, x + 1), which is then - reused for every subsequent curvature matrix build. + of the real-space grid. The preload is real, so it is transformed with a real + FFT (`rfft2`) to obtain a kernel of shape (2y, x + 1), which is then reused for + every subsequent curvature matrix build. That transform is a cached property + (`Khat` on JAX, `khat_np` on NumPy), computed on first use rather than here, so + constructing the operator imports no backend it is not asked for. Parameters ---------- @@ -974,16 +1052,14 @@ def from_nufft_precision_operator( Returns ------- InterferometerSparseOperator - Immutable cached state object containing shapes and FFT kernel `Khat`, - of shape (2y, x + 1) and complex dtype. + Immutable cached state object containing the preload, the grid shapes and the + lazily built FFT kernels of shape (2y, x + 1) and complex dtype. Raises ------ ValueError If `nufft_precision_operator` does not have even shape in both dimensions. """ - import jax.numpy as jnp - H2, W2 = nufft_precision_operator.shape if (H2 % 2) != 0 or (W2 % 2) != 0: raise ValueError( @@ -994,20 +1070,33 @@ def from_nufft_precision_operator( x_shape = W2 // 2 M = y_shape * x_shape - Khat = jnp.fft.rfft2(nufft_precision_operator) + # C-contiguous NumPy, so the NumPy and numba CPU paths can index the preload + # directly. This is a no-op (the same object) for an array that is already a + # C-contiguous NumPy array, and materialises a JAX-built preload onto the host. + nufft_precision_operator = np.ascontiguousarray(nufft_precision_operator) return InterferometerSparseOperator( + nufft_precision_operator=nufft_precision_operator, dirty_image=dirty_image, y_shape=y_shape, x_shape=x_shape, M=M, batch_size=int(batch_size), w_dtype=nufft_precision_operator.dtype, - Khat=Khat, - col_offsets=jnp.arange(int(batch_size), dtype=jnp.int32), ) - def apply_operator(self, Fbatch_flat): + @staticmethod + def _is_jax(xp) -> bool: + """ + Returns `True` if `xp` is the JAX array module. + + Every public method branches on this rather than on `xp is np`, so that any + NumPy-API-compatible module (`numpy`, and the `numpy` re-exports the library + passes around) takes the NumPy branch, while `jax.numpy` takes the JAX branch. + """ + return xp.__name__.startswith("jax") + + def apply_operator(self, Fbatch_flat, xp=np): """ Apply the interferometer W~ operator to a batch of vectors. @@ -1035,28 +1124,51 @@ def apply_operator(self, Fbatch_flat): ---------- Fbatch_flat Array of shape (M, B) representing B vectors on the rectangular grid. + xp + The array module the operator is applied with: `numpy` (default) uses + `scipy.fft`, `jax.numpy` uses `jax.numpy.fft`. Returns ------- ndarray Array of shape (M, B) equal to W~ applied to the batch. """ - import jax.numpy as jnp - y_shape, x_shape = self.y_shape, self.x_shape M = y_shape * x_shape - Khat = self.Khat + + if self._is_jax(xp): + import jax.numpy as jnp + + Khat = self.Khat + + B = Fbatch_flat.shape[1] + F_img = Fbatch_flat.T.reshape((B, y_shape, x_shape)) + F_pad = jnp.pad(F_img, ((0, 0), (0, y_shape), (0, x_shape))) + Fhat = jnp.fft.rfft2(F_pad) + Ghat = Fhat * Khat[None, :, :] + G_pad = jnp.fft.irfft2(Ghat, s=(2 * y_shape, 2 * x_shape)) + G = G_pad[:, :y_shape, :x_shape] + return G.reshape((B, M)).T + + import scipy.fft + + Fbatch_flat = np.asarray(Fbatch_flat, dtype=np.float64) B = Fbatch_flat.shape[1] - F_img = Fbatch_flat.T.reshape((B, y_shape, x_shape)) - F_pad = jnp.pad(F_img, ((0, 0), (0, y_shape), (0, x_shape))) - Fhat = jnp.fft.rfft2(F_pad) - Ghat = Fhat * Khat[None, :, :] - G_pad = jnp.fft.irfft2(Ghat, s=(2 * y_shape, 2 * x_shape)) + + # The pad is written as a zeroed buffer with the block copied into its top-left + # corner rather than as `np.pad`, so only one (B, 2y, 2x) array is allocated. + F_pad = np.zeros((B, 2 * y_shape, 2 * x_shape), dtype=np.float64) + F_pad[:, :y_shape, :x_shape] = Fbatch_flat.T.reshape((B, y_shape, x_shape)) + + Fhat = scipy.fft.rfft2(F_pad, axes=(-2, -1)) + Fhat *= self.khat_np[None, :, :] + G_pad = scipy.fft.irfft2(Fhat, s=(2 * y_shape, 2 * x_shape), axes=(-2, -1)) + G = G_pad[:, :y_shape, :x_shape] return G.reshape((B, M)).T - def curvature_matrix_diag_from(self, rows, cols, vals, *, S: int): + def curvature_matrix_diag_from(self, rows, cols, vals, *, S: int, xp=np): """ Compute the diagonal (mapper-mapper) curvature matrix block F = Aᵀ W~ A. @@ -1085,12 +1197,21 @@ def curvature_matrix_diag_from(self, rows, cols, vals, *, S: int): These should already be produced by `mapper.sparse_triplets_curvature`. S Number of source pixels / parameters for this mapper. + xp + The array module the block assembly runs on: `numpy` (default) assembles `A` + as a `scipy.sparse` CSC matrix and loops the blocks in Python, `jax.numpy` + runs the `fori_loop` / `segment_sum` route. Returns ------- ndarray Curvature matrix of shape (S, S), symmetric. """ + if not self._is_jax(xp): + return self._curvature_matrix_diag_from_np( + rows=rows, cols=cols, vals=vals, S=S + ) + import jax.numpy as jnp from jax import lax from jax.ops import segment_sum @@ -1117,7 +1238,7 @@ def body(block_i, C): F = jnp.zeros((M, B), dtype=jnp.float64) F = F.at[rows, bc].add(v) - G = self.apply_operator(F) # (M, B) + G = self.apply_operator(F, xp=jnp) # (M, B) contrib = vals[:, None] * G[rows, :] Cblock = segment_sum(contrib, cols, num_segments=S) # (S, B) @@ -1131,8 +1252,84 @@ def body(block_i, C): C = C_pad[:, :S] return 0.5 * (C + C.T) + def _sparse_matrix_from(self, rows, cols, vals, *, S: int): + """ + Returns the mapping operator `A` of shape (M, S) as a `scipy.sparse` CSC matrix, + built from its COO triplets. + + `csc_matrix((vals, (rows, cols)))` *sums* duplicate `(row, col)` entries, which is + required: a Delaunay mapper produces one triplet per (sub-pixel, vertex) pair, so + several triplets can share a cell, exactly as the JAX branch's `.at[].add` scatter + accumulates them. + + Out-of-range columns are dropped. `mapper_util.sparse_triplets_from` pads every + sub-pixel's interpolation stencil to the longest one, so an unused slot arrives as + `col = -1` with `val = 0.0`. The JAX branch drops those implicitly (its `in_block` + mask is false for a negative column, and `segment_sum` drops out-of-bound segment + ids), whereas `csc_matrix` raises on a negative index — so the same entries are + dropped explicitly here. They carry zero weight, so dropping them changes nothing + but the error. + + Parameters + ---------- + rows, cols, vals + COO triplets encoding `A`: extent-grid flat row indices, source pixel column + indices, and mapping weights. + S + Number of source pixels / parameters, i.e. the column count of `A`. + """ + from scipy.sparse import csc_matrix + + S = int(S) + + rows = np.asarray(rows, dtype=np.int64) + cols = np.asarray(cols, dtype=np.int64) + vals = np.asarray(vals, dtype=np.float64) + + in_range = (cols >= 0) & (cols < S) + + if not in_range.all(): + rows = rows[in_range] + cols = cols[in_range] + vals = vals[in_range] + + return csc_matrix((vals, (rows, cols)), shape=(self.M, S)) + + def _curvature_matrix_diag_from_np(self, rows, cols, vals, *, S: int) -> np.ndarray: + """ + The NumPy/scipy body of `curvature_matrix_diag_from`: `F = Aᵀ W~ A`. + + `A` is assembled once as a CSC matrix, then swept in column blocks of + `batch_size`: each block is densified, operated on with `apply_operator`, and + projected back with `Aᵀ` as a sparse-dense product (the NumPy counterpart of the + JAX branch's `segment_sum`). + + The block loop is a plain Python `for` rather than the JAX branch's `fori_loop`, + so no padding, masking or `dynamic_update_slice` is needed: those exist purely to + keep the traced shapes static. + """ + A = self._sparse_matrix_from(rows=rows, cols=cols, vals=vals, S=S) + AT = A.T.tocsr() + + S = int(S) + B = int(self.batch_size) + + C = np.zeros((S, S), dtype=np.float64) + + for start in range(0, S, B): + stop = min(start + B, S) + + F = A[:, start:stop].toarray() + G = self.apply_operator(F, xp=np) + + C[:, start:stop] = AT @ G + + # The JAX branch symmetrises its result, so this one must too or the two disagree + # at round-off in the last digits rather than to the pin the tests hold. + return 0.5 * (C + C.T) + def curvature_matrix_off_diag_from( - self, rows0, cols0, vals0, rows1, cols1, vals1, *, S0: int, S1: int + self, rows0, cols0, vals0, rows1, cols1, vals1, *, S0: int, S1: int, xp=np ): """ Compute the off-diagonal (mapper-mapper) curvature block F01 = A0ᵀ W~ A1. @@ -1165,6 +1362,8 @@ def curvature_matrix_off_diag_from( Number of source pixels / parameters for mapper 0. S1 Number of source pixels / parameters for mapper 1. + xp + The array module the block assembly runs on: `numpy` (default) or `jax.numpy`. Returns ------- @@ -1175,8 +1374,21 @@ def curvature_matrix_off_diag_from( ----- - The result is *not* symmetrized here because it is not square in general. The symmetric counterpart is F10 = F01ᵀ, because A0 and A1 share the same W~. - - Padding to `S1_pad = ceil(S1/B)*B` ensures `dynamic_update_slice` is always legal. + - Padding to `S1_pad = ceil(S1/B)*B` ensures `dynamic_update_slice` is always legal + on the JAX branch; the NumPy branch's Python block loop needs no padding. """ + if not self._is_jax(xp): + return self._curvature_matrix_off_diag_from_np( + rows0=rows0, + cols0=cols0, + vals0=vals0, + rows1=rows1, + cols1=cols1, + vals1=vals1, + S0=S0, + S1=S1, + ) + import jax.numpy as jnp from jax import lax from jax.ops import segment_sum @@ -1207,7 +1419,7 @@ def body(block_i, F01): F = jnp.zeros((M, B), dtype=jnp.float64) F = F.at[rows1, bc].add(v) - G = self.apply_operator(F) # (M, B) + G = self.apply_operator(F, xp=jnp) # (M, B) contrib = vals0[:, None] * G[rows0, :] block = segment_sum(contrib, cols0, num_segments=S0) @@ -1220,7 +1432,40 @@ def body(block_i, F01): F01_pad = lax.fori_loop(0, n_blocks, body, F01_0) return F01_pad[:, :S1] - def operated_matrix_slim_from(self, matrix_slim, extent_index_for_masked_pixel): + def _curvature_matrix_off_diag_from_np( + self, rows0, cols0, vals0, rows1, cols1, vals1, *, S0: int, S1: int + ) -> np.ndarray: + """ + The NumPy/scipy body of `curvature_matrix_off_diag_from`: `F01 = A0ᵀ W~ A1`. + + The same block sweep as `_curvature_matrix_diag_from_np`, but over the columns of + `A1` and projected back with `A0ᵀ`. The result is rectangular in general, so — + unlike the diagonal block — it is not symmetrised. + """ + A0 = self._sparse_matrix_from(rows=rows0, cols=cols0, vals=vals0, S=S0) + A1 = self._sparse_matrix_from(rows=rows1, cols=cols1, vals=vals1, S=S1) + + A0T = A0.T.tocsr() + + S0 = int(S0) + S1 = int(S1) + B = int(self.batch_size) + + F01 = np.zeros((S0, S1), dtype=np.float64) + + for start in range(0, S1, B): + stop = min(start + B, S1) + + F = A1[:, start:stop].toarray() + G = self.apply_operator(F, xp=np) + + F01[:, start:stop] = A0T @ G + + return F01 + + def operated_matrix_slim_from( + self, matrix_slim, extent_index_for_masked_pixel, xp=np + ): """ Apply the interferometer W~ operator to columns defined on the *slim masked* grid. @@ -1235,12 +1480,30 @@ def operated_matrix_slim_from(self, matrix_slim, extent_index_for_masked_pixel): `mapping_matrix` of an `AbstractLinearObjFuncList`). extent_index_for_masked_pixel Array of shape (M_pix,) mapping slim masked pixel indices to extent-grid flat indices. + xp + The array module the operator is applied with: `numpy` (default) or `jax.numpy`. Returns ------- ndarray Array of shape (M_pix, n_cols) equal to W~ applied to each column. """ + if not self._is_jax(xp): + matrix_slim = np.asarray(matrix_slim, dtype=np.float64) + extent_index_for_masked_pixel = np.asarray( + extent_index_for_masked_pixel, dtype=np.int64 + ) + + # A scatter-*set*, not an add: `extent_index_for_masked_pixel` is one extent + # cell per masked pixel and therefore has no repeats, so NumPy fancy-index + # assignment matches the JAX branch's `.at[].set` exactly. + grid_flat = np.zeros((self.M, matrix_slim.shape[1]), dtype=np.float64) + grid_flat[extent_index_for_masked_pixel, :] = matrix_slim + + return self.apply_operator(grid_flat, xp=np)[ + extent_index_for_masked_pixel, : + ] + import jax.numpy as jnp matrix_slim = jnp.asarray(matrix_slim, dtype=jnp.float64) @@ -1251,7 +1514,7 @@ def operated_matrix_slim_from(self, matrix_slim, extent_index_for_masked_pixel): grid_flat = jnp.zeros((self.M, matrix_slim.shape[1]), dtype=jnp.float64) grid_flat = grid_flat.at[extent_index_for_masked_pixel, :].set(matrix_slim) - return self.apply_operator(grid_flat)[extent_index_for_masked_pixel, :] + return self.apply_operator(grid_flat, xp=jnp)[extent_index_for_masked_pixel, :] def curvature_matrix_off_diag_func_list_from( self, @@ -1262,6 +1525,7 @@ def curvature_matrix_off_diag_func_list_from( vals, # triplets where rows are EXTENT indices *, S: int, + xp=np, ): """ Compute the mapper–linear-function off-diagonal block Aᵀ W~ B. @@ -1300,6 +1564,8 @@ def curvature_matrix_off_diag_func_list_from( - `vals` are mapping weights, shape (nnz,) S Number of source pixels / parameters in the mapper. + xp + The array module the operator is applied with: `numpy` (default) or `jax.numpy`. Returns ------- @@ -1311,6 +1577,27 @@ def curvature_matrix_off_diag_func_list_from( - No `batch_size` sweep is required because the operator is applied to `n_funcs` columns (typically a handful) rather than to all S source pixels. """ + if not self._is_jax(xp): + curvature_weights = np.asarray(curvature_weights, dtype=np.float64) + extent_index_for_masked_pixel = np.asarray( + extent_index_for_masked_pixel, dtype=np.int64 + ) + + n_funcs = curvature_weights.shape[1] + + # 1) scatter slim -> extent(flat) (a set, the indices are unique) + grid_flat = np.zeros((self.M, n_funcs), dtype=np.float64) + grid_flat[extent_index_for_masked_pixel, :] = curvature_weights + + # 2) apply W~ on the extent grid + operated = self.apply_operator(grid_flat, xp=np) # (M, n_funcs) + + # 3) project onto the mapper's source pixels with `Aᵀ`, the sparse-matrix + # counterpart of the JAX branch's `segment_sum` over `cols`. + A = self._sparse_matrix_from(rows=rows, cols=cols, vals=vals, S=S) + + return A.T.tocsr() @ operated # (S, n_funcs) + import jax.numpy as jnp from jax.ops import segment_sum @@ -1332,7 +1619,7 @@ def curvature_matrix_off_diag_func_list_from( ) # 2) apply W~ on the extent grid - operated = self.apply_operator(grid_flat) # (M, n_funcs) + operated = self.apply_operator(grid_flat, xp=jnp) # (M, n_funcs) # 3) gather at the mapper's rows (extent coords) and accumulate to source pixels contrib = vals[:, None] * operated[rows, :] @@ -1343,6 +1630,7 @@ def curvature_matrix_func_list_from( curvature_weights_0, # (M_pix, n_funcs_0) curvature_weights_1, # (M_pix, n_funcs_1) extent_index_for_masked_pixel, # (M_pix,) slim -> extent(flat) + xp=np, ): """ Compute a linear-function–linear-function curvature block B0ᵀ W~ B1. @@ -1362,12 +1650,25 @@ def curvature_matrix_func_list_from( on the slim masked grid, of shape (M_pix, n_funcs). extent_index_for_masked_pixel Array of shape (M_pix,) mapping slim masked pixel indices to extent-grid flat indices. + xp + The array module the operator is applied with: `numpy` (default) or `jax.numpy`. Returns ------- ndarray Curvature block of shape (n_funcs_0, n_funcs_1). """ + if not self._is_jax(xp): + curvature_weights_0 = np.asarray(curvature_weights_0, dtype=np.float64) + + operated = self.operated_matrix_slim_from( + matrix_slim=curvature_weights_1, + extent_index_for_masked_pixel=extent_index_for_masked_pixel, + xp=np, + ) + + return curvature_weights_0.T @ operated + import jax.numpy as jnp curvature_weights_0 = jnp.asarray(curvature_weights_0, dtype=jnp.float64) @@ -1375,6 +1676,7 @@ def curvature_matrix_func_list_from( operated = self.operated_matrix_slim_from( matrix_slim=curvature_weights_1, extent_index_for_masked_pixel=extent_index_for_masked_pixel, + xp=jnp, ) return curvature_weights_0.T @ operated diff --git a/autoarray/inversion/inversion/interferometer/sparse.py b/autoarray/inversion/inversion/interferometer/sparse.py index 1e1f11aba..8d7185e14 100644 --- a/autoarray/inversion/inversion/interferometer/sparse.py +++ b/autoarray/inversion/inversion/interferometer/sparse.py @@ -170,11 +170,6 @@ def curvature_matrix(self) -> np.ndarray: ) if len(self.no_regularization_index_list) > 0: - if self._xp is np: - # The sparse operator returns JAX arrays, which the NumPy in-place diagonal - # update below cannot write to. - curvature_matrix = np.array(curvature_matrix) - curvature_matrix = inversion_util.curvature_matrix_with_added_to_diag_from( curvature_matrix=curvature_matrix, value=self.settings.no_regularization_add_to_curvature_diag_value, @@ -205,6 +200,7 @@ def curvature_matrix_diag(self) -> np.ndarray: cols=cols, vals=vals, S=mapper.params, + xp=self._xp, ) @property @@ -233,6 +229,7 @@ def _curvature_matrix_mapper_diag(self) -> Optional[np.ndarray]: cols=cols, vals=vals, S=mapper.params, + xp=self._xp, ) start, end = mapper_param_range_list[mapper_index] @@ -263,6 +260,7 @@ def _curvature_matrix_off_diag_from( vals1=vals_1, S0=mapper_0.params, S1=mapper_1.params, + xp=self._xp, ) @property @@ -359,6 +357,7 @@ def _curvature_matrix_func_list_and_mapper(self) -> np.ndarray: cols=cols, vals=vals, S=mapper.params, + xp=self._xp, ) if self._xp is np: @@ -384,6 +383,7 @@ def _curvature_matrix_func_list_and_mapper(self) -> np.ndarray: curvature_weights_0=mapping_matrix_list[index_0], curvature_weights_1=mapping_matrix_list[index_1], extent_index_for_masked_pixel=extent_index_for_masked_pixel, + xp=self._xp, ) if self._xp is np: diff --git a/test_autoarray/inversion/inversion/interferometer/test_interferometer.py b/test_autoarray/inversion/inversion/interferometer/test_interferometer.py index b55cae30f..7cf1bf1cc 100644 --- a/test_autoarray/inversion/inversion/interferometer/test_interferometer.py +++ b/test_autoarray/inversion/inversion/interferometer/test_interferometer.py @@ -708,3 +708,161 @@ def test__interferometer_sparse_operator__no_regularization_value_added_to_diag( assert curvature_matrix[0, 0] == pytest.approx(curvature_func[0, 0] + value, 1.0e-8) assert curvature_matrix[1, 1] == pytest.approx(curvature_func[1, 1] + value, 1.0e-8) assert curvature_matrix[0, 1] == pytest.approx(curvature_func[0, 1], 1.0e-8) + + +def _sparse_np_vs_jax_setup(mask, n_visibilities, seed, pixels, shape): + """ + Returns the dense dataset, the sparse-operator dataset and a regularized Delaunay mapper + for the end-to-end NumPy/JAX parity test below. + + The operator is built with `batch_size=4` so the block sweep runs more than one block + and ends on a partial one, and with the NumPy brute-force preload builder so the two + inversions are handed a byte-identical operator to start from. + """ + rng = np.random.default_rng(seed=seed) + + 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), + real_space_mask=mask, + transformer_class=aa.TransformerDFT, + ) + + dataset_sparse = dataset.apply_sparse_operator( + nufft_precision_operator=dataset.psf_precision_operator_from(method="numpy"), + batch_size=4, + ) + + mapper = _mapper_from( + mask=mask, + pixels=pixels, + shape=shape, + regularization=aa.reg.Constant(coefficient=1.0), + ) + + return dataset, dataset_sparse, mapper + + +def test__interferometer_sparse_operator__numpy_inversion_matches_jax_inversion(): + """ + End-to-end: `InversionInterferometerSparse(xp=np)` — which now assembles every curvature + block with scipy rather than JAX — must reproduce the `xp=jnp` inversion. + + `curvature_matrix` and `data_vector` come straight off the operator and are pinned + exactly. `reconstruction` and the log-determinant terms come out of a linear solve, whose + NumPy and JAX implementations differ by more than the matrices they are handed do; the + control at the end shows that spread is the solver pair's, by measuring the identical + difference on the *dense* mapping inversion, which this change does not touch. + """ + pytest.importorskip("jax") + + import jax.numpy as jnp + + cases = [ + ( + aa.Mask2D( + mask=[ + [True, True, True, True, True, True, True], + [True, True, True, True, True, True, True], + [True, True, True, False, True, True, True], + [True, True, False, False, False, True, True], + [True, True, True, False, True, True, True], + [True, True, True, True, True, True, True], + [True, True, True, True, True, True, True], + ], + pixel_scales=2.0, + ), + 5, + 0, + 9, + (3, 3), + ), + ( + aa.Mask2D.circular(shape_native=(12, 12), pixel_scales=1.0, radius=4.0), + 64, + 11, + 16, + (4, 4), + ), + ] + + for mask, n_visibilities, seed, pixels, shape in cases: + dataset, dataset_sparse, mapper = _sparse_np_vs_jax_setup( + mask=mask, + n_visibilities=n_visibilities, + seed=seed, + pixels=pixels, + shape=shape, + ) + + inversion_np = aa.Inversion( + dataset=dataset_sparse, linear_obj_list=[mapper], xp=np + ) + inversion_jax = aa.Inversion( + dataset=dataset_sparse, linear_obj_list=[mapper], xp=jnp + ) + + assert isinstance(inversion_np, aa.InversionInterferometerSparse) + assert isinstance(inversion_jax, aa.InversionInterferometerSparse) + + for name in ("curvature_matrix", "data_vector"): + reference = np.asarray(getattr(inversion_jax, name)) + + np.testing.assert_allclose( + np.asarray(getattr(inversion_np, name)), + reference, + rtol=1.0e-10, + atol=1.0e-10 * np.abs(reference).max(), + err_msg=name, + ) + + reconstruction = np.asarray(inversion_jax.reconstruction) + + np.testing.assert_allclose( + np.asarray(inversion_np.reconstruction), + reconstruction, + rtol=1.0e-7, + atol=1.0e-7 * np.abs(reconstruction).max(), + ) + + for name in ( + "log_det_curvature_reg_matrix_term", + "log_det_regularization_matrix_term", + ): + reference = float(getattr(inversion_jax, name)) + + np.testing.assert_allclose( + float(getattr(inversion_np, name)), + reference, + rtol=1.0e-10, + atol=1.0e-10 * abs(reference), + err_msg=name, + ) + + # The control: the dense mapping inversion runs no operator code at all, so whatever + # `xp=np` and `xp=jnp` disagree by there is the linear solver's, not the sparse + # operator's. The sparse path must not be worse. + reconstruction_dense_np = np.asarray( + aa.Inversion( + dataset=dataset, linear_obj_list=[mapper], xp=np + ).reconstruction + ) + reconstruction_dense_jax = np.asarray( + aa.Inversion( + dataset=dataset, linear_obj_list=[mapper], xp=jnp + ).reconstruction + ) + + difference_dense = np.abs( + reconstruction_dense_np - reconstruction_dense_jax + ).max() + difference_sparse = np.abs( + np.asarray(inversion_np.reconstruction) - reconstruction + ).max() + + assert difference_sparse <= max(10.0 * difference_dense, 1.0e-14) 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 26cba69a4..8a59b8a03 100644 --- a/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py +++ b/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py @@ -1,3 +1,5 @@ +import os + import autoarray as aa import numpy as np import pytest @@ -700,3 +702,426 @@ def test__nufft_precision_operator_from__nufftax_absent_falls_back_to_the_numpy_ aa.util.inversion_interferometer.nufft_precision_operator_via_nufft_from( **inputs ) + + +def _numpy_backend_fixtures(): + """ + The two `InterferometerSparseOperator` fixtures the NumPy/JAX parity tests below run on, + each built with `batch_size=4`. + + The 7x7 / K=5 case is the module's shared shape fixture. The 12x12 / K=64 case is the + seeded one the `rfft2` pin uses: its source-pixel count exceeds `batch_size`, so the + block sweep runs more than one block and finishes on a partial one -- the branch the JAX + path needs `dynamic_update_slice` and a column mask for, and the NumPy path a plain + Python loop. + """ + fixtures = [] + + for mask, n_visibilities, seed in ( + (_mask_7x7(), 5, 3), + ( + aa.Mask2D.circular(shape_native=(12, 12), pixel_scales=1.0, radius=4.0), + 64, + 11, + ), + ): + dataset, rng = _dataset_from( + mask=mask, n_visibilities=n_visibilities, seed=seed + ) + + operator = dataset.apply_sparse_operator( + nufft_precision_operator=dataset.psf_precision_operator_from( + method="numpy" + ), + batch_size=4, + ).sparse_operator + + fixtures.append((operator, mask, rng)) + + return fixtures + + +def _assert_numpy_matches_jax(result, result_via_jax): + """ + Asserts a NumPy-branch result matches the JAX branch at the module's exact pin. + + Both branches evaluate the same real-FFT convolution and the same sparse triple product + in float64, so they agree to floating-point round-off and nothing looser is warranted. + `atol` is peak-scaled so that entries orders below the peak -- which carry no relative + accuracy of their own -- do not turn round-off into a failure. + """ + result_via_jax = np.asarray(result_via_jax) + + np.testing.assert_allclose( + np.asarray(result), + result_via_jax, + rtol=1.0e-10, + atol=1.0e-10 * np.abs(result_via_jax).max(), + ) + + +def _triplets_with_duplicates(rng, M, S, nnz): + """ + Returns COO triplets whose row/column index ranges are small enough that + `(row, col)` pairs repeat, so the assembly has to *sum* duplicate entries. The JAX + branch does this with `.at[].add`, the NumPy branch with `scipy.sparse`'s COO + duplicate summation; a NumPy branch that overwrote instead would fail the pin. + """ + rows = rng.integers(0, min(M, 6), size=nnz) + cols = rng.integers(0, S, size=nnz) + vals = rng.normal(size=nnz) + + return rows, cols, vals + + +def test__interferometer_sparse_operator__numpy_branch_matches_jax_branch(): + """ + Every public method of `InterferometerSparseOperator` takes an `xp` and must return the + same matrix on either backend: a CPU fit passing `xp=np` runs the NumPy/scipy bodies + instead of the JAX ones, and that must be a change of backend only. + """ + pytest.importorskip("jax") + + import jax.numpy as jnp + + for operator, mask, rng in _numpy_backend_fixtures(): + M = operator.M + S = 11 + S1 = 9 + + assert S > operator.batch_size + + extent_index_for_masked_pixel = np.array(mask.extent_index_for_masked_pixel) + + # apply_operator + Fbatch = rng.normal(size=(M, 7)) + + _assert_numpy_matches_jax( + operator.apply_operator(Fbatch, xp=np), + operator.apply_operator(jnp.asarray(Fbatch), xp=jnp), + ) + + # curvature_matrix_diag_from + rows, cols, vals = _triplets_with_duplicates(rng, M=M, S=S, nnz=40) + + _assert_numpy_matches_jax( + operator.curvature_matrix_diag_from( + rows=rows, cols=cols, vals=vals, S=S, xp=np + ), + operator.curvature_matrix_diag_from( + rows=rows, cols=cols, vals=vals, S=S, xp=jnp + ), + ) + + # curvature_matrix_off_diag_from + rows_1, cols_1, vals_1 = _triplets_with_duplicates(rng, M=M, S=S1, nnz=30) + + _assert_numpy_matches_jax( + operator.curvature_matrix_off_diag_from( + rows0=rows, + cols0=cols, + vals0=vals, + rows1=rows_1, + cols1=cols_1, + vals1=vals_1, + S0=S, + S1=S1, + xp=np, + ), + operator.curvature_matrix_off_diag_from( + rows0=rows, + cols0=cols, + vals0=vals, + rows1=rows_1, + cols1=cols_1, + vals1=vals_1, + S0=S, + S1=S1, + xp=jnp, + ), + ) + + # operated_matrix_slim_from + matrix_slim = rng.normal(size=(mask.pixels_in_mask, 3)) + + _assert_numpy_matches_jax( + operator.operated_matrix_slim_from( + matrix_slim=matrix_slim, + extent_index_for_masked_pixel=extent_index_for_masked_pixel, + xp=np, + ), + operator.operated_matrix_slim_from( + matrix_slim=matrix_slim, + extent_index_for_masked_pixel=extent_index_for_masked_pixel, + xp=jnp, + ), + ) + + # curvature_matrix_off_diag_func_list_from + curvature_weights = rng.normal(size=(mask.pixels_in_mask, 3)) + + _assert_numpy_matches_jax( + operator.curvature_matrix_off_diag_func_list_from( + curvature_weights=curvature_weights, + extent_index_for_masked_pixel=extent_index_for_masked_pixel, + rows=rows, + cols=cols, + vals=vals, + S=S, + xp=np, + ), + operator.curvature_matrix_off_diag_func_list_from( + curvature_weights=curvature_weights, + extent_index_for_masked_pixel=extent_index_for_masked_pixel, + rows=rows, + cols=cols, + vals=vals, + S=S, + xp=jnp, + ), + ) + + # curvature_matrix_func_list_from + curvature_weights_0 = rng.normal(size=(mask.pixels_in_mask, 2)) + + _assert_numpy_matches_jax( + operator.curvature_matrix_func_list_from( + curvature_weights_0=curvature_weights_0, + curvature_weights_1=curvature_weights, + extent_index_for_masked_pixel=extent_index_for_masked_pixel, + xp=np, + ), + operator.curvature_matrix_func_list_from( + curvature_weights_0=curvature_weights_0, + curvature_weights_1=curvature_weights, + extent_index_for_masked_pixel=extent_index_for_masked_pixel, + xp=jnp, + ), + ) + + +def _delaunay_triplets_from(mask, over_sample_size): + """ + Returns the COO triplets of a real Delaunay mapper on `mask`, alongside its parameter + count, exactly as `InversionInterferometerSparse._sparse_triplets_curvature_from` + builds them. + + These are the triplets the production path actually hands the operator, and they carry + two properties hand-written triplets do not: `mapper_util.sparse_triplets_from` pads + each sub-pixel's interpolation stencil to the longest, emitting `col = -1` with weight + `0.0` for the unused slots, and at `over_sample_size > 1` several sub-pixels of the same + image pixel hit the same source pixel, so `(row, col)` pairs repeat. + """ + grid = aa.Grid2D.from_mask(mask=mask, over_sample_size=over_sample_size) + + image_mesh_grid = aa.image_mesh.Overlay(shape=(4, 4)).image_plane_mesh_grid_from( + mask=mask, adapt_data=None + ) + interpolator = aa.mesh.Delaunay(pixels=16).interpolator_from( + source_plane_data_grid=grid, + source_plane_mesh_grid=image_mesh_grid, + ) + mapper = aa.Mapper(interpolator=interpolator) + + rows, cols, vals = aa.util.mapper.sparse_triplets_from( + pix_indexes_for_sub=mapper.pix_indexes_for_sub_slim_index, + pix_weights_for_sub=mapper.pix_weights_for_sub_slim_index, + slim_index_for_sub=mapper.slim_index_for_sub_slim_index, + fft_index_for_masked_pixel=mask.extent_index_for_masked_pixel, + sub_fraction_slim=mapper.over_sampler.sub_fraction.array, + return_rows_slim=False, + xp=np, + ) + + return np.asarray(rows), np.asarray(cols), np.asarray(vals), mapper.params + + +def test__interferometer_sparse_operator__numpy_branch_matches_jax_branch__delaunay_triplets(): + """ + The same parity, on the triplets a real over-sampled Delaunay mapper produces: padded + `col = -1` entries and repeated `(row, col)` pairs, both of which the JAX branch handles + implicitly (an out-of-range column is dropped, a repeat is accumulated by `.at[].add`) + and the NumPy branch must handle explicitly. + """ + pytest.importorskip("jax") + + import jax.numpy as jnp + + mask = aa.Mask2D.circular(shape_native=(12, 12), pixel_scales=1.0, radius=4.0) + + dataset, _ = _dataset_from(mask=mask, n_visibilities=64, seed=11) + + operator = dataset.apply_sparse_operator( + nufft_precision_operator=dataset.psf_precision_operator_from(method="numpy"), + batch_size=4, + ).sparse_operator + + rows, cols, vals, S = _delaunay_triplets_from(mask=mask, over_sample_size=2) + + # The fixture only tests what it contains: a mapper whose stencils happened not to be + # padded, or whose sub-pixels happened not to collide, would leave both behaviours + # untested and the test would still pass. + assert (cols < 0).any() + + pairs, counts = np.unique( + np.stack([rows[cols >= 0], cols[cols >= 0]], axis=1), axis=0, return_counts=True + ) + assert (counts > 1).any() + + assert S > operator.batch_size + + _assert_numpy_matches_jax( + operator.curvature_matrix_diag_from( + rows=rows, cols=cols, vals=vals, S=S, xp=np + ), + operator.curvature_matrix_diag_from( + rows=rows, cols=cols, vals=vals, S=S, xp=jnp + ), + ) + + +def test__interferometer_sparse_operator__numpy_branch_runs_with_jax_unimportable(): + """ + The point of the `xp` branch: a CPU user must be able to build the operator and run every + NumPy body in an environment where JAX is not installed at all. + + Deliberately carries no `importorskip("jax")`, so the `unittest-nojax` CI leg runs it. + """ + import sys + + monkeypatch = pytest.MonkeyPatch() + + mask = _mask_7x7() + + dataset, rng = _dataset_from(mask=mask, n_visibilities=5, seed=3) + + try: + # `None` in `sys.modules` is the documented way to make an import fail: `import jax` + # then raises `ImportError` rather than finding the installed package. + monkeypatch.setitem(sys.modules, "jax", None) + monkeypatch.setitem(sys.modules, "jax.numpy", None) + + # The brute-force NumPy builder, because the default `"nufft"` builder runs on JAX. + operator = dataset.apply_sparse_operator( + nufft_precision_operator=dataset.psf_precision_operator_from( + method="numpy" + ), + batch_size=4, + ).sparse_operator + + M = operator.M + S = 11 + + extent_index_for_masked_pixel = np.array(mask.extent_index_for_masked_pixel) + + rows, cols, vals = _triplets_with_duplicates(rng, M=M, S=S, nnz=40) + rows_1, cols_1, vals_1 = _triplets_with_duplicates(rng, M=M, S=5, nnz=20) + + curvature_weights = rng.normal(size=(mask.pixels_in_mask, 3)) + + assert operator.apply_operator(rng.normal(size=(M, 3)), xp=np).shape == (M, 3) + assert operator.curvature_matrix_diag_from( + rows=rows, cols=cols, vals=vals, S=S, xp=np + ).shape == (S, S) + assert operator.curvature_matrix_off_diag_from( + rows0=rows, + cols0=cols, + vals0=vals, + rows1=rows_1, + cols1=cols_1, + vals1=vals_1, + S0=S, + S1=5, + xp=np, + ).shape == (S, 5) + assert operator.operated_matrix_slim_from( + matrix_slim=curvature_weights, + extent_index_for_masked_pixel=extent_index_for_masked_pixel, + xp=np, + ).shape == (mask.pixels_in_mask, 3) + assert operator.curvature_matrix_off_diag_func_list_from( + curvature_weights=curvature_weights, + extent_index_for_masked_pixel=extent_index_for_masked_pixel, + rows=rows, + cols=cols, + vals=vals, + S=S, + xp=np, + ).shape == (S, 3) + assert operator.curvature_matrix_func_list_from( + curvature_weights_0=curvature_weights, + curvature_weights_1=curvature_weights, + extent_index_for_masked_pixel=extent_index_for_masked_pixel, + xp=np, + ).shape == (3, 3) + + # Without this the test would also pass in an environment where JAX imports fine, + # and would therefore stop testing anything the day the block above broke. + with pytest.raises(ImportError): + operator.Khat + finally: + monkeypatch.undo() + + +def test__interferometer_sparse_operator__numpy_branch_imports_no_jax(): + """ + Stronger than the guard above, and the actual complaint the `xp` branch answers: JAX must + not merely be unnecessary, it must never be *imported*. Importing it costs seconds and + allocates a backend, and the operator's JAX state (`Khat`, `col_offsets`) is lazy purely + so that a NumPy fit never pays for it. + + Run in a subprocess because `jax` is in `sys.modules` for the rest of this session as + soon as any other test imports it. + """ + import subprocess + import sys + import textwrap + + script = textwrap.dedent( + """ + import os + import sys + + import numpy as np + + import autoarray as aa + + mask = aa.Mask2D.circular(shape_native=(8, 8), pixel_scales=1.0, radius=3.0) + + rng = np.random.default_rng(seed=1) + + dataset = aa.Interferometer( + data=aa.Visibilities(visibilities=rng.normal(size=(5, 2))), + noise_map=aa.VisibilitiesNoiseMap(visibilities=np.ones((5, 2))), + uv_wavelengths=rng.normal(size=(5, 2)), + real_space_mask=mask, + transformer_class=aa.TransformerDFT, + ) + + operator = dataset.apply_sparse_operator( + method="numpy", batch_size=4 + ).sparse_operator + + M = operator.M + + operator.apply_operator(rng.normal(size=(M, 3)), xp=np) + operator.curvature_matrix_diag_from( + rows=rng.integers(0, M, 20), + cols=rng.integers(0, 6, 20), + vals=rng.normal(size=20), + S=6, + xp=np, + ) + + assert "jax" not in sys.modules, sorted( + name for name in sys.modules if name.startswith("jax") + ) + """ + ) + + environment = dict(os.environ) + environment["PYAUTO_DISABLE_JAX"] = "1" + + subprocess.run( + [sys.executable, "-c", script], check=True, env=environment, timeout=600 + )