autoarray 2026.5.8.2__tar.gz → 2026.5.14.2__tar.gz

{autoarray-2026.5.8.2 → autoarray-2026.5.14.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: autoarray
-Version: 2026.5.8.2
+Version: 2026.5.14.2
 Summary: PyAuto Data Structures
 Author-email: James Nightingale <James.Nightingale@newcastle.ac.uk>, Richard Hayes <richard@rghsoftware.co.uk>
 License: MIT
@@ -18,7 +18,7 @@ Classifier: Programming Language :: Python :: 3.13
 Requires-Python: >=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: autoconf==2026.5.8.2
+Requires-Dist: autoconf==2026.5.14.2
 Requires-Dist: astropy<=7.2.0,>=5.0
 Requires-Dist: decorator>=4.0.0
 Requires-Dist: dill>=0.3.1.1
@@ -32,6 +32,7 @@ Requires-Dist: autoconf[jax]; extra == "jax"
 Provides-Extra: optional
 Requires-Dist: autoarray[jax]; extra == "optional"
 Requires-Dist: numba; extra == "optional"
+Requires-Dist: nufftax; extra == "optional"
 Requires-Dist: pynufft; extra == "optional"
 Requires-Dist: tensorflow-probability==0.25.0; extra == "optional"
 Provides-Extra: test
@@ -40,6 +41,7 @@ Provides-Extra: dev
 Requires-Dist: pytest; extra == "dev"
 Requires-Dist: black; extra == "dev"
 Requires-Dist: numba; extra == "dev"
+Requires-Dist: nufftax; extra == "dev"
 Requires-Dist: pynufft==2022.2.2; extra == "dev"
 Dynamic: license-file
 

{autoarray-2026.5.8.2 → autoarray-2026.5.14.2}/autoarray/__init__.py

@@ -60,6 +60,7 @@ from .mask.mask_1d import Mask1D
 from .mask.mask_2d import Mask2D
 from .operators.transformer import TransformerDFT
 from .operators.transformer import TransformerNUFFT
+from .operators.transformer import TransformerNUFFTPyNUFFT
 from .operators.over_sampling.decorator import over_sample
 from .operators.contour import Grid2DContour
 from .layout.layout import Layout1D
@@ -77,6 +78,7 @@ from .inversion.mesh.mesh_geometry.delaunay import MeshGeometryDelaunay
 from .inversion.mesh.interpolator.rectangular import InterpolatorRectangular
 from .inversion.mesh.interpolator.delaunay import InterpolatorDelaunay
 from .operators.convolver import Convolver
+from .operators.interp_2d import interp_2d
 from .structures.vectors.uniform import VectorYX2D
 from .structures.vectors.irregular import VectorYX2DIrregular
 from .structures.triangles.abstract import AbstractTriangles
@@ -101,4 +103,4 @@ from autoconf.fitsable import hdu_list_for_output_from
 
 conf.instance.register(__file__)
 
-__version__ = "2026.5.8.2"
+__version__ = "2026.5.14.2"

{autoarray-2026.5.8.2 → autoarray-2026.5.14.2}/autoarray/dataset/interferometer/dataset.py

@@ -251,6 +251,29 @@ class Interferometer(AbstractDataset):
             enabling efficient pixelized source reconstruction via the sparse linear algebra formalism.
         """
 
+        if isinstance(self.transformer, TransformerNUFFT):
+            raise NotImplementedError(
+                "\n--------------------\n"
+                "`apply_sparse_operator` is not yet supported with the default "
+                "`TransformerNUFFT` (nufftax-backed) transformer.\n\n"
+                "The sparse-operator path consumes the dirty image returned by "
+                "`transformer.image_from(use_adjoint_scaling=True)` together with "
+                "the NUFFT precision operator; their relative scale matters. The "
+                "new `TransformerNUFFT` returns the strict mathematical adjoint "
+                "(matching `TransformerDFT`), whereas the legacy pynufft adjoint "
+                "applies an internal Kaiser-Bessel kernel deconvolution. The two "
+                "scales differ by a non-constant factor, so feeding the new "
+                "dirty image into the existing sparse-operator solver would "
+                "silently give wrong answers.\n\n"
+                "Workarounds:\n"
+                "  - Build the dataset with `transformer_class=TransformerDFT` "
+                "(the JAX-likelihood scripts do this today), or\n"
+                "  - Build the dataset with "
+                "`transformer_class=TransformerNUFFTPyNUFFT` to keep the legacy "
+                "pynufft adjoint scale (requires `pip install pynufft`).\n"
+                "----------------------"
+            )
+
         if nufft_precision_operator is None:
 
             logger.info(

{autoarray-2026.5.8.2 → autoarray-2026.5.14.2}/autoarray/inversion/inversion/abstract.py

@@ -802,7 +802,8 @@ class AbstractInversion:
     def regularization_weights_mapper_dict(self) -> Dict[LinearObj, np.ndarray]:
         regularization_weights_dict = {}
 
-        for index, mapper in enumerate(self.cls_list_from(cls=Mapper)):
+        for mapper in self.cls_list_from(cls=Mapper):
+            index = self.linear_obj_list.index(mapper)
             regularization_weights_dict[mapper] = self.regularization_weights_from(
                 index=index,
             )

{autoarray-2026.5.8.2 → autoarray-2026.5.14.2}/autoarray/operators/convolver.py

@@ -133,6 +133,13 @@ class ConvolverState:
 
         self.fft_kernel = np.fft.rfft2(self.kernel.native.array, s=self.fft_shape)
         self.fft_kernel_mapping = np.expand_dims(self.fft_kernel, 2)
+        # Pre-cached complex64 view for the use_mixed_precision=True path of
+        # convolved_image_from. Cast once here so the FFT branch does not
+        # repeat the astype per JIT trace — it would otherwise produce a fresh
+        # numpy buffer each call, which on CPU costs more than the fp32 FFT
+        # saves. convolved_mapping_matrix_from intentionally does NOT use a
+        # complex64 kernel — see that method's body for why.
+        self.fft_kernel_c64 = self.fft_kernel.astype(np.complex64)
 
 
 class Convolver:
@@ -532,17 +539,23 @@ class Convolver:
 
         state = self.state_from(mask=image.mask)
 
+        # When use_mixed_precision is on, the FFT runs in complex64 end-to-end:
+        # the input cube is allocated as float32, rfft2 emits complex64, the
+        # precomputed (complex128) kernel is cast on the fly, and irfft2
+        # returns float32 natively. No trailing astype is needed.
+        real_dtype = jnp.float32 if use_mixed_precision else jnp.float64
+
         # Build combined native image in the FFT dtype
-        image_both_native = xp.zeros(state.fft_shape, dtype=jnp.float64)
+        image_both_native = xp.zeros(state.fft_shape, dtype=real_dtype)
 
         image_both_native = image_both_native.at[state.mask.slim_to_native_tuple].set(
-            jnp.asarray(image.array, dtype=jnp.float64)
+            jnp.asarray(image.array, dtype=real_dtype)
         )
 
         if blurring_image is not None:
             image_both_native = image_both_native.at[
                 state.blurring_mask.slim_to_native_tuple
-            ].set(jnp.asarray(blurring_image.array, dtype=jnp.float64))
+            ].set(jnp.asarray(blurring_image.array, dtype=real_dtype))
         else:
             warnings.warn(
                 "No blurring_image provided. Only the direct image will be convolved. "
@@ -554,9 +567,14 @@ class Convolver:
             image_both_native, s=state.fft_shape, axes=(0, 1)
         )
 
+        # Pick the precomputed kernel matching the FFT dtype. ConvolverState
+        # caches both complex128 (default) and complex64 (mixed precision) at
+        # init time, so this is a constant lookup rather than a per-call cast.
+        fft_kernel = state.fft_kernel_c64 if use_mixed_precision else state.fft_kernel
+
         # Multiply by PSF in Fourier space and invert
         blurred_image_full = xp.fft.irfft2(
-            state.fft_kernel * fft_image_native, s=state.fft_shape, axes=(0, 1)
+            fft_kernel * fft_image_native, s=state.fft_shape, axes=(0, 1)
         )
         ky, kx = self.kernel.shape_native  # (21, 21)
         off_y = (ky - 1) // 2
@@ -572,15 +590,11 @@ class Convolver:
             blurred_image_full, start_indices, state.fft_shape
         )
 
-        # Return slim form; optionally cast for downstream stability
+        # Return slim form; dtype already matches use_mixed_precision via the
+        # FFT path, so no explicit downcast.
         blurred_slim = blurred_image_native[state.mask.slim_to_native_tuple]
 
-        blurred_image = Array2D(values=blurred_slim, mask=image.mask)
-
-        if use_mixed_precision:
-            blurred_image = blurred_image.astype(jnp.float32)
-
-        return blurred_image
+        return Array2D(values=blurred_slim, mask=image.mask)
 
     def convolved_mapping_matrix_from(
         self,
@@ -677,7 +691,19 @@ class Convolver:
         # -------------------------------------------------------------------------
         # Mixed precision handling
         # -------------------------------------------------------------------------
-        fft_complex_dtype = jnp.complex64 if use_mixed_precision else jnp.complex128
+        # mapping_matrix_native_from honors use_mixed_precision and produces a
+        # fp32 native cube. rfft2 of that cube emits complex64. We deliberately
+        # multiply by the complex128 precomputed kernel below, which upcasts
+        # the product back to complex128 so the irfft2 returns float64. This
+        # asymmetry is intentional: pixelization meshes with K >> 40 source
+        # pixels accumulate enough fp32 round-off through the NNLS active-set
+        # / log-determinant that the figure_of_merit drifts by O(1) units
+        # (verified on the delaunay_mge regression). The fp32 input cube and
+        # complex64 forward FFT still buy us a faster scatter and slightly
+        # cheaper rfft2; keeping the kernel multiply in complex128 preserves
+        # the precision the downstream linear algebra needs.
+        # convolved_image_from (used by light profiles) takes the full fp32
+        # path because its 40-column linear systems are well-conditioned.
 
         # -------------------------------------------------------------------------
         # Build native cube on the *native mask grid*

autoarray-2026.5.14.2/autoarray/operators/interp_2d.py

@@ -0,0 +1,74 @@
+"""
+2D regular-grid bilinear interpolation with NumPy and JAX backends.
+
+Mirrors the 1D dispatch pattern in
+``autoarray.inversion.mesh.interpolator.rectangular_spline``.
+
+NumPy path delegates to ``scipy.interpolate.RegularGridInterpolator``
+(``bounds_error=False``, configurable ``fill_value``). JAX path is built on
+``jax.scipy.ndimage.map_coordinates(order=1, mode='constant')``.
+"""
+
+import numpy as np
+
+
+def _interp_2d_numpy(points, x_axis, y_axis, values, fill_value=0.0):
+    """
+    Evaluate ``values`` (regular grid on (x_axis, y_axis)) at ``(N, 2)`` query
+    points using bilinear interpolation. Out-of-bounds queries return
+    ``fill_value``.
+
+    Matches the wiring used by ``DatasetInterp`` in PyAutoGalaxy: pass the
+    axes and the values array unchanged, do not transpose.
+    """
+    from scipy import interpolate
+
+    interpolator = interpolate.RegularGridInterpolator(
+        points=(x_axis, y_axis),
+        values=values,
+        bounds_error=False,
+        fill_value=fill_value,
+    )
+    return interpolator(points)
+
+
+def _interp_2d_jax(points, x_axis, y_axis, values, fill_value=0.0):
+    """
+    JAX-compatible counterpart to ``_interp_2d_numpy``. Uses
+    ``jax.scipy.ndimage.map_coordinates(order=1, mode='constant')``.
+
+    The query ``points`` are in world coordinates with column 0 matching
+    ``x_axis`` and column 1 matching ``y_axis`` (mirroring the scipy
+    interpolator's ``(x, y)`` argument order). World coords are converted to
+    pixel-fractional indices into ``values`` using axis spacings derived from
+    the first two axis points.
+    """
+    import jax.numpy as jnp
+    from jax.scipy.ndimage import map_coordinates
+
+    x_axis = jnp.asarray(x_axis)
+    y_axis = jnp.asarray(y_axis)
+    values = jnp.asarray(values)
+    pts = jnp.asarray(points)
+
+    # World -> pixel-fractional indices.
+    dx = x_axis[1] - x_axis[0]
+    dy = y_axis[1] - y_axis[0]
+    col_coords = (pts[:, 0] - x_axis[0]) / dx  # column 0 of points -> x-axis
+    row_coords = (pts[:, 1] - y_axis[0]) / dy  # column 1 of points -> y-axis
+
+    # map_coordinates indexes values[i, j] where i is the first-axis index
+    # and j the second. The scipy call uses points=(x_axis, y_axis) so the
+    # first axis of `values` is interpreted as the x-axis.
+    coords = jnp.stack([col_coords, row_coords], axis=0)
+
+    return map_coordinates(
+        values, coords, order=1, mode="constant", cval=fill_value
+    )
+
+
+def interp_2d(points, x_axis, y_axis, values, fill_value=0.0, xp=np):
+    """Dispatcher -- picks ``_interp_2d_numpy`` or ``_interp_2d_jax`` based on ``xp``."""
+    if xp is np:
+        return _interp_2d_numpy(points, x_axis, y_axis, values, fill_value=fill_value)
+    return _interp_2d_jax(points, x_axis, y_axis, values, fill_value=fill_value)

{autoarray-2026.5.8.2 → autoarray-2026.5.14.2}/autoarray/operators/transformer.py

@@ -23,16 +23,36 @@ from autoarray.structures.arrays import array_2d_util
 from autoarray.operators import transformer_util
 
 
+try:
+    import nufftax as _nufftax
+except ModuleNotFoundError:
+    _nufftax = None
+
+
 def pynufft_exception():
     raise ModuleNotFoundError(
         "\n--------------------\n"
-        "You are attempting to perform interferometer analysis.\n\n"
+        "You are attempting to perform interferometer analysis with the legacy "
+        "pynufft-backed `TransformerNUFFTPyNUFFT`.\n\n"
         "However, the optional library PyNUFFT (https://github.com/jyhmiinlin/pynufft) is not installed.\n\n"
         "Install it via the command `pip install pynufft==2022.2.2`.\n\n"
         "----------------------"
     )
 
 
+def nufftax_exception():
+    raise ModuleNotFoundError(
+        "\n--------------------\n"
+        "You are attempting to perform interferometer analysis with the default "
+        "JAX-native `TransformerNUFFT`.\n\n"
+        "However, the optional library nufftax (https://github.com/GragasLab/nufftax) is not installed.\n\n"
+        "Install it via the command `pip install nufftax`.\n\n"
+        "If you want to use the legacy pynufft backend instead, pass "
+        "`transformer_class=TransformerNUFFTPyNUFFT` and install pynufft.\n\n"
+        "----------------------"
+    )
+
+
 class TransformerDFT:
     def __init__(
         self,
@@ -175,13 +195,18 @@ class TransformerDFT:
         )
 
 
-class TransformerNUFFT(NUFFT_cpu):
+class TransformerNUFFTPyNUFFT(NUFFT_cpu):
     def __init__(
         self, uv_wavelengths: np.ndarray, real_space_mask: Mask2D, xp=np, **kwargs
     ):
         """
         Performs the Non-Uniform Fast Fourier Transform (NUFFT) for interferometric image reconstruction.
 
+        Legacy pynufft-backed transformer. The default `TransformerNUFFT` is now backed by `nufftax`
+        (JAX-native, differentiable, ~zero gridding error) — this class is retained so users who depend
+        on pynufft's specific gridding behaviour can opt in by passing
+        `transformer_class=TransformerNUFFTPyNUFFT`.
+
         This transformer uses the PyNUFFT library to efficiently compute the Fourier transform
         of an image defined on a regular real-space grid to a set of non-uniform uv-plane (Fourier space)
         coordinates, as is typical in radio interferometry.
@@ -226,7 +251,7 @@ class TransformerNUFFT(NUFFT_cpu):
         if isinstance(self, NUFFTPlaceholder):
             pynufft_exception()
 
-        super(TransformerNUFFT, self).__init__()
+        super(TransformerNUFFTPyNUFFT, self).__init__()
 
         self.uv_wavelengths = uv_wavelengths
         self.real_space_mask = real_space_mask
@@ -469,3 +494,227 @@ class TransformerNUFFT(NUFFT_cpu):
                 )
 
         return transformed_mapping_matrix
+
+
+class TransformerNUFFT:
+    def __init__(
+        self,
+        uv_wavelengths: np.ndarray,
+        real_space_mask: Mask2D,
+        eps: float = 1e-12,
+        xp=np,
+        **kwargs,
+    ):
+        """
+        JAX-native Non-Uniform FFT for image -> visibilities, backed by `nufftax`.
+
+        This is the default `TransformerNUFFT` in PyAutoArray. It uses the
+        `nufftax` library (https://github.com/GragasLab/nufftax), a pure-JAX
+        NUFFT implementation that supports `jax.jit`, `jax.grad`, and
+        `jax.vmap`. It replaces the legacy `TransformerNUFFTPyNUFFT` (which
+        wraps the non-differentiable `pynufft` library) as the default backend.
+
+        Convention recipe (matches `TransformerDFT` to ~1e-13 relative across
+        odd/even/non-square image sizes):
+
+            image_flipped = image[::-1, :]
+            x = 2 * pi * u_lambda * pixel_scale_rad
+            y = 2 * pi * v_lambda * pixel_scale_rad
+            offset_x = 0.5 if N_x is even else 0.0
+            offset_y = 0.5 if N_y is even else 0.0
+            shift = exp(-i * (offset_x * x + offset_y * y))
+            visibilities = nufftax.nufft2d2(x, y, image_flipped, eps, -1) * shift
+
+        The `shift` factor is the half-pixel correction between autoarray's
+        grid centre at index `(N - 1) / 2` and nufftax's mode-0 at index
+        `N // 2`; pynufft applies this internally, nufftax does not.
+
+        Parameters
+        ----------
+        uv_wavelengths
+            The (u, v) coordinates of the measured visibilities in wavelengths,
+            shape `(n_vis, 2)`.
+        real_space_mask
+            The 2D mask defining the real-space image grid.
+        eps
+            Requested NUFFT precision passed to nufftax. Defaults to `1e-12`
+            (effectively machine precision); relax to `1e-9` or `1e-6` for
+            faster execution if marginal accuracy is acceptable.
+        xp
+            Accepted for signature compatibility with the legacy class; not
+            stored. The active backend is selected per-call via the `xp`
+            argument to `visibilities_from` / `image_from`.
+
+        Attributes
+        ----------
+        grid
+            The real-space pixel grid in radians (computed from the mask).
+        total_visibilities
+            Number of measured visibilities.
+        total_image_pixels
+            Number of unmasked pixels in the image grid.
+        adjoint_scaling
+            Scaling factor available for callers who want to apply an
+            optional normalisation to the adjoint output. Provided for
+            parity with the legacy class.
+        """
+        from astropy import units
+
+        if _nufftax is None:
+            nufftax_exception()
+
+        self.uv_wavelengths = uv_wavelengths.astype("float")
+        self.real_space_mask = real_space_mask
+        self.grid = Grid2D.from_mask(mask=self.real_space_mask).in_radians
+        self.eps = eps
+        self.native_index_for_slim_index = copy.copy(
+            real_space_mask.derive_indexes.native_for_slim.astype("int")
+        )
+
+        pixel_scale_rad = self.grid.pixel_scales[0] * units.arcsec.to(units.rad)
+        # nufft2d2 frequency arguments:
+        #   x is paired with the column-axis mode (image x)
+        #   y is paired with the row-axis    mode (image y)
+        # Both must lie in [-pi, pi); the 2*pi*Δ_rad scaling makes uv_lambda
+        # land in that range for any sane uv-coverage.
+        self._x = 2.0 * np.pi * self.uv_wavelengths[:, 0] * pixel_scale_rad
+        self._y = 2.0 * np.pi * self.uv_wavelengths[:, 1] * pixel_scale_rad
+
+        n_y, n_x = self.real_space_mask.shape_native
+        offset_x = 0.5 if n_x % 2 == 0 else 0.0
+        offset_y = 0.5 if n_y % 2 == 0 else 0.0
+        self._shift = np.exp(-1j * (offset_x * self._x + offset_y * self._y))
+
+        self.total_visibilities = uv_wavelengths.shape[0]
+        self.total_image_pixels = real_space_mask.pixels_in_mask
+        self.adjoint_scaling = (2.0 * n_y) * (2.0 * n_x)
+
+    def _forward_native(self, image_native_2d, xp=np):
+        """Run nufft2d2 on a 2D native-shape image array, returning visibilities."""
+        if xp.__name__.startswith("jax"):
+            import jax.numpy as jnp
+
+            img = jnp.asarray(image_native_2d)[::-1, :].astype(jnp.complex128)
+            x = jnp.asarray(self._x)
+            y = jnp.asarray(self._y)
+            shift = jnp.asarray(self._shift)
+            return _nufftax.nufft2d2(x, y, img, self.eps, -1) * shift
+
+        img = image_native_2d[::-1, :].astype(np.complex128)
+        out = _nufftax.nufft2d2(self._x, self._y, img, self.eps, -1) * self._shift
+        return np.array(np.asarray(out))
+
+    def visibilities_from(self, image, xp=np) -> Visibilities:
+        """
+        Forward NUFFT: real-space image -> visibilities at the configured uv points.
+
+        For numpy callers (`xp=np`) the result is materialised back to numpy
+        before being wrapped in `Visibilities`. For JAX callers (`xp=jnp`)
+        the result stays as a `jax.Array` so it can flow through `jax.jit`
+        / `jax.grad` / `jax.vmap` without device round-trips.
+        """
+        if xp.__name__.startswith("jax"):
+            import jax.numpy as jnp
+
+            image_native = jnp.zeros(image.mask.shape, dtype=image.dtype)
+            image_native = image_native.at[image.mask.slim_to_native_tuple].set(
+                image.array
+            )
+        else:
+            image_native = image.native.array
+
+        return Visibilities(visibilities=self._forward_native(image_native, xp=xp))
+
+    def image_from(
+        self,
+        visibilities: Visibilities,
+        use_adjoint_scaling: bool = False,
+        xp=np,
+    ) -> Array2D:
+        """
+        Adjoint NUFFT: visibilities -> real-space (dirty) image.
+
+        Implemented as `nufftax.nufft2d1` with `conj(shift)` applied to the
+        visibilities and a final row-flip to return to autoarray's native
+        orientation. The real part is taken to discard imaginary residue,
+        matching the legacy class' behaviour.
+
+        Note that this is the **mathematical adjoint** of `visibilities_from`,
+        with no kernel deconvolution applied. The dirty image therefore
+        differs in absolute scale from the legacy `TransformerNUFFTPyNUFFT`
+        adjoint (which applies pynufft's internal IFFT and kernel
+        deconvolution). The structure of the dirty image is the same, and
+        the values match `TransformerDFT.image_from` exactly.
+
+        **Scale-sensitive callers**: `Interferometer.apply_sparse_operator`
+        consumes the dirty-image scale together with a precision operator; it
+        is currently incompatible with this class and raises
+        `NotImplementedError`. Use `TransformerDFT` or
+        `TransformerNUFFTPyNUFFT` if you need the sparse-operator path.
+        """
+        n_y, n_x = self.real_space_mask.shape_native
+        n_modes = (n_x, n_y)  # nufftax wants (n1, n2) = (N_x, N_y)
+
+        if xp.__name__.startswith("jax"):
+            import jax.numpy as jnp
+
+            x = jnp.asarray(self._x)
+            y = jnp.asarray(self._y)
+            shift_conj = jnp.asarray(np.conj(self._shift))
+            c = jnp.asarray(visibilities.array) * shift_conj
+            f = _nufftax.nufft2d1(x, y, c, n_modes, self.eps, +1)
+            image = jnp.real(f)[::-1, :]
+        else:
+            c = visibilities.array * np.conj(self._shift)
+            f = _nufftax.nufft2d1(self._x, self._y, c, n_modes, self.eps, +1)
+            image = np.array(np.asarray(f)[::-1, :].real)
+
+        if use_adjoint_scaling:
+            image = image * self.adjoint_scaling
+
+        return Array2D(values=image, mask=self.real_space_mask)
+
+    def transform_mapping_matrix(self, mapping_matrix, xp=np):
+        """
+        Apply the forward NUFFT to each column of a mapping matrix.
+
+        All columns are scattered into a single batched native-shape image
+        of shape ``(n_src, N_y, N_x)`` and passed through nufft2d2 in one
+        call (nufft2d2 supports batched ``f``). This avoids the
+        per-column Python loop that, under ``jax.jit``, would unroll into
+        ``n_src`` separate NUFFT invocations and blow up the JIT graph
+        for pixelization-heavy fits (notably double-source-plane).
+        """
+        n_src = mapping_matrix.shape[1]
+        rows, cols = self.real_space_mask.slim_to_native_tuple
+        n_y, n_x = self.real_space_mask.shape_native
+
+        if xp.__name__.startswith("jax"):
+            import jax.numpy as jnp
+
+            mm_T = jnp.asarray(mapping_matrix).T.astype(jnp.complex128)
+            source_images = jnp.zeros((n_src, n_y, n_x), dtype=jnp.complex128)
+            source_images = source_images.at[
+                jnp.arange(n_src)[:, None],
+                jnp.asarray(rows)[None, :],
+                jnp.asarray(cols)[None, :],
+            ].set(mm_T)
+            flipped = source_images[:, ::-1, :]
+            x = jnp.asarray(self._x)
+            y = jnp.asarray(self._y)
+            shift = jnp.asarray(self._shift)
+            # nufft2d2 returns shape (n_trans, M); transpose to (M, n_src).
+            vis_batched = (
+                _nufftax.nufft2d2(x, y, flipped, self.eps, -1) * shift[None, :]
+            )
+            return vis_batched.T
+
+        mm_T = np.asarray(mapping_matrix).T.astype(np.complex128)
+        source_images = np.zeros((n_src, n_y, n_x), dtype=np.complex128)
+        source_images[np.arange(n_src)[:, None], rows[None, :], cols[None, :]] = mm_T
+        flipped = source_images[:, ::-1, :]
+        vis_batched = (
+            _nufftax.nufft2d2(self._x, self._y, flipped, self.eps, -1)
+            * self._shift[None, :]
+        )
+        return np.array(np.asarray(vis_batched).T)