From 7901c6df4c1375b9569b82437f343a6115ee4e86 Mon Sep 17 00:00:00 2001
From: Jammy2211 <JNightingale2211@gmail.com>
Date: Wed, 21 Oct 2026 11:59:05 +0100
Subject: [PATCH 1/5] feat: primary beam helper and beamed forward transforms

Add autoarray/dataset/interferometer/primary_beam.py with
primary_beam_gaussian_from (FWHM = pb_factor c / nu / D) and the slim
conversion/validation helper. TransformerDFT and TransformerNUFFT take
primary_beam=None, expose with_primary_beam (shallow copy) and multiply
by P on the slim grid in visibilities_from / transform_mapping_matrix
(numpy and JAX); image_from stays the un-beamed adjoint.

Refs #640

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
---
 autoarray/__init__.py                         |   1 +
 .../dataset/interferometer/primary_beam.py    | 136 ++++++++++++++++++
 autoarray/operators/transformer.py            | 112 ++++++++++++++-
 3 files changed, 242 insertions(+), 7 deletions(-)
 create mode 100644 autoarray/dataset/interferometer/primary_beam.py

diff --git a/autoarray/__init__.py b/autoarray/__init__.py
index 0445a636f..d1b7f08c6 100644
--- a/autoarray/__init__.py
+++ b/autoarray/__init__.py
@@ -17,6 +17,7 @@
 from .dataset.imaging.simulator import SimulatorImaging
 from .dataset.interferometer.dataset import Interferometer
 from .dataset.interferometer.simulator import SimulatorInterferometer
+from .dataset.interferometer.primary_beam import primary_beam_gaussian_from
 from .dataset.dataset_model import DatasetModel
 from .fit.fit_dataset import AbstractFit
 from .fit.fit_dataset import FitDataset
diff --git a/autoarray/dataset/interferometer/primary_beam.py b/autoarray/dataset/interferometer/primary_beam.py
new file mode 100644
index 000000000..7a7ebbb35
--- /dev/null
+++ b/autoarray/dataset/interferometer/primary_beam.py
@@ -0,0 +1,136 @@
+"""
+Helpers for building the primary beam `P` of an interferometer dataset.
+
+The primary beam is the direction-dependent sensitivity of each antenna's dish: the
+interferometer measures the visibilities of the apparent sky `P x I`, not of the true sky
+`I`. It is passed to `Interferometer(primary_beam=...)` as an image on the real-space
+mask, after which the transformers and the sparse (w-tilde) inversion apply it so every
+light profile, linear light profile and pixelization models the true sky.
+
+A user-supplied beam image (e.g. one exported by CASA) is the general API; the Gaussian
+here is the standard approximation to the main lobe of a uniformly illuminated dish.
+"""
+
+import numpy as np
+from typing import Tuple
+
+from autoarray.mask.mask_2d import Mask2D
+from autoarray.structures.arrays.uniform_2d import Array2D
+from autoarray.structures.grids.uniform_2d import Grid2D
+
+SPEED_OF_LIGHT_M_PER_S = 299792458.0
+
+ARCSEC_PER_RADIAN = 180.0 / np.pi * 3600.0
+
+
+def primary_beam_fwhm_arcsec_from(
+    frequency_hz: float, dish_diameter_m: float, pb_factor: float = 1.13
+) -> float:
+    """
+    Returns the full width at half maximum, in arcseconds, of the Gaussian approximation to
+    the primary beam of a dish of diameter `D` observing at frequency `nu`:
+
+        FWHM = pb_factor * lambda / D = pb_factor * c / (nu * D)  [radians]
+
+    Parameters
+    ----------
+    frequency_hz
+        The observing frequency in Hz.
+    dish_diameter_m
+        The antenna dish diameter in metres.
+    pb_factor
+        The FWHM in units of `lambda / D`; `1.13` is the value for ALMA's 12 m antennas
+        (a tapered illumination broadens the beam beyond the uniform-dish value of 1.02).
+    """
+    wavelength_m = SPEED_OF_LIGHT_M_PER_S / frequency_hz
+
+    return float(pb_factor * wavelength_m / dish_diameter_m * ARCSEC_PER_RADIAN)
+
+
+def primary_beam_gaussian_from(
+    mask: Mask2D,
+    frequency_hz: float,
+    dish_diameter_m: float,
+    pb_factor: float = 1.13,
+    centre: Tuple[float, float] = (0.0, 0.0),
+) -> Array2D:
+    """
+    Returns a circular Gaussian primary beam `P` on the unmasked pixels of a real-space mask,
+    with peak value 1 at `centre` and FWHM `pb_factor * c / frequency_hz / dish_diameter_m`
+    (converted to arcseconds):
+
+        P(r) = exp(-4 ln(2) r^2 / FWHM^2)
+
+    so `P = 0.5` at `r = FWHM / 2`. Coordinates are the pixel centres of `Grid2D.from_mask`
+    (arcseconds), so the beam is evaluated on exactly the grid the transformers use.
+
+    Parameters
+    ----------
+    mask
+        The real-space mask of the interferometer dataset.
+    frequency_hz
+        The observing frequency in Hz.
+    dish_diameter_m
+        The antenna dish diameter in metres.
+    pb_factor
+        The FWHM in units of `lambda / D` (default `1.13`).
+    centre
+        The `(y, x)` arcsecond position of the beam peak (the pointing centre relative to the
+        image origin).
+    """
+    fwhm = primary_beam_fwhm_arcsec_from(
+        frequency_hz=frequency_hz,
+        dish_diameter_m=dish_diameter_m,
+        pb_factor=pb_factor,
+    )
+
+    grid = Grid2D.from_mask(mask=mask).array
+
+    r_squared = (grid[:, 0] - centre[0]) ** 2 + (grid[:, 1] - centre[1]) ** 2
+
+    values = np.exp(-4.0 * np.log(2.0) * r_squared / fwhm**2)
+
+    return Array2D(values=values, mask=mask)
+
+
+def primary_beam_slim_from(primary_beam, mask: Mask2D) -> np.ndarray:
+    """
+    Returns a primary beam as a float64 NumPy array on the slim (unmasked, 1D) grid of `mask`,
+    the representation the transformers and the sparse operator multiply by.
+
+    `primary_beam` may be an `Array2D` (on `mask`, or on another mask of the same native
+    shape, in which case its native values are read at `mask`'s unmasked pixels) or a plain
+    array of either the native shape `mask.shape_native` or the slim shape
+    `(mask.pixels_in_mask,)`.
+
+    Raises
+    ------
+    exc.DatasetException
+        If the primary beam's shape matches neither the native nor the slim shape of `mask`.
+    """
+    from autoarray import exc
+
+    if isinstance(primary_beam, Array2D):
+        if primary_beam.mask is mask or (
+            primary_beam.mask.shape_native == mask.shape_native
+            and np.array_equal(np.asarray(primary_beam.mask), np.asarray(mask))
+        ):
+            return np.asarray(primary_beam.slim.array, dtype=np.float64)
+
+        values = np.asarray(primary_beam.native.array, dtype=np.float64)
+    else:
+        values = np.asarray(getattr(primary_beam, "array", primary_beam))
+        values = values.astype(np.float64)
+
+    if values.shape == tuple(mask.shape_native):
+        return np.ascontiguousarray(values[mask.slim_to_native_tuple])
+
+    if values.shape == (mask.pixels_in_mask,):
+        return np.ascontiguousarray(values)
+
+    raise exc.DatasetException(
+        f"The primary beam has shape {values.shape}, which matches neither the native shape "
+        f"{tuple(mask.shape_native)} nor the slim shape ({mask.pixels_in_mask},) of the "
+        f"interferometer's real-space mask. Pass the primary beam as an image on the "
+        f"real-space mask (e.g. via `aa.primary_beam_gaussian_from(mask=real_space_mask, ...)`)."
+    )
diff --git a/autoarray/operators/transformer.py b/autoarray/operators/transformer.py
index 3323aea5b..9bfefdf01 100644
--- a/autoarray/operators/transformer.py
+++ b/autoarray/operators/transformer.py
@@ -11,6 +11,7 @@
 
 from autoarray.structures.arrays import array_2d_util
 from autoarray.operators import transformer_util
+from autoarray.dataset.interferometer.primary_beam import primary_beam_slim_from
 
 # nufftax pulls in jax at import (~0.7s), which sessions that never touch an
 # interferometer transformer should not pay for — deferred to _load_nufftax(),
@@ -106,11 +107,68 @@ def nufftax_exception():
     )
 
 
-class TransformerDFT:
+class _PrimaryBeamMixin:
+    """
+    The primary-beam handling shared by `TransformerDFT` and `TransformerNUFFT`.
+
+    The primary beam `P` is stored on the slim (unmasked) real-space grid as
+    `primary_beam_slim`. When set, the forward transforms `visibilities_from` and
+    `transform_mapping_matrix` multiply their input by `P` on the slim grid before the
+    Fourier transform, so they compute `F P i` (the visibilities of the apparent sky) from the
+    true-sky image `i`. The adjoint `image_from` is left un-beamed: the dirty image, dirty
+    beam and the w-tilde precision operator are the same with and without `P`.
+    """
+
+    primary_beam_slim: Optional[np.ndarray] = None
+
+    def _primary_beam_slim_set(self, primary_beam):
+        self.primary_beam_slim = (
+            None
+            if primary_beam is None
+            else primary_beam_slim_from(
+                primary_beam=primary_beam, mask=self.real_space_mask
+            )
+        )
+
+    def with_primary_beam(self, primary_beam):
+        """
+        Returns a shallow copy of this transformer whose forward transforms apply the primary
+        beam `primary_beam` (an `Array2D` on the real-space mask, or a native / slim array), or
+        none if it is `None`.
+
+        This transformer is not modified, so a transformer shared between datasets (e.g. one
+        returned by a `transformer_class` lambda) is safe. Returns `self` when the primary beam
+        is `None` and this transformer has none, so the un-beamed path is unchanged.
+        """
+        if primary_beam is None and self.primary_beam_slim is None:
+            return self
+
+        transformer = copy.copy(self)
+        transformer._primary_beam_slim_set(primary_beam=primary_beam)
+
+        return transformer
+
+    def _image_slim_beamed_from(self, image_slim, xp=np):
+        """Returns `P * image_slim`, or `image_slim` unchanged when there is no primary beam."""
+        if self.primary_beam_slim is None:
+            return image_slim
+
+        return image_slim * xp.asarray(self.primary_beam_slim)
+
+    def _mapping_matrix_beamed_from(self, mapping_matrix, xp=np):
+        """Returns `diag(P) @ mapping_matrix`, or `mapping_matrix` unchanged when there is no primary beam."""
+        if self.primary_beam_slim is None:
+            return mapping_matrix
+
+        return mapping_matrix * xp.asarray(self.primary_beam_slim)[:, None]
+
+
+class TransformerDFT(_PrimaryBeamMixin):
     def __init__(
         self,
         uv_wavelengths: np.ndarray,
         real_space_mask: Mask2D,
+        primary_beam=None,
     ):
         """
         A direct Fourier transform (DFT) operator for radio interferometric imaging.
@@ -128,6 +186,11 @@ def __init__(
             The (u, v) coordinates in wavelengths of the measured visibilities.
         real_space_mask
             The real-space mask that defines the image grid and which pixels are valid.
+        primary_beam
+            The primary beam `P` of the antennas, an `Array2D` on `real_space_mask` (or a native /
+            slim array). When set, `visibilities_from` and `transform_mapping_matrix` multiply their
+            input by `P` before the transform; `image_from` stays the un-beamed adjoint. `None`
+            (default) applies no primary beam.
 
         Attributes
         ----------
@@ -153,6 +216,8 @@ def __init__(
         self.total_visibilities = uv_wavelengths.shape[0]
         self.total_image_pixels = self.real_space_mask.pixels_in_mask
 
+        self._primary_beam_slim_set(primary_beam=primary_beam)
+
     def visibilities_from(self, image: Array2D, xp=np) -> Visibilities:
         """
         Computes the visibilities from a real-space image using the direct Fourier transform (DFT).
@@ -166,13 +231,16 @@ def visibilities_from(self, image: Array2D, xp=np) -> Visibilities:
             The real-space image to be transformed to the uv-plane. Must be defined on the
             same grid and mask as this transformer's `real_space_mask`.
 
+        If the transformer has a `primary_beam_slim` the image is multiplied by it on the slim
+        grid first, so the visibilities are those of the apparent sky `P * image`.
+
         Returns
         -------
         The complex visibilities resulting from the Fourier transform of the input image.
         """
 
         visibilities = transformer_util.visibilities_from(
-            image_1d=image.slim.array,
+            image_1d=self._image_slim_beamed_from(image.slim.array, xp=xp),
             grid_radians=self.grid.array,
             uv_wavelengths=self.uv_wavelengths,
             xp=xp,
@@ -188,6 +256,10 @@ def image_from(self, visibilities: Visibilities, xp=np) -> Array2D:
         complex visibilities back into image space. This is typically used as the first step
         in inverse imaging algorithms like CLEAN or regularized reconstruction.
 
+        The primary beam is deliberately not applied: this is the adjoint of the un-beamed
+        transform, so the dirty image, dirty beam and the w-tilde quantities built from it do
+        not depend on `primary_beam_slim`.
+
         Parameters
         ----------
         visibilities
@@ -229,18 +301,19 @@ def transform_mapping_matrix(self, mapping_matrix: np.ndarray, xp=np) -> np.ndar
         Returns
         -------
         A 2D complex-valued array of shape (n_visibilities, n_source_pixels) that maps source-plane basis
-        functions directly to the visibilities.
+        functions directly to the visibilities. If the transformer has a `primary_beam_slim` every
+        column is multiplied by it first, giving `F P M`.
         """
 
         return transformer_util.transformed_mapping_matrix_from(
-            mapping_matrix=mapping_matrix,
+            mapping_matrix=self._mapping_matrix_beamed_from(mapping_matrix, xp=xp),
             grid_radians=self.grid.array,
             uv_wavelengths=self.uv_wavelengths,
             xp=xp,
         )
 
 
-class TransformerNUFFT:
+class TransformerNUFFT(_PrimaryBeamMixin):
     def __init__(
         self,
         uv_wavelengths: np.ndarray,
@@ -248,6 +321,7 @@ def __init__(
         eps: float = 1e-12,
         chunk_size: Optional[int] = None,
         xp=np,
+        primary_beam=None,
         **kwargs,
     ):
         """
@@ -299,6 +373,11 @@ def __init__(
             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`.
+        primary_beam
+            The primary beam `P` of the antennas, an `Array2D` on `real_space_mask` (or a native /
+            slim array). When set, `visibilities_from` and `transform_mapping_matrix` multiply their
+            input by `P` on the slim grid before the scatter to the native grid; `image_from` stays
+            the un-beamed adjoint. `None` (default) applies no primary beam.
 
         Attributes
         ----------
@@ -345,6 +424,8 @@ def __init__(
         self.total_visibilities = uv_wavelengths.shape[0]
         self.total_image_pixels = real_space_mask.pixels_in_mask
 
+        self._primary_beam_slim_set(primary_beam=primary_beam)
+
     def _forward_native(self, image_native_2d, xp=np):
         """Run nufft2d2 on a 2D native-shape image array, returning visibilities.
 
@@ -412,16 +493,24 @@ def visibilities_from(self, image, xp=np) -> Visibilities:
         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 the transformer has a `primary_beam_slim` the image is multiplied by it on the
+        slim grid before the scatter to the native grid, so the visibilities are those of
+        `P * image`.
         """
         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
+                self._image_slim_beamed_from(image.array, xp=xp)
             )
-        else:
+        elif self.primary_beam_slim is None:
             image_native = image.native.array
+        else:
+            image_slim = self._image_slim_beamed_from(image.slim.array, xp=np)
+            image_native = np.zeros(image.mask.shape, dtype=image_slim.dtype)
+            image_native[image.mask.slim_to_native_tuple] = image_slim
 
         return Visibilities(visibilities=self._forward_native(image_native, xp=xp))
 
@@ -442,6 +531,10 @@ def image_from(
         `TransformerDFT.image_from` exactly, which is what makes
         `Interferometer.apply_sparse_operator` scale-consistent across both
         transformers.
+
+        The primary beam is deliberately not applied: this is the adjoint of the
+        un-beamed transform, so the dirty image, dirty beam and the w-tilde quantities
+        built from it do not depend on `primary_beam_slim`.
         """
         _load_nufftax()
 
@@ -515,9 +608,14 @@ def transform_mapping_matrix(self, mapping_matrix, xp=np):
         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).
+
+        If the transformer has a `primary_beam_slim` every column is multiplied by it on
+        the slim grid before the scatter, giving `F P M`.
         """
         _load_nufftax()
 
+        mapping_matrix = self._mapping_matrix_beamed_from(mapping_matrix, xp=xp)
+
         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

From 942da1565f149556a9077cc960b4376c78e0cfcf Mon Sep 17 00:00:00 2001
From: Jammy2211 <JNightingale2211@gmail.com>
Date: Wed, 21 Oct 2026 11:59:05 +0100
Subject: [PATCH 2/5] feat: primary beam on Interferometer, sparse operator and
 simulator

Interferometer takes primary_beam (native or slim, validated against the
real-space mask, warns when any unmasked value is 0), applies it to its
transformer via with_primary_beam and forwards it through the pooled-noise
branch, apply_sparse_operator, apply_sparse_operator_from_chunks,
from_sparse_terms, from_stream and from_fits (primary_beam_path).
InterferometerSparseOperator gains primary_beam_slim; sparse_profile_terms_from
applies image <- P image. SimulatorInterferometer takes primary_beam.

Refs #640

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
---
 autoarray/dataset/interferometer/dataset.py   | 87 ++++++++++++++++++-
 autoarray/dataset/interferometer/simulator.py | 12 ++-
 .../inversion_interferometer_util.py          | 30 +++++++
 3 files changed, 127 insertions(+), 2 deletions(-)

diff --git a/autoarray/dataset/interferometer/dataset.py b/autoarray/dataset/interferometer/dataset.py
index 3f6cd084f..287d7b3d8 100644
--- a/autoarray/dataset/interferometer/dataset.py
+++ b/autoarray/dataset/interferometer/dataset.py
@@ -1,4 +1,6 @@
+import dataclasses
 import logging
+import warnings
 import numpy as np
 from typing import Optional, Tuple
 
@@ -25,6 +27,7 @@ def disable_jax():
 
 from autoarray.dataset.abstract.dataset import AbstractDataset
 from autoarray.dataset.grids import GridsDataset
+from autoarray.dataset.interferometer.primary_beam import primary_beam_slim_from
 from autoarray.inversion.inversion.interferometer.inversion_interferometer_util import (
     InterferometerSparseOperator,
 )
@@ -54,6 +57,7 @@ def __init__(
         sparse_operator: Optional[InterferometerSparseOperator] = None,
         raise_error_dft_visibilities_limit: bool = True,
         sparse_terms: Optional["inversion_interferometer_util.SparseTerms"] = None,
+        primary_beam=None,
     ):
         """
         An interferometer dataset, containing the visibilities data, noise-map, real-space msk, Fourier transformer and
@@ -112,6 +116,17 @@ def __init__(
             `from_stream`, `from_sparse_terms` and `apply_sparse_operator_from_chunks`). It
             carries the naturally weighted dirty image and beam and the provenance of the
             accumulation.
+        primary_beam
+            The primary beam `P` of the antennas, an `Array2D` on `real_space_mask` (or a plain
+            array of its native or slim shape; see `aa.primary_beam_gaussian_from`). The
+            visibilities are then modelled as those of the apparent sky `P x I`, so every light
+            profile, linear light profile and pixelization reconstructs the true sky `I`: the
+            transformer applies `P` in its forward transforms and the sparse operator carries it
+            (as `primary_beam_slim`) for the sparse (w-tilde) inversion. The dirty image, dirty
+            beam and w-tilde precision operator do not depend on `P`. A warning is emitted if any
+            unmasked value of `P` is 0 (source pixels seen only there are unconstrained by the
+            data and set by the regularization); no floor is applied. `None` (default) applies no
+            primary beam.
 
         Array-free datasets
         -------------------
@@ -133,11 +148,34 @@ def __init__(
 
         self.uv_wavelengths = uv_wavelengths
 
+        self.primary_beam = primary_beam
+        self.primary_beam_slim = (
+            None
+            if primary_beam is None
+            else primary_beam_slim_from(primary_beam=primary_beam, mask=real_space_mask)
+        )
+
+        if self.primary_beam_slim is not None and np.any(self.primary_beam_slim == 0.0):
+            warnings.warn(
+                f"The primary beam of this Interferometer is 0 at "
+                f"{int(np.sum(self.primary_beam_slim == 0.0))} unmasked real-space pixel(s). The "
+                f"data carry no information about the sky there, so the curvature matrix of an "
+                f"inversion is rank-deficient for any source pixel seen only through those pixels "
+                f"(its value is then set by the regularization). No floor is applied; consider "
+                f"masking the pixels where the primary beam is 0."
+            )
+
         if uv_wavelengths is not None and transformer_class is not None:
             self.transformer = transformer_class(
                 uv_wavelengths=uv_wavelengths,
                 real_space_mask=real_space_mask,
             )
+            # The dataset is the source of truth for the primary beam: the transformer is
+            # replaced by a (shallow) copy carrying it, so a transformer shared through a
+            # `transformer_class` lambda is never modified.
+            self.transformer = self.transformer.with_primary_beam(
+                self.primary_beam_slim
+            )
         else:
             self.transformer = None
 
@@ -150,6 +188,12 @@ def __init__(
         self.sparse_operator = sparse_operator
         self.sparse_terms = sparse_terms
 
+        if sparse_operator is not None and self.primary_beam_slim is not None:
+            if getattr(sparse_operator, "primary_beam_slim", None) is None:
+                self.sparse_operator = dataclasses.replace(
+                    sparse_operator, primary_beam_slim=self.primary_beam_slim
+                )
+
         if raise_error_dft_visibilities_limit and self.uv_wavelengths is not None:
             if (
                 self.uv_wavelengths.shape[0] > 10000
@@ -178,6 +222,8 @@ def from_fits(
         uv_wavelengths_hdu=0,
         transformer_class=TransformerNUFFT,
         raise_error_dft_visibilities_limit: bool = True,
+        primary_beam_path=None,
+        primary_beam_hdu=0,
     ):
         """
         Load an interferometer dataset from multiple .fits files.
@@ -220,6 +266,12 @@ def from_fits(
             `TransformerDFT` and the dataset has more than 10,000 visibilities. Set to False to
             opt into the slow DFT path at ALMA-scale (e.g. when profiling the JAX-traceable
             DFT path before a JIT-friendly NUFFT is available).
+        primary_beam_path
+            The path to a .fits file containing the primary beam image, of the native shape of
+            `real_space_mask` (see `Interferometer`'s `primary_beam`). `None` (default) loads no
+            primary beam.
+        primary_beam_hdu
+            The HDU index in the primary beam .fits file from which it is loaded.
 
         Returns
         -------
@@ -237,6 +289,13 @@ def from_fits(
             file_path=uv_wavelengths_path, hdu=uv_wavelengths_hdu
         )
 
+        primary_beam = None
+
+        if primary_beam_path is not None:
+            primary_beam = ndarray_via_fits_from(
+                file_path=primary_beam_path, hdu=primary_beam_hdu
+            )
+
         return Interferometer(
             real_space_mask=real_space_mask,
             data=visibilities,
@@ -244,6 +303,7 @@ def from_fits(
             uv_wavelengths=uv_wavelengths,
             transformer_class=transformer_class,
             raise_error_dft_visibilities_limit=raise_error_dft_visibilities_limit,
+            primary_beam=primary_beam,
         )
 
     @classmethod
@@ -253,6 +313,7 @@ def from_sparse_terms(
         real_space_mask: Mask2D,
         *,
         batch_size: int = 128,
+        primary_beam=None,
     ) -> "Interferometer":
         """
         Build an array-free `Interferometer` from an accumulated `SparseTerms` record.
@@ -276,6 +337,10 @@ def from_sparse_terms(
             The real-space `Mask2D` the terms were accumulated on.
         batch_size
             The number of source-pixel columns processed per batch by the sparse operator.
+        primary_beam
+            The primary beam on `real_space_mask` (see `Interferometer`). The `SparseTerms` do
+            not depend on it, so the same terms serve any primary beam; it is carried on the
+            sparse operator as `primary_beam_slim`.
         """
         shape_native = getattr(terms, "shape_native", None)
 
@@ -310,6 +375,13 @@ def from_sparse_terms(
             terms,
             real_space_mask=real_space_mask,
             batch_size=batch_size,
+            primary_beam_slim=(
+                None
+                if primary_beam is None
+                else primary_beam_slim_from(
+                    primary_beam=primary_beam, mask=real_space_mask
+                )
+            ),
         )
 
         return cls(
@@ -320,6 +392,7 @@ def from_sparse_terms(
             transformer_class=None,
             sparse_operator=sparse_operator,
             sparse_terms=terms,
+            primary_beam=primary_beam,
         )
 
     @classmethod
@@ -340,6 +413,7 @@ def from_stream(
         pool_noise_map: bool = False,
         oversample: Optional[int] = None,
         oversample_pad: float = 0.25,
+        primary_beam=None,
     ) -> "Interferometer":
         """
         Build an array-free `Interferometer` by accumulating a stream of visibility chunks,
@@ -393,6 +467,9 @@ def from_stream(
             `sparse_terms` (`precision_operator_fine`, `dirty_image_fine`) and need the NUFFT
             transformer. Memory grows as `q^2` (~5-6 GB peak at 400 pixels, `q = 8`); see
             `sparse_terms_from_chunks`, "Fine grids".
+        primary_beam
+            The primary beam on `real_space_mask` (see `Interferometer`), passed to
+            `from_sparse_terms`. The accumulated terms do not depend on it.
 
         Raises
         ------
@@ -420,7 +497,10 @@ def from_stream(
         )
 
         return cls.from_sparse_terms(
-            terms, real_space_mask=real_space_mask, batch_size=batch_size
+            terms,
+            real_space_mask=real_space_mask,
+            batch_size=batch_size,
+            primary_beam=primary_beam,
         )
 
     @property
@@ -606,6 +686,7 @@ def apply_sparse_operator(
                 ),
                 uv_wavelengths=self.uv_wavelengths,
                 transformer_class=lambda uv_wavelengths, real_space_mask: self.transformer,
+                primary_beam=self.primary_beam,
             )
 
             pooling_record.log()
@@ -694,6 +775,7 @@ def apply_sparse_operator(
             batch_size=batch_size,
             data_term=data_term,
             noise_normalization=noise_normalization,
+            primary_beam_slim=self.primary_beam_slim,
         )
 
         return Interferometer(
@@ -703,6 +785,7 @@ def apply_sparse_operator(
             uv_wavelengths=self.uv_wavelengths,
             transformer_class=lambda uv_wavelengths, real_space_mask: self.transformer,
             sparse_operator=sparse_operator,
+            primary_beam=self.primary_beam,
         )
 
     def apply_sparse_operator_from_chunks(
@@ -860,6 +943,7 @@ def transformer_class(uv_wavelengths, real_space_mask):
             terms,
             real_space_mask=self.real_space_mask,
             batch_size=batch_size,
+            primary_beam_slim=self.primary_beam_slim,
         )
 
         return Interferometer(
@@ -870,6 +954,7 @@ def transformer_class(uv_wavelengths, real_space_mask):
             transformer_class=lambda uv_wavelengths, real_space_mask: self.transformer,
             sparse_operator=sparse_operator,
             sparse_terms=terms,
+            primary_beam=self.primary_beam,
         )
 
     def psf_precision_operator_from(
diff --git a/autoarray/dataset/interferometer/simulator.py b/autoarray/dataset/interferometer/simulator.py
index f1c26191b..754512758 100644
--- a/autoarray/dataset/interferometer/simulator.py
+++ b/autoarray/dataset/interferometer/simulator.py
@@ -54,6 +54,8 @@ def _register_interferometer_pytrees() -> None:
             "over_sample_size_pixelization",
             "noise_covariance_matrix",
             "sparse_operator",
+            "primary_beam",
+            "primary_beam_slim",
         ),
     )
 
@@ -70,6 +72,7 @@ def __init__(
         noise_if_add_noise_false=0.1,
         noise_seed=-1,
         use_jax: bool = False,
+        primary_beam=None,
     ):
         """
         Simulates observations of `Interferometer` data, including transforming a real-space image to
@@ -117,6 +120,11 @@ def __init__(
             the simulator's internal complex-Gaussian noise generation routes
             through ``jax.random``. The returned ``Interferometer`` carries
             ``jax.Array`` visibilities. Mirror of ``SimulatorImaging.use_jax``.
+        primary_beam
+            The primary beam `P` of the antennas, an `Array2D` on the image's mask (or a native /
+            slim array). The simulated visibilities are then those of the apparent sky
+            `P * image`, and the returned `Interferometer` carries `P` so a fit models the true
+            sky. `None` (default) applies no primary beam.
         """
 
         self.uv_wavelengths = uv_wavelengths
@@ -126,6 +134,7 @@ def __init__(
         self.noise_if_add_noise_false = noise_if_add_noise_false
         self.noise_seed = noise_seed
         self.use_jax = use_jax
+        self.primary_beam = primary_beam
 
     @property
     def _xp(self):
@@ -172,7 +181,7 @@ def via_image_from(self, image, xp=None):
 
         transformer = self.transformer_class(
             uv_wavelengths=self.uv_wavelengths, real_space_mask=image.mask
-        )
+        ).with_primary_beam(self.primary_beam)
 
         visibilities = transformer.visibilities_from(image=image, xp=xp)
 
@@ -204,4 +213,5 @@ def via_image_from(self, image, xp=None):
             uv_wavelengths=transformer.uv_wavelengths,
             real_space_mask=image.mask,
             transformer_class=self.transformer_class,
+            primary_beam=self.primary_beam,
         )
diff --git a/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py b/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py
index 76a3c8d1c..6539c1c89 100644
--- a/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py
+++ b/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py
@@ -1250,6 +1250,7 @@ class InterferometerSparseOperator:
     w_dtype: np.dtype
     data_term: Optional[float] = None
     noise_normalization: Optional[float] = None
+    primary_beam_slim: Optional[np.ndarray] = None
     """
     Cached FFT operator state for fast interferometer curvature-matrix assembly.
 
@@ -1327,6 +1328,12 @@ class InterferometerSparseOperator:
         visibilities, equal to `fit_util.noise_normalization_complex_from(noise_map)`. When
         set, `FitInterferometer.noise_normalization` reads it instead of reducing over the
         noise-map every likelihood call. `None` falls back to the array reduction.
+    primary_beam_slim
+        The primary beam `P` of the dataset on the slim masked real-space grid, or `None`. The
+        operator itself (`W~`, the dirty image and both scalars) is independent of `P`; the
+        sparse inversion reads it to form the beamed response `A' = P A`, giving the curvature
+        matrix `A^T P W~ P A` and data vector `A^T P d~`, and `sparse_profile_terms_from`
+        applies it to profile images.
 
     Cached properties
     -----------------
@@ -1401,6 +1408,7 @@ def from_nufft_precision_operator(
         batch_size: int = 128,
         data_term: Optional[float] = None,
         noise_normalization: Optional[float] = None,
+        primary_beam_slim: Optional[np.ndarray] = None,
     ):
         """
         Construct an `InterferometerSparseOperator` from a curvature-preload array.
@@ -1429,6 +1437,9 @@ def from_nufft_precision_operator(
             Number of source-pixel columns processed per block when assembling
             the curvature matrix. Higher values typically improve GPU efficiency
             but increase intermediate memory usage.
+        primary_beam_slim
+            The dataset's primary beam on the slim masked grid, or `None` (see the class
+            docstring). Stored as a float64 NumPy array.
 
         Returns
         -------
@@ -1468,6 +1479,11 @@ def from_nufft_precision_operator(
             noise_normalization=(
                 None if noise_normalization is None else float(noise_normalization)
             ),
+            primary_beam_slim=(
+                None
+                if primary_beam_slim is None
+                else np.asarray(primary_beam_slim, dtype=np.float64)
+            ),
         )
 
     @classmethod
@@ -1477,6 +1493,7 @@ def from_sparse_terms(
         *,
         real_space_mask,
         batch_size: int = 128,
+        primary_beam_slim: Optional[np.ndarray] = None,
     ) -> "InterferometerSparseOperator":
         """
         Construct an `InterferometerSparseOperator` from the per-visibility sums of a
@@ -1498,6 +1515,9 @@ def from_sparse_terms(
         batch_size
             Number of source-pixel columns processed per block when assembling the curvature
             matrix (see `from_nufft_precision_operator`).
+        primary_beam_slim
+            The dataset's primary beam on the slim masked grid, or `None`. The `SparseTerms`
+            are independent of it, so the same terms serve any primary beam.
         """
         from autoarray.structures.arrays.uniform_2d import Array2D
 
@@ -1511,6 +1531,7 @@ def from_sparse_terms(
             batch_size=batch_size,
             data_term=terms.data_term,
             noise_normalization=terms.noise_normalization,
+            primary_beam_slim=primary_beam_slim,
         )
 
     @staticmethod
@@ -2136,6 +2157,10 @@ def sparse_profile_terms_from(
     The data term is a difference of large numbers at high signal-to-noise (`data_term >> chi^2`); in float64
     this is accurate to ~1e-15 relative to `data_term`, ample for likelihood comparisons.
 
+    If the operator carries a `primary_beam_slim` `P`, the visibilities of `i_p` are `F P i_p`, so `i_p` is
+    replaced by `P i_p` throughout: `sparse_dirty_image = d~ - W~ P i_p` and
+    `data_term - 2 (P i_p)^T d~ + (P i_p)^T W~ (P i_p)`, and the returned `operated_image` is `W~ P i_p`.
+
     Every operation is an `xp` array operation, so under `jax.jit` the returned scalar is traced with `i_p`.
 
     Parameters
@@ -2162,6 +2187,11 @@ def sparse_profile_terms_from(
     """
     image = getattr(image, "array", image)
 
+    primary_beam_slim = getattr(sparse_operator, "primary_beam_slim", None)
+
+    if primary_beam_slim is not None:
+        image = image * xp.asarray(primary_beam_slim)
+
     operated_image = sparse_operator.operated_matrix_slim_from(
         matrix_slim=image[:, None],
         extent_index_for_masked_pixel=extent_index_for_masked_pixel,

From 5ab769169cdafa8ef3ec596b2204eba3f03bb241 Mon Sep 17 00:00:00 2001
From: Jammy2211 <JNightingale2211@gmail.com>
Date: Wed, 21 Oct 2026 11:59:05 +0100
Subject: [PATCH 3/5] feat: primary beam in the sparse w-tilde interferometer
 inversion

mapper_util.sparse_triplets_from gains slim_weights; the sparse inversion
passes the primary beam so curvature = A^T P W~ P A and data vector =
A^T P d~, scales linear-function mapping matrices by P, and keeps the
mapped reconstructed image un-beamed (true sky). dirty_model_image_natural_from
applies P. The numba direct_conv route cannot apply P, so the factory skips
it for beamed datasets and the class raises. Document that P is the caller's
job for operated_mapping_matrix_override.

Refs #640

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
---
 autoarray/fit/fit_interferometer.py           |  8 +++
 autoarray/inversion/inversion/factory.py      |  6 ++-
 .../inversion/interferometer/abstract.py      |  4 +-
 .../inversion/interferometer/sparse.py        | 52 +++++++++++++++++--
 .../inversion/interferometer_numba/sparse.py  | 10 ++++
 autoarray/inversion/linear_obj/linear_obj.py  |  4 +-
 autoarray/inversion/mappers/mapper_util.py    | 12 +++++
 7 files changed, 89 insertions(+), 7 deletions(-)

diff --git a/autoarray/fit/fit_interferometer.py b/autoarray/fit/fit_interferometer.py
index 66c46c9f1..5c088480a 100644
--- a/autoarray/fit/fit_interferometer.py
+++ b/autoarray/fit/fit_interferometer.py
@@ -32,6 +32,9 @@ def dirty_model_image_natural_from(dataset, image) -> Array2D:
         The `Interferometer` dataset, which must carry a `sparse_operator`.
     image
         The model image `m` on the slim masked real-space grid of the dataset's `real_space_mask`.
+        If the dataset's `sparse_operator` carries a `primary_beam_slim` `P`, `m` is the true-sky
+        image and the result is `W~ P m / sum(w)`, the natural dirty image of the visibilities
+        `F P m`.
     """
     sparse_operator = getattr(dataset, "sparse_operator", None)
 
@@ -51,6 +54,11 @@ def dirty_model_image_natural_from(dataset, image) -> Array2D:
 
     image = np.asarray(getattr(image, "array", image), dtype=np.float64)
 
+    primary_beam_slim = getattr(sparse_operator, "primary_beam_slim", None)
+
+    if primary_beam_slim is not None:
+        image = image * primary_beam_slim
+
     operated_image = sparse_operator.operated_matrix_slim_from(
         matrix_slim=image[:, None],
         extent_index_for_masked_pixel=dataset.real_space_mask.extent_index_for_masked_pixel,
diff --git a/autoarray/inversion/inversion/factory.py b/autoarray/inversion/inversion/factory.py
index f47c22cc5..8c7a7e1de 100644
--- a/autoarray/inversion/inversion/factory.py
+++ b/autoarray/inversion/inversion/factory.py
@@ -216,7 +216,11 @@ def inversion_interferometer_from(
     # need no mapper). `_use_interferometer_numba` returns False whenever a func-list is present.
     if dataset.sparse_operator is not None:
 
-        if _use_interferometer_numba(
+        # The numba `direct_conv` kernel uses the mapper weights as-is and cannot apply a primary
+        # beam, so a beamed dataset always takes the FFT sparse route.
+        if getattr(
+            dataset.sparse_operator, "primary_beam_slim", None
+        ) is None and _use_interferometer_numba(
             linear_obj_list=linear_obj_list,
             settings=settings,
             xp=xp,
diff --git a/autoarray/inversion/inversion/interferometer/abstract.py b/autoarray/inversion/inversion/interferometer/abstract.py
index edca3928c..037d7e406 100644
--- a/autoarray/inversion/inversion/interferometer/abstract.py
+++ b/autoarray/inversion/inversion/interferometer/abstract.py
@@ -107,7 +107,9 @@ def operated_mapping_matrix_list(self) -> List[np.ndarray]:
         A linear object may have a `operated_mapping_matrix_override` property, which bypasses the `mapping_matrix`
         computation and transformer operation and is directly placed in the `operated_mapping_matrix_list`. Because
         the override bypasses the transformer it must already be in the data's visibility space, with (complex)
-        shape [total_visibilities, params] (e.g. computed via an analytic Fourier transform).
+        shape [total_visibilities, params] (e.g. computed via an analytic Fourier transform). The dataset's primary
+        beam is not applied to the override either (the transformer applies it to every other linear object), so
+        an override for a dataset with a `primary_beam` must already include it; that is the caller's job.
         """
         operated_mapping_matrix_list = []
 
diff --git a/autoarray/inversion/inversion/interferometer/sparse.py b/autoarray/inversion/inversion/interferometer/sparse.py
index d30970853..d49c81bf1 100644
--- a/autoarray/inversion/inversion/interferometer/sparse.py
+++ b/autoarray/inversion/inversion/interferometer/sparse.py
@@ -100,6 +100,10 @@ def data_vector(self) -> np.ndarray:
         the inversion's data differs from them (e.g. a `DatasetInterface` whose data has the visibilities
         of ordinary light profiles subtracted), the dataset's `sparse_dirty_image` of that data is used
         instead, so that `D` is consistent with the data the chi-squared is computed from.
+
+        If the dataset's `sparse_operator` carries a `primary_beam_slim` `P`, every linear object's
+        real-space mapping matrix `L` is replaced by `P L` (the mapper triplets carry `P` in their
+        `vals`), giving `D = Lᵀ P d~`.
         """
         dirty_image = getattr(self.dataset, "sparse_dirty_image", None)
 
@@ -121,7 +125,10 @@ def data_vector(self) -> np.ndarray:
                 )
             else:
                 data_vector_list.append(
-                    self._xp.dot(linear_obj.mapping_matrix.T, dirty_image)
+                    self._xp.dot(
+                        self._mapping_matrix_beamed_from(linear_obj.mapping_matrix).T,
+                        dirty_image,
+                    )
                 )
 
         if len(data_vector_list) == 1:
@@ -129,7 +136,28 @@ def data_vector(self) -> np.ndarray:
 
         return self._xp.concatenate(data_vector_list)
 
-    def _sparse_triplets_slim_from(self, mapper: Mapper):
+    @property
+    def _primary_beam_slim(self):
+        """
+        The primary beam `P` on the slim masked grid carried by the dataset's `sparse_operator`
+        (`None` if it has none). Read via `getattr` so a `DatasetInterface` (e.g. PyAutoGalaxy's),
+        whose `sparse_operator` is the dataset's own, works unchanged.
+        """
+        return getattr(self.dataset.sparse_operator, "primary_beam_slim", None)
+
+    def _mapping_matrix_beamed_from(self, mapping_matrix):
+        """
+        Returns `diag(P) @ mapping_matrix` for a linear function's real-space mapping matrix, or the
+        mapping matrix unchanged when there is no primary beam.
+        """
+        primary_beam_slim = self._primary_beam_slim
+
+        if primary_beam_slim is None:
+            return mapping_matrix
+
+        return mapping_matrix * self._xp.asarray(primary_beam_slim)[:, None]
+
+    def _sparse_triplets_slim_from(self, mapper: Mapper, apply_primary_beam: bool = True):
         """
         Returns the sparse COO triplets `(rows, cols, vals)` of a mapper's real-space mapping
         operator `A`, with `rows` indexed on the slim (masked-pixel) grid the dirty image lives on.
@@ -147,6 +175,10 @@ def _sparse_triplets_slim_from(self, mapper: Mapper):
         ----------
         mapper
             The mapper whose mapping operator is expressed in sparse triplet form.
+        apply_primary_beam
+            If `True` (default) and the dataset has a primary beam `P`, the triplets are those of the
+            beamed operator `P A` (used by the data vector); `False` gives those of `A` (used for the
+            true-sky reconstructed image).
         """
         xp = self._xp
 
@@ -157,6 +189,7 @@ def _sparse_triplets_slim_from(self, mapper: Mapper):
             fft_index_for_masked_pixel=self.mask.extent_index_for_masked_pixel,
             sub_fraction_slim=mapper.over_sampler.sub_fraction.array,
             return_rows_slim=True,
+            slim_weights=self._primary_beam_slim if apply_primary_beam else None,
             xp=xp,
         )
 
@@ -209,6 +242,9 @@ def _sparse_triplets_curvature_from(self, mapper: Mapper):
         This is the interferometer counterpart of `Mapper.sparse_triplets_curvature`, which
         indexes rows on the imaging FFT grid and therefore cannot be reused here.
 
+        If the dataset has a primary beam `P` the triplets are those of the beamed operator `P A`,
+        so every curvature block built from them is `Aᵀ P W~ P A` with the FFT kernel untouched.
+
         Parameters
         ----------
         mapper
@@ -221,6 +257,7 @@ def _sparse_triplets_curvature_from(self, mapper: Mapper):
             fft_index_for_masked_pixel=self.mask.extent_index_for_masked_pixel,
             sub_fraction_slim=mapper.over_sampler.sub_fraction.array,
             return_rows_slim=False,
+            slim_weights=self._primary_beam_slim,
             xp=self._xp,
         )
 
@@ -441,8 +478,11 @@ def _curvature_matrix_func_list_and_mapper(self) -> np.ndarray:
             cls=AbstractLinearObjFuncList
         )
 
+        # With a primary beam `P` each linear function's response is `P B_k`, matching the beamed
+        # mapper triplets, so every block is formed with the beamed operator.
         mapping_matrix_list = [
-            linear_func.mapping_matrix for linear_func in linear_func_list
+            self._mapping_matrix_beamed_from(linear_func.mapping_matrix)
+            for linear_func in linear_func_list
         ]
 
         for mapper_index, mapper in enumerate(mapper_list):
@@ -590,7 +630,11 @@ def _mapped_reconstructed_data_mapper_from(
         reconstruction
             The mapper's block of the `reconstruction` (its source pixel values `s`).
         """
-        rows, cols, vals = self._sparse_triplets_slim_from(mapper=mapper)
+        # The reconstructed image is the true sky `A s`, so the primary beam is not applied; the
+        # model visibilities `F P A s` pick it up in the transformer.
+        rows, cols, vals = self._sparse_triplets_slim_from(
+            mapper=mapper, apply_primary_beam=False
+        )
 
         contrib = vals * reconstruction[cols]
 
diff --git a/autoarray/inversion/inversion/interferometer_numba/sparse.py b/autoarray/inversion/inversion/interferometer_numba/sparse.py
index d7919f0ce..ddaaf6a95 100644
--- a/autoarray/inversion/inversion/interferometer_numba/sparse.py
+++ b/autoarray/inversion/inversion/interferometer_numba/sparse.py
@@ -105,6 +105,16 @@ def _check_preconditions(self) -> None:
         NaN mapping weights are not a precondition: they propagate into `F` and give the
         same NaN evidence as the FFT route, which is the intended rejection of the sample.
         """
+        if getattr(self.dataset.sparse_operator, "primary_beam_slim", None) is not None:
+            raise exc.InversionException(
+                "`InversionInterferometerSparseNumba` was passed a dataset with a primary "
+                "beam (`sparse_operator.primary_beam_slim`). The `direct_conv` kernel uses "
+                "`pix_weights_for_sub_slim_index` as-is, so it would silently compute the "
+                "un-beamed F.\n\n"
+                "Use `InversionInterferometerSparse`, which folds the primary beam into the "
+                "mapper triplets."
+            )
+
         if self.has(cls=AbstractLinearObjFuncList):
             raise exc.InversionException(
                 "A linear-function list (e.g. a linear light profile or MGE basis) was "
diff --git a/autoarray/inversion/linear_obj/linear_obj.py b/autoarray/inversion/linear_obj/linear_obj.py
index 52372adb2..a9210b756 100644
--- a/autoarray/inversion/linear_obj/linear_obj.py
+++ b/autoarray/inversion/linear_obj/linear_obj.py
@@ -141,7 +141,9 @@ def operated_mapping_matrix_override(self) -> Optional[np.ndarray]:
 
         - `Interferometer`: a complex matrix of dimensions (total_visibilities, total_parameters), e.g. the
           visibilities of each linear object's parameter computed via an analytic Fourier transform. The
-          transformer (NUFFT / DFT) is not applied to the override.
+          transformer (NUFFT / DFT) is not applied to the override, and neither is the dataset's primary
+          beam: if the dataset has a `primary_beam`, applying it (computing the visibilities of `P` times
+          the linear object's image) is the caller's job.
 
         Returns
         -------
diff --git a/autoarray/inversion/mappers/mapper_util.py b/autoarray/inversion/mappers/mapper_util.py
index 149ccf864..abaef95df 100644
--- a/autoarray/inversion/mappers/mapper_util.py
+++ b/autoarray/inversion/mappers/mapper_util.py
@@ -275,6 +275,7 @@ def sparse_triplets_from(
     sub_fraction_slim,  # (N_unmasked,)
     *,
     return_rows_slim: bool = True,
+    slim_weights=None,
     xp=np,
 ):
     """
@@ -297,6 +298,11 @@ def sparse_triplets_from(
         Mapping slim pixel -> rectangular FFT-grid pixel index (N_unmasked,)
     sub_fraction_slim
         Oversampling normalization per slim pixel (N_unmasked,)
+    slim_weights
+        Optional per-slim-pixel weights (N_unmasked,) multiplied into `vals`, i.e. the triplets
+        of `diag(w) A` rather than `A`. The interferometer sparse inversion passes the primary
+        beam here so the curvature and data vector use the beamed response `P A`. `None`
+        (default) applies no weighting.
     xp
         Backend module (np or jnp)
 
@@ -331,6 +337,9 @@ def sparse_triplets_from(
         slim_rows = slim_index_for_sub[sub_ids]  # int32
         vals = vals * sub_fraction_slim[slim_rows]  # float64
 
+        if slim_weights is not None:
+            vals = vals * np.asarray(slim_weights, dtype=np.float64)[slim_rows]
+
         if return_rows_slim:
             return slim_rows, cols, vals
 
@@ -358,6 +367,9 @@ def sparse_triplets_from(
     slim_rows = slim_index_for_sub[sub_ids]
     vals = vals * sub_fraction_slim[slim_rows]
 
+    if slim_weights is not None:
+        vals = vals * xp.asarray(slim_weights, dtype=xp.float64)[slim_rows]
+
     if return_rows_slim:
         return slim_rows, cols, vals
 

From fb880e0267f8aaa5efe70836c0e1b2745919ccee Mon Sep 17 00:00:00 2001
From: Jammy2211 <JNightingale2211@gmail.com>
Date: Wed, 21 Oct 2026 12:06:15 +0100
Subject: [PATCH 4/5] fix: keep the numba interferometer route for a primary
 beam of ones

primary_beam_is_trivial (None or every value 1) gates the numba direct_conv
route and its precondition, so primary_beam=ones stays bit-for-bit
primary_beam=None on the default CPU route.

Refs #640

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
---
 autoarray/inversion/inversion/factory.py      | 19 +++++++++++++++----
 .../inversion/interferometer_numba/sparse.py  |  8 ++++++--
 2 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/autoarray/inversion/inversion/factory.py b/autoarray/inversion/inversion/factory.py
index 8c7a7e1de..fd4af6932 100644
--- a/autoarray/inversion/inversion/factory.py
+++ b/autoarray/inversion/inversion/factory.py
@@ -169,6 +169,16 @@ def inversion_imaging_from(
     )
 
 
+def primary_beam_is_trivial(primary_beam_slim) -> bool:
+    """
+    Whether a primary beam leaves the interferometer response unchanged: `None`, or every value exactly 1.
+    """
+    if primary_beam_slim is None:
+        return True
+
+    return bool(np.all(np.asarray(primary_beam_slim) == 1.0))
+
+
 def inversion_interferometer_from(
     dataset: Union[Interferometer, DatasetInterface],
     linear_obj_list: List[LinearObj],
@@ -217,10 +227,11 @@ def inversion_interferometer_from(
     if dataset.sparse_operator is not None:
 
         # The numba `direct_conv` kernel uses the mapper weights as-is and cannot apply a primary
-        # beam, so a beamed dataset always takes the FFT sparse route.
-        if getattr(
-            dataset.sparse_operator, "primary_beam_slim", None
-        ) is None and _use_interferometer_numba(
+        # beam, so a dataset with a non-trivial one (any value other than 1) always takes the FFT
+        # sparse route. `P = ones` keeps the numba route, so it is bit-for-bit `P = None`.
+        if primary_beam_is_trivial(
+            getattr(dataset.sparse_operator, "primary_beam_slim", None)
+        ) and _use_interferometer_numba(
             linear_obj_list=linear_obj_list,
             settings=settings,
             xp=xp,
diff --git a/autoarray/inversion/inversion/interferometer_numba/sparse.py b/autoarray/inversion/inversion/interferometer_numba/sparse.py
index ddaaf6a95..05beed804 100644
--- a/autoarray/inversion/inversion/interferometer_numba/sparse.py
+++ b/autoarray/inversion/inversion/interferometer_numba/sparse.py
@@ -105,10 +105,14 @@ def _check_preconditions(self) -> None:
         NaN mapping weights are not a precondition: they propagate into `F` and give the
         same NaN evidence as the FFT route, which is the intended rejection of the sample.
         """
-        if getattr(self.dataset.sparse_operator, "primary_beam_slim", None) is not None:
+        from autoarray.inversion.inversion.factory import primary_beam_is_trivial
+
+        if not primary_beam_is_trivial(
+            getattr(self.dataset.sparse_operator, "primary_beam_slim", None)
+        ):
             raise exc.InversionException(
                 "`InversionInterferometerSparseNumba` was passed a dataset with a primary "
-                "beam (`sparse_operator.primary_beam_slim`). The `direct_conv` kernel uses "
+                "beam (`sparse_operator.primary_beam_slim`) that is not 1 everywhere. The `direct_conv` kernel uses "
                 "`pix_weights_for_sub_slim_index` as-is, so it would silently compute the "
                 "un-beamed F.\n\n"
                 "Use `InversionInterferometerSparse`, which folds the primary beam into the "

From 8c73a9fb104b553442ae8cd59b3f645b77340702 Mon Sep 17 00:00:00 2001
From: Jammy2211 <JNightingale2211@gmail.com>
Date: Wed, 21 Oct 2026 12:06:15 +0100
Subject: [PATCH 5/5] test: primary beam witness, parity and propagation tests

Bit-for-bit P=None / P=ones parity (numpy and jit), brute-force
F = L^T P W~ P L and D = L^T P d~ on the 7x7 Delaunay fixture for every
block structure, dense vs sparse, DFT vs NUFFT, array-free vs in-memory,
transformer, dataset propagation / validation / warning, simulator,
Gaussian helper and slim_weights tests.

Refs #640

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
---
 .../dataset/interferometer/test_dataset.py    | 228 +++++++
 .../interferometer/test_primary_beam.py       |  75 +++
 .../dataset/interferometer/test_simulator.py  |  89 +++
 .../interferometer/test_interferometer.py     | 556 ++++++++++++++++++
 .../test_inversion_interferometer_util.py     |  66 +++
 test_autoarray/operators/test_transformer.py  | 209 +++++++
 6 files changed, 1223 insertions(+)
 create mode 100644 test_autoarray/dataset/interferometer/test_primary_beam.py

diff --git a/test_autoarray/dataset/interferometer/test_dataset.py b/test_autoarray/dataset/interferometer/test_dataset.py
index 23e4447f7..a7f003449 100644
--- a/test_autoarray/dataset/interferometer/test_dataset.py
+++ b/test_autoarray/dataset/interferometer/test_dataset.py
@@ -1253,3 +1253,231 @@ def test__apply_sparse_operator_from_chunks__pool_noise_map__raises(mask_2d_7x7)
         dataset.apply_sparse_operator_from_chunks(
             _chunks_of(dataset, [0, 40]), pool_noise_map=True, method="numpy"
         )
+
+
+def _primary_beam_from(mask, seed=8):
+    """A non-trivial primary beam on the unmasked pixels of `mask`, with values in (0.2, 1]."""
+    return aa.Array2D(
+        values=np.random.default_rng(seed=seed).uniform(
+            0.2, 1.0, size=mask.pixels_in_mask
+        ),
+        mask=mask,
+    )
+
+
+def _assert_carries_primary_beam(dataset, primary_beam_slim):
+    assert np.array_equal(dataset.primary_beam_slim, primary_beam_slim)
+    assert np.array_equal(dataset.sparse_operator.primary_beam_slim, primary_beam_slim)
+
+    if dataset.transformer is not None:
+        assert np.array_equal(dataset.transformer.primary_beam_slim, primary_beam_slim)
+
+
+def test__primary_beam__native_and_slim_accepted_and_applied_to_the_transformer(
+    mask_2d_7x7,
+):
+    primary_beam = _primary_beam_from(mask_2d_7x7)
+
+    for value in (
+        primary_beam,
+        primary_beam.native,
+        primary_beam.array,
+        primary_beam.native.array,
+    ):
+        dataset = _random_interferometer(mask_2d_7x7, transformer.TransformerDFT)
+        dataset = aa.Interferometer(
+            data=dataset.data,
+            noise_map=dataset.noise_map,
+            uv_wavelengths=dataset.uv_wavelengths,
+            real_space_mask=mask_2d_7x7,
+            transformer_class=transformer.TransformerDFT,
+            primary_beam=value,
+        )
+
+        assert dataset.primary_beam_slim.shape == (mask_2d_7x7.pixels_in_mask,)
+        assert np.array_equal(dataset.primary_beam_slim, primary_beam.array)
+        assert np.array_equal(
+            dataset.transformer.primary_beam_slim, primary_beam.array
+        )
+
+    dataset = _random_interferometer(mask_2d_7x7, transformer.TransformerDFT)
+
+    assert dataset.primary_beam is None
+    assert dataset.primary_beam_slim is None
+    assert dataset.transformer.primary_beam_slim is None
+
+
+def test__primary_beam__shape_mismatch__raises(mask_2d_7x7):
+    dataset = _random_interferometer(mask_2d_7x7, transformer.TransformerDFT)
+
+    for value in (np.ones((5, 5)), np.ones(mask_2d_7x7.pixels_in_mask + 1)):
+        with pytest.raises(aa.exc.DatasetException):
+            aa.Interferometer(
+                data=dataset.data,
+                noise_map=dataset.noise_map,
+                uv_wavelengths=dataset.uv_wavelengths,
+                real_space_mask=mask_2d_7x7,
+                transformer_class=transformer.TransformerDFT,
+                primary_beam=value,
+            )
+
+
+def test__primary_beam__zero_unmasked_value__warns(mask_2d_7x7):
+    dataset = _random_interferometer(mask_2d_7x7, transformer.TransformerDFT)
+
+    primary_beam = _primary_beam_from(mask_2d_7x7).array.copy()
+    primary_beam[0] = 0.0
+
+    with pytest.warns(UserWarning, match="primary beam"):
+        dataset_beamed = aa.Interferometer(
+            data=dataset.data,
+            noise_map=dataset.noise_map,
+            uv_wavelengths=dataset.uv_wavelengths,
+            real_space_mask=mask_2d_7x7,
+            transformer_class=transformer.TransformerDFT,
+            primary_beam=primary_beam,
+        )
+
+    # No floor is applied.
+    assert dataset_beamed.primary_beam_slim[0] == 0.0
+
+    import warnings
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+
+        aa.Interferometer(
+            data=dataset.data,
+            noise_map=dataset.noise_map,
+            uv_wavelengths=dataset.uv_wavelengths,
+            real_space_mask=mask_2d_7x7,
+            transformer_class=transformer.TransformerDFT,
+            primary_beam=_primary_beam_from(mask_2d_7x7),
+        )
+
+
+def test__primary_beam__dirty_image_and_dirty_beam_unchanged(mask_2d_7x7):
+    dataset = _random_interferometer(mask_2d_7x7, transformer.TransformerDFT)
+
+    dataset_beamed = aa.Interferometer(
+        data=dataset.data,
+        noise_map=dataset.noise_map,
+        uv_wavelengths=dataset.uv_wavelengths,
+        real_space_mask=mask_2d_7x7,
+        transformer_class=transformer.TransformerDFT,
+        primary_beam=_primary_beam_from(mask_2d_7x7),
+    )
+
+    for name in ("dirty_image", "dirty_image_natural", "dirty_beam"):
+        assert np.array_equal(
+            getattr(dataset_beamed, name).array, getattr(dataset, name).array
+        ), name
+
+    sparse = dataset.apply_sparse_operator(use_jax=False).sparse_operator
+    sparse_beamed = dataset_beamed.apply_sparse_operator(use_jax=False).sparse_operator
+
+    assert np.array_equal(sparse_beamed.dirty_image, sparse.dirty_image)
+    assert np.array_equal(
+        sparse_beamed.nufft_precision_operator, sparse.nufft_precision_operator
+    )
+    assert sparse_beamed.data_term == sparse.data_term
+    assert sparse.primary_beam_slim is None
+
+
+def test__primary_beam__propagates_through_every_reconstruction_site(
+    mask_2d_7x7, tmp_path
+):
+    pytest.importorskip("nufftax")
+
+    primary_beam = _primary_beam_from(mask_2d_7x7)
+
+    dataset = _random_interferometer(mask_2d_7x7, transformer.TransformerNUFFT)
+    dataset = aa.Interferometer(
+        data=dataset.data,
+        noise_map=dataset.noise_map,
+        uv_wavelengths=dataset.uv_wavelengths,
+        real_space_mask=mask_2d_7x7,
+        transformer_class=transformer.TransformerNUFFT,
+        primary_beam=primary_beam,
+    )
+
+    # apply_sparse_operator.
+    _assert_carries_primary_beam(
+        dataset.apply_sparse_operator(use_jax=False), primary_beam.array
+    )
+
+    # The pooled-noise branch.
+    _assert_carries_primary_beam(
+        dataset.apply_sparse_operator(use_jax=False, pool_noise_map=True),
+        primary_beam.array,
+    )
+
+    # apply_sparse_operator_from_chunks.
+    chunks = _chunks_of(dataset, [0, 13, 40])
+
+    _assert_carries_primary_beam(
+        dataset.apply_sparse_operator_from_chunks(chunks), primary_beam.array
+    )
+
+    # from_stream and from_sparse_terms (array-free).
+    dataset_stream = aa.Interferometer.from_stream(
+        _chunks_of(dataset, [0, 13, 40]), mask_2d_7x7, primary_beam=primary_beam
+    )
+
+    _assert_carries_primary_beam(dataset_stream, primary_beam.array)
+
+    dataset_terms = aa.Interferometer.from_sparse_terms(
+        dataset_stream.sparse_terms, mask_2d_7x7, primary_beam=primary_beam.native
+    )
+
+    _assert_carries_primary_beam(dataset_terms, primary_beam.array)
+
+    # from_fits with a primary_beam_path.
+    from astropy.io import fits
+
+    paths = {}
+
+    for name, values in (
+        ("data", np.stack([dataset.data.array.real, dataset.data.array.imag], 1)),
+        (
+            "noise_map",
+            np.stack([dataset.noise_map.array.real, dataset.noise_map.array.imag], 1),
+        ),
+        ("uv_wavelengths", dataset.uv_wavelengths),
+        ("primary_beam", primary_beam.native.array),
+    ):
+        paths[name] = tmp_path / f"{name}.fits"
+        fits.PrimaryHDU(data=values).writeto(paths[name], overwrite=True)
+
+    dataset_fits = aa.Interferometer.from_fits(
+        data_path=paths["data"],
+        noise_map_path=paths["noise_map"],
+        uv_wavelengths_path=paths["uv_wavelengths"],
+        real_space_mask=mask_2d_7x7,
+        primary_beam_path=paths["primary_beam"],
+    )
+
+    assert np.array_equal(dataset_fits.primary_beam_slim, primary_beam.array)
+    assert np.array_equal(
+        dataset_fits.transformer.primary_beam_slim, primary_beam.array
+    )
+
+    # A precomputed operator without a primary beam picks up the dataset's.
+    operator = (
+        _random_interferometer(mask_2d_7x7, transformer.TransformerNUFFT)
+        .apply_sparse_operator(use_jax=False)
+        .sparse_operator
+    )
+
+    dataset_cached = aa.Interferometer(
+        data=dataset.data,
+        noise_map=dataset.noise_map,
+        uv_wavelengths=dataset.uv_wavelengths,
+        real_space_mask=mask_2d_7x7,
+        transformer_class=transformer.TransformerNUFFT,
+        sparse_operator=operator,
+        primary_beam=primary_beam,
+    )
+
+    _assert_carries_primary_beam(dataset_cached, primary_beam.array)
+    assert operator.primary_beam_slim is None
diff --git a/test_autoarray/dataset/interferometer/test_primary_beam.py b/test_autoarray/dataset/interferometer/test_primary_beam.py
new file mode 100644
index 000000000..d21e521e4
--- /dev/null
+++ b/test_autoarray/dataset/interferometer/test_primary_beam.py
@@ -0,0 +1,75 @@
+import numpy as np
+import pytest
+
+import autoarray as aa
+from autoarray.dataset.interferometer import primary_beam as primary_beam_util
+
+
+def test__primary_beam_gaussian_from__peak_1_at_centre_and_half_at_half_fwhm():
+    frequency_hz = 230.0e9
+    dish_diameter_m = 12.0
+
+    fwhm = primary_beam_util.primary_beam_fwhm_arcsec_from(
+        frequency_hz=frequency_hz, dish_diameter_m=dish_diameter_m
+    )
+
+    # 1.13 lambda / D for ALMA's 12 m antennas at 230 GHz is ~25.3 arcsec.
+    assert fwhm == pytest.approx(
+        1.13 * 299792458.0 / frequency_hz / dish_diameter_m * 206264.80624709636,
+        rel=1.0e-12,
+    )
+    assert fwhm == pytest.approx(25.3, abs=0.1)
+
+    # An odd 41 x 41 grid whose pixel scale puts a pixel centre exactly at x = FWHM / 2.
+    pixel_scale = (fwhm / 2.0) / 10.0
+
+    mask = aa.Mask2D.all_false(shape_native=(41, 41), pixel_scales=pixel_scale)
+
+    primary_beam = aa.primary_beam_gaussian_from(
+        mask=mask, frequency_hz=frequency_hz, dish_diameter_m=dish_diameter_m
+    )
+
+    assert isinstance(primary_beam, aa.Array2D)
+    assert primary_beam.native[20, 20] == pytest.approx(1.0, abs=1.0e-14)
+    assert primary_beam.native[20, 30] == pytest.approx(0.5, rel=1.0e-12)
+    assert primary_beam.native[10, 20] == pytest.approx(0.5, rel=1.0e-12)
+    assert np.max(primary_beam.array) == pytest.approx(1.0, abs=1.0e-14)
+
+    # An offset centre moves the peak.
+    primary_beam = aa.primary_beam_gaussian_from(
+        mask=mask,
+        frequency_hz=frequency_hz,
+        dish_diameter_m=dish_diameter_m,
+        centre=(0.0, 10.0 * pixel_scale),
+    )
+
+    assert primary_beam.native[20, 30] == pytest.approx(1.0, abs=1.0e-14)
+    assert primary_beam.native[20, 20] == pytest.approx(0.5, rel=1.0e-12)
+
+
+def test__primary_beam_gaussian_from__pb_factor_scales_the_fwhm():
+    fwhm_113 = primary_beam_util.primary_beam_fwhm_arcsec_from(
+        frequency_hz=100.0e9, dish_diameter_m=12.0
+    )
+    fwhm_102 = primary_beam_util.primary_beam_fwhm_arcsec_from(
+        frequency_hz=100.0e9, dish_diameter_m=12.0, pb_factor=1.02
+    )
+
+    assert fwhm_102 / fwhm_113 == pytest.approx(1.02 / 1.13, rel=1.0e-12)
+
+
+def test__primary_beam_gaussian_from__masked_grid_matches_the_unmasked_values():
+    mask_full = aa.Mask2D.all_false(shape_native=(9, 9), pixel_scales=1.0)
+    mask = aa.Mask2D.circular(shape_native=(9, 9), pixel_scales=1.0, radius=3.0)
+
+    beam_full = aa.primary_beam_gaussian_from(
+        mask=mask_full, frequency_hz=1.0e12, dish_diameter_m=12.0
+    )
+    beam = aa.primary_beam_gaussian_from(
+        mask=mask, frequency_hz=1.0e12, dish_diameter_m=12.0
+    )
+
+    assert beam.shape_slim == mask.pixels_in_mask
+    np.testing.assert_allclose(
+        beam.array, beam_full.native.array[mask.slim_to_native_tuple], rtol=1.0e-14
+    )
diff --git a/test_autoarray/dataset/interferometer/test_simulator.py b/test_autoarray/dataset/interferometer/test_simulator.py
index 283a7416f..f8d9a40cd 100644
--- a/test_autoarray/dataset/interferometer/test_simulator.py
+++ b/test_autoarray/dataset/interferometer/test_simulator.py
@@ -52,3 +52,92 @@ def test__setup_with_noise(uv_wavelengths_7x2, transformer_7x7_7):
     assert dataset.data[0] == pytest.approx(-0.005364 - 2.36682j, 1.0e-4)
 
     assert (dataset.noise_map == 0.1 + 0.1j * np.ones((7,))).all()
+
+
+def _primary_beam_simulator_setup(transformer_class):
+    mask = aa.Mask2D.all_false(shape_native=(5, 6), pixel_scales=0.3)
+
+    rng = np.random.default_rng(seed=6)
+    uv_wavelengths = rng.normal(size=(17, 2)) * 1.0e5
+
+    image = aa.Array2D(values=rng.normal(size=mask.pixels_in_mask), mask=mask)
+    primary_beam = aa.Array2D(
+        values=rng.uniform(0.2, 1.0, size=mask.pixels_in_mask), mask=mask
+    )
+
+    def simulator_from(primary_beam):
+        return aa.SimulatorInterferometer(
+            exposure_time=1.0,
+            transformer_class=transformer_class,
+            uv_wavelengths=uv_wavelengths,
+            noise_sigma=None,
+            primary_beam=primary_beam,
+        )
+
+    return mask, image, primary_beam, simulator_from
+
+
+def test__primary_beam__simulated_visibilities_equal_unbeamed_simulation_of_beamed_image():
+    transformer_classes = [aa.TransformerDFT]
+
+    try:
+        import nufftax  # noqa: F401
+
+        transformer_classes.append(aa.TransformerNUFFT)
+    except ModuleNotFoundError:
+        pass
+
+    for transformer_class in transformer_classes:
+        mask, image, primary_beam, simulator_from = _primary_beam_simulator_setup(
+            transformer_class
+        )
+
+        dataset = simulator_from(primary_beam).via_image_from(image=image)
+        dataset_unbeamed = simulator_from(None).via_image_from(
+            image=aa.Array2D(values=image.array * primary_beam.array, mask=mask)
+        )
+
+        np.testing.assert_allclose(
+            dataset.data.array, dataset_unbeamed.data.array, rtol=1.0e-12, atol=1.0e-12
+        )
+
+        # The simulated dataset carries the primary beam, so fitting it models the true sky.
+        assert np.array_equal(dataset.primary_beam_slim, primary_beam.array)
+        assert np.array_equal(dataset.transformer.primary_beam_slim, primary_beam.array)
+        assert dataset_unbeamed.primary_beam_slim is None
+
+        # With P = None the simulation is bit-for-bit the one without the argument.
+        dataset_none = aa.SimulatorInterferometer(
+            exposure_time=1.0,
+            transformer_class=transformer_class,
+            uv_wavelengths=simulator_from(None).uv_wavelengths,
+            noise_sigma=None,
+        ).via_image_from(image=image)
+
+        assert np.array_equal(
+            simulator_from(None).via_image_from(image=image).data.array,
+            dataset_none.data.array,
+        )
+
+
+def test__primary_beam__simulator_jax_jit_matches_numpy():
+    jax = pytest.importorskip("jax")
+    import jax.numpy as jnp
+
+    mask, image, primary_beam, simulator_from = _primary_beam_simulator_setup(
+        aa.TransformerDFT
+    )
+
+    simulator = simulator_from(primary_beam)
+
+    dataset_numpy = simulator.via_image_from(image=image)
+
+    data_jax = jax.jit(
+        lambda values: simulator.via_image_from(
+            image=aa.Array2D(values=values, mask=mask), xp=jnp
+        ).data.array
+    )(jnp.asarray(image.array))
+
+    np.testing.assert_allclose(
+        np.asarray(data_jax), dataset_numpy.data.array, rtol=1.0e-10, atol=1.0e-10
+    )
diff --git a/test_autoarray/inversion/inversion/interferometer/test_interferometer.py b/test_autoarray/inversion/inversion/interferometer/test_interferometer.py
index a683a8408..37d84fa7a 100644
--- a/test_autoarray/inversion/inversion/interferometer/test_interferometer.py
+++ b/test_autoarray/inversion/inversion/interferometer/test_interferometer.py
@@ -2116,3 +2116,559 @@ def images_and_reconstruction_from(dirty_image):
             inversion=inversion_np,
             rtol=1.0e-12,
         )
+
+
+# ----------------------------------------------------------------------------------------
+# Primary beam
+# ----------------------------------------------------------------------------------------
+
+
+_PRIMARY_BEAM_SETTINGS = aa.Settings(no_regularization_add_to_curvature_diag_value=0.0)
+
+
+def _primary_beam_parity_setup(seed=0, primary_beam="random", transformer_class=None):
+    """
+    The 7x7 / Delaunay setup of `_sparse_parity_setup` with a primary beam on the dataset, returning
+    the mask, a random generator, the dense and sparse-operator datasets and the primary beam.
+
+    `primary_beam="random"` draws a non-trivial beam with values in (0.2, 1]; anything else is passed
+    to the dataset as is (`None`, ones, ...). The sparse operator is built with the NumPy brute-force
+    preload builder, so the brute-force comparison below is limited only by float64 round-off.
+    """
+    mask = 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,
+    )
+
+    n_visibilities = 5
+    rng = np.random.default_rng(seed=seed)
+    data = aa.Visibilities(
+        visibilities=rng.normal(size=(n_visibilities, 2)).astype(np.float64)
+    )
+    sigma = rng.uniform(0.5, 2.0, size=n_visibilities)
+    noise_map = aa.VisibilitiesNoiseMap(visibilities=np.stack([sigma, sigma], axis=1))
+    uv_wavelengths = rng.normal(size=(n_visibilities, 2)).astype(np.float64)
+
+    if isinstance(primary_beam, str):
+        primary_beam = aa.Array2D(
+            values=np.random.default_rng(seed=seed + 100).uniform(
+                0.2, 1.0, size=mask.pixels_in_mask
+            ),
+            mask=mask,
+        )
+
+    dataset = aa.Interferometer(
+        data=data,
+        noise_map=noise_map,
+        uv_wavelengths=uv_wavelengths,
+        real_space_mask=mask,
+        transformer_class=transformer_class or aa.TransformerDFT,
+        primary_beam=primary_beam,
+    )
+
+    dataset_sparse = dataset.apply_sparse_operator(
+        nufft_precision_operator=dataset.psf_precision_operator_from(method="numpy"),
+        batch_size=4,
+    )
+
+    return mask, rng, dataset, dataset_sparse, primary_beam
+
+
+def _primary_beam_linear_obj_lists(mask, rng):
+    """The linear-object lists exercising every block: one mapper, func list + mapper, two mappers,
+    func list + two mappers and a func list on its own."""
+    mapper_0 = _mapper_from(
+        mask=mask, pixels=9, shape=(3, 3), regularization=aa.reg.Constant(1.0)
+    )
+    mapper_1 = _mapper_from(
+        mask=mask, pixels=16, shape=(4, 4), regularization=aa.reg.Constant(2.0)
+    )
+    linear_obj = aa.m.MockLinearObjFuncList(
+        parameters=2, mapping_matrix=rng.normal(size=(mask.pixels_in_mask, 2))
+    )
+
+    return {
+        "mapper": [mapper_0],
+        "func_list_and_mapper": [linear_obj, mapper_0],
+        "x2_mappers": [mapper_0, mapper_1],
+        "func_list_and_x2_mappers": [linear_obj, mapper_0, mapper_1],
+        "func_list_only": [linear_obj],
+    }
+
+
+def _brute_force_w_tilde_and_dirty_image_from(dataset):
+    """
+    `W~ = Re(F^H N^-1 F)` and `d~ = Re(F^H N^-1 d)` computed densely from the un-beamed DFT response of
+    every image pixel, `T = transform_mapping_matrix(I)`.
+    """
+    transformer = aa.TransformerDFT(
+        uv_wavelengths=dataset.uv_wavelengths, real_space_mask=dataset.real_space_mask
+    )
+
+    assert transformer.primary_beam_slim is None
+
+    response = transformer.transform_mapping_matrix(
+        mapping_matrix=np.eye(dataset.real_space_mask.pixels_in_mask)
+    )
+
+    weights_real = dataset.noise_map.array.real**-2.0
+    weights_imag = dataset.noise_map.array.imag**-2.0
+
+    w_tilde = (response.real.T * weights_real) @ response.real + (
+        response.imag.T * weights_imag
+    ) @ response.imag
+    dirty_image = response.real.T @ (
+        dataset.data.array.real * weights_real
+    ) + response.imag.T @ (dataset.data.array.imag * weights_imag)
+
+    return w_tilde, dirty_image
+
+
+def test__primary_beam__sparse_curvature_and_data_vector_match_brute_force():
+    """
+    The witness: with a non-trivial primary beam `P`, the sparse (w-tilde) inversion's curvature matrix is
+    `Lᵀ diag(P) W~ diag(P) L` and its data vector `Lᵀ diag(P) d~`, with `W~` and `d~` built densely from the
+    un-beamed DFT, for every block structure (mappers, linear function lists and both).
+    """
+    mask, rng, dataset, dataset_sparse, primary_beam = _primary_beam_parity_setup()
+
+    w_tilde, dirty_image = _brute_force_w_tilde_and_dirty_image_from(dataset)
+
+    # The operator's dirty image is independent of P and is the brute-force d~.
+    np.testing.assert_allclose(
+        dataset_sparse.sparse_operator.dirty_image,
+        dirty_image,
+        rtol=1.0e-10,
+        atol=1.0e-10 * np.abs(dirty_image).max(),
+    )
+
+    P = primary_beam.array
+
+    for name, linear_obj_list in _primary_beam_linear_obj_lists(mask, rng).items():
+        inversion = aa.Inversion(
+            dataset=dataset_sparse,
+            linear_obj_list=linear_obj_list,
+            settings=_PRIMARY_BEAM_SETTINGS,
+        )
+
+        assert type(inversion) is aa.InversionInterferometerSparse, name
+
+        mapping_matrix = np.hstack(
+            [np.asarray(linear_obj.mapping_matrix) for linear_obj in linear_obj_list]
+        )
+        beamed = P[:, None] * mapping_matrix
+
+        curvature_matrix = beamed.T @ w_tilde @ beamed
+        data_vector = beamed.T @ dirty_image
+
+        np.testing.assert_allclose(
+            np.asarray(inversion.curvature_matrix),
+            curvature_matrix,
+            rtol=1.0e-10,
+            atol=1.0e-10 * np.abs(curvature_matrix).max(),
+            err_msg=name,
+        )
+        np.testing.assert_allclose(
+            np.asarray(inversion.data_vector),
+            data_vector,
+            rtol=1.0e-10,
+            atol=1.0e-10 * np.abs(data_vector).max(),
+            err_msg=name,
+        )
+
+        # And P genuinely changes them: the un-beamed brute force differs.
+        assert not np.allclose(
+            inversion.data_vector, mapping_matrix.T @ dirty_image, rtol=1.0e-3
+        ), name
+
+
+def test__primary_beam__none_and_ones__bit_for_bit_unchanged__numpy():
+    """
+    `primary_beam=None` and `primary_beam=ones` leave the sparse and mapping curvature matrices, data
+    vectors and `fast_chi_squared` bit-for-bit those of a dataset built without a primary beam.
+    """
+    mask, rng, dataset_ref, dataset_sparse_ref, _ = _primary_beam_parity_setup(
+        primary_beam=None
+    )
+
+    for primary_beam in (None, np.ones(mask.pixels_in_mask)):
+        _, _, dataset, dataset_sparse, _ = _primary_beam_parity_setup(
+            primary_beam=primary_beam
+        )
+
+        for name, linear_obj_list in _primary_beam_linear_obj_lists(
+            mask, np.random.default_rng(seed=1)
+        ).items():
+            for dataset_test, dataset_reference in (
+                (dataset_sparse, dataset_sparse_ref),
+                (dataset, dataset_ref),
+            ):
+                inversion = aa.Inversion(
+                    dataset=dataset_test, linear_obj_list=linear_obj_list
+                )
+                inversion_reference = aa.Inversion(
+                    dataset=dataset_reference, linear_obj_list=linear_obj_list
+                )
+
+                assert type(inversion) is type(inversion_reference)
+
+                for quantity in ("curvature_matrix", "data_vector", "reconstruction"):
+                    assert np.array_equal(
+                        np.asarray(getattr(inversion, quantity)),
+                        np.asarray(getattr(inversion_reference, quantity)),
+                    ), (name, quantity)
+
+                assert inversion.fast_chi_squared == inversion_reference.fast_chi_squared
+
+                for linear_obj in linear_obj_list:
+                    assert np.array_equal(
+                        inversion.mapped_reconstructed_data_dict[linear_obj].array,
+                        inversion_reference.mapped_reconstructed_data_dict[
+                            linear_obj
+                        ].array,
+                    )
+
+
+def _primary_beam_jit_quantities_from(dataset_sparse, mapper, image, xp):
+    """
+    The data vector, curvature matrix and `fast_chi_squared` of a sparse inversion of the visibilities
+    with the profile image `image` subtracted (via `sparse_profile_terms_from`), on a `DatasetInterface`
+    with no visibilities -- the likelihood path a `jax.jit` fit traces.
+    """
+    mask = dataset_sparse.real_space_mask
+
+    _, sparse_dirty_image, data_term = (
+        aa.util.inversion_interferometer.sparse_profile_terms_from(
+            sparse_operator=dataset_sparse.sparse_operator,
+            image=image,
+            extent_index_for_masked_pixel=mask.extent_index_for_masked_pixel,
+            xp=xp,
+        )
+    )
+
+    inversion = aa.Inversion(
+        dataset=aa.DatasetInterface(
+            data=None,
+            noise_map=dataset_sparse.noise_map,
+            grids=dataset_sparse.grids,
+            transformer=dataset_sparse.transformer,
+            sparse_operator=dataset_sparse.sparse_operator,
+            sparse_dirty_image=sparse_dirty_image,
+            data_term=data_term,
+        ),
+        linear_obj_list=[mapper],
+        xp=xp,
+    )
+
+    return inversion.data_vector, inversion.curvature_matrix, inversion.fast_chi_squared
+
+
+def test__primary_beam__none_and_ones__bit_for_bit_unchanged__jax_jit():
+    jax = pytest.importorskip("jax")
+
+    import jax.numpy as jnp
+
+    mask, _, dataset_ref, dataset_sparse_ref, _ = _primary_beam_parity_setup(
+        primary_beam=None
+    )
+
+    mapper = _mapper_from(
+        mask=mask, pixels=9, shape=(3, 3), regularization=aa.reg.Constant(1.0)
+    )
+    image = jnp.asarray(np.random.default_rng(seed=5).normal(size=mask.pixels_in_mask))
+
+    def jitted(dataset_sparse):
+        return [
+            np.asarray(value)
+            for value in jax.jit(
+                lambda i: _primary_beam_jit_quantities_from(
+                    dataset_sparse, mapper, i, xp=jnp
+                )
+            )(image)
+        ]
+
+    reference = jitted(dataset_sparse_ref)
+
+    mapping_reference = aa.Inversion(
+        dataset=dataset_ref, linear_obj_list=[mapper], xp=jnp
+    )
+
+    for primary_beam in (None, np.ones(mask.pixels_in_mask)):
+        _, _, dataset, dataset_sparse, _ = _primary_beam_parity_setup(
+            primary_beam=primary_beam
+        )
+
+        for value, value_reference in zip(jitted(dataset_sparse), reference):
+            assert np.array_equal(value, value_reference)
+
+        mapping = aa.Inversion(dataset=dataset, linear_obj_list=[mapper], xp=jnp)
+
+        for quantity in ("curvature_matrix", "data_vector"):
+            assert np.array_equal(
+                np.asarray(getattr(mapping, quantity)),
+                np.asarray(getattr(mapping_reference, quantity)),
+            )
+
+
+def test__primary_beam__jax_jit_matches_numpy_and_profile_terms_apply_P():
+    """
+    With a non-trivial primary beam the jitted sparse likelihood path matches NumPy, and the
+    profile-subtracted terms are those of the beamed profile image `P i`: the dirty image
+    `d~ - W~ P i` and data term of the un-beamed operator applied to `P i`.
+    """
+    jax = pytest.importorskip("jax")
+
+    import jax.numpy as jnp
+
+    mask, _, _, dataset_sparse, primary_beam = _primary_beam_parity_setup()
+    _, _, _, dataset_sparse_unbeamed, _ = _primary_beam_parity_setup(primary_beam=None)
+
+    mapper = _mapper_from(
+        mask=mask, pixels=9, shape=(3, 3), regularization=aa.reg.Constant(1.0)
+    )
+    image = np.random.default_rng(seed=5).normal(size=mask.pixels_in_mask)
+
+    numpy = _primary_beam_jit_quantities_from(dataset_sparse, mapper, image, xp=np)
+    jitted = jax.jit(
+        lambda i: _primary_beam_jit_quantities_from(dataset_sparse, mapper, i, xp=jnp)
+    )(jnp.asarray(image))
+
+    for value, value_numpy in zip(jitted, numpy):
+        value_numpy = np.asarray(value_numpy)
+        np.testing.assert_allclose(
+            np.asarray(value),
+            value_numpy,
+            rtol=1.0e-9,
+            atol=1.0e-9 * np.abs(value_numpy).max(),
+        )
+
+    extent_index = mask.extent_index_for_masked_pixel
+
+    terms = aa.util.inversion_interferometer.sparse_profile_terms_from(
+        sparse_operator=dataset_sparse.sparse_operator,
+        image=image,
+        extent_index_for_masked_pixel=extent_index,
+    )
+    terms_unbeamed = aa.util.inversion_interferometer.sparse_profile_terms_from(
+        sparse_operator=dataset_sparse_unbeamed.sparse_operator,
+        image=image * primary_beam.array,
+        extent_index_for_masked_pixel=extent_index,
+    )
+
+    for value, value_unbeamed in zip(terms, terms_unbeamed):
+        np.testing.assert_allclose(value, value_unbeamed, rtol=1.0e-12, atol=1.0e-12)
+
+
+def test__primary_beam__sparse_matches_mapping_and_fast_chi_squared_matches_residuals():
+    """
+    With a primary beam, the sparse inversion reproduces the dense (mapping) one -- whose transformer applies
+    `P` in `transform_mapping_matrix` -- and its `fast_chi_squared` equals the chi-squared of the residual
+    visibilities of `mapped_reconstructed_operated_data`, which the transformer forms as `F P (L s)` from the
+    un-beamed (true-sky) reconstructed image `L s`.
+    """
+    mask, rng, dataset, dataset_sparse, primary_beam = _primary_beam_parity_setup()
+
+    for name, linear_obj_list in _primary_beam_linear_obj_lists(mask, rng).items():
+        if name != "func_list_only":
+            _assert_sparse_matches_mapping(
+                dataset=dataset,
+                dataset_sparse=dataset_sparse,
+                linear_obj_list=linear_obj_list,
+            )
+
+        inversion_sparse = aa.Inversion(
+            dataset=dataset_sparse,
+            linear_obj_list=linear_obj_list,
+            settings=_PRIMARY_BEAM_SETTINGS,
+        )
+        inversion_mapping = aa.Inversion(
+            dataset=dataset,
+            linear_obj_list=linear_obj_list,
+            settings=_PRIMARY_BEAM_SETTINGS,
+        )
+
+        for quantity in ("curvature_matrix", "data_vector"):
+            reference = np.asarray(getattr(inversion_mapping, quantity))
+
+            np.testing.assert_allclose(
+                np.asarray(getattr(inversion_sparse, quantity)),
+                reference,
+                rtol=1.0e-10,
+                atol=1.0e-10 * np.abs(reference).max(),
+                err_msg=(name, quantity),
+            )
+
+        # The reconstructed image is the true sky `L s` (no P).
+        reconstruction_dict = inversion_sparse.source_quantity_dict_from(
+            source_quantity=inversion_sparse.reconstruction
+        )
+
+        for linear_obj in linear_obj_list:
+            np.testing.assert_allclose(
+                inversion_sparse.mapped_reconstructed_data_dict[linear_obj].array,
+                np.asarray(linear_obj.mapping_matrix) @ reconstruction_dict[linear_obj],
+                rtol=1.0e-12,
+                atol=1.0e-12,
+            )
+
+        residual_map = aa.util.fit.residual_map_from(
+            data=dataset.data,
+            model_data=inversion_sparse.mapped_reconstructed_operated_data,
+        )
+        chi_squared = aa.util.fit.chi_squared_complex_from(
+            chi_squared_map=aa.util.fit.chi_squared_map_complex_from(
+                residual_map=residual_map, noise_map=dataset.noise_map
+            )
+        )
+
+        # The unregularized linear function list solves to coefficients ~1e4 (no positive-only solver
+        # in the test config), so `s^T F s - 2 s^T D + d^T d` cancels to ~1e-7 relative, with or without P.
+        assert float(inversion_sparse.fast_chi_squared) == pytest.approx(
+            float(chi_squared), rel=1.0e-6
+        ), name
+
+
+def test__primary_beam__dft_and_nufft_sparse_inversions_match():
+    pytest.importorskip("nufftax")
+
+    mask, rng, _, dataset_sparse_dft, _ = _primary_beam_parity_setup()
+    _, _, _, dataset_sparse_nufft, _ = _primary_beam_parity_setup(
+        transformer_class=aa.TransformerNUFFT
+    )
+
+    assert np.array_equal(
+        dataset_sparse_nufft.transformer.primary_beam_slim,
+        dataset_sparse_dft.transformer.primary_beam_slim,
+    )
+
+    for name, linear_obj_list in _primary_beam_linear_obj_lists(mask, rng).items():
+        inversion_dft = aa.Inversion(
+            dataset=dataset_sparse_dft,
+            linear_obj_list=linear_obj_list,
+            settings=_PRIMARY_BEAM_SETTINGS,
+        )
+        inversion_nufft = aa.Inversion(
+            dataset=dataset_sparse_nufft,
+            linear_obj_list=linear_obj_list,
+            settings=_PRIMARY_BEAM_SETTINGS,
+        )
+
+        for quantity in ("curvature_matrix", "data_vector"):
+            reference = np.asarray(getattr(inversion_dft, quantity))
+
+            np.testing.assert_allclose(
+                np.asarray(getattr(inversion_nufft, quantity)),
+                reference,
+                rtol=1.0e-8,
+                atol=1.0e-8 * np.abs(reference).max(),
+                err_msg=(name, quantity),
+            )
+
+        np.testing.assert_allclose(
+            inversion_nufft.mapped_reconstructed_operated_data.array,
+            inversion_dft.mapped_reconstructed_operated_data.array,
+            rtol=1.0e-6,
+            atol=1.0e-8,
+            err_msg=name,
+        )
+
+
+def test__primary_beam__array_free_from_sparse_terms_matches_in_memory():
+    pytest.importorskip("nufftax")
+
+    mask, dataset_memory_unbeamed, dataset_stream_unbeamed, mapper = _array_free_setup()
+
+    primary_beam = aa.Array2D(
+        values=np.random.default_rng(seed=9).uniform(0.2, 1.0, size=mask.pixels_in_mask),
+        mask=mask,
+    )
+
+    dataset_memory = aa.Interferometer(
+        data=dataset_memory_unbeamed.data,
+        noise_map=dataset_memory_unbeamed.noise_map,
+        uv_wavelengths=dataset_memory_unbeamed.uv_wavelengths,
+        real_space_mask=mask,
+        transformer_class=aa.TransformerNUFFT,
+        primary_beam=primary_beam,
+    ).apply_sparse_operator()
+
+    dataset_terms = aa.Interferometer.from_sparse_terms(
+        dataset_stream_unbeamed.sparse_terms, mask, primary_beam=primary_beam
+    )
+
+    inversion_terms = aa.Inversion(dataset=dataset_terms, linear_obj_list=[mapper])
+    inversion_memory = aa.Inversion(dataset=dataset_memory, linear_obj_list=[mapper])
+
+    _assert_array_free_matches_in_memory(inversion_terms, inversion_memory, rel=1.0e-8)
+
+    # The primary beam changes the fit.
+    assert float(inversion_memory.fast_chi_squared) != pytest.approx(
+        float(
+            aa.Inversion(
+                dataset=dataset_memory_unbeamed, linear_obj_list=[mapper]
+            ).fast_chi_squared
+        ),
+        rel=1.0e-6,
+    )
+
+
+def test__primary_beam__dirty_model_image_natural_applies_P():
+    mask, _, _, dataset_sparse, primary_beam = _primary_beam_parity_setup()
+    _, _, _, dataset_sparse_unbeamed, _ = _primary_beam_parity_setup(primary_beam=None)
+
+    image = np.random.default_rng(seed=2).normal(size=mask.pixels_in_mask)
+
+    np.testing.assert_allclose(
+        aa.fit.fit_interferometer.dirty_model_image_natural_from(
+            dataset=dataset_sparse, image=image
+        ).array,
+        aa.fit.fit_interferometer.dirty_model_image_natural_from(
+            dataset=dataset_sparse_unbeamed, image=image * primary_beam.array
+        ).array,
+        rtol=1.0e-12,
+        atol=1.0e-12,
+    )
+
+
+def test__primary_beam__numba_route_is_skipped_and_direct_construction_raises():
+    pytest.importorskip("numba")
+
+    from autoarray.inversion.inversion.interferometer_numba.sparse import (
+        InversionInterferometerSparseNumba,
+    )
+
+    mask, _, _, dataset_sparse, _ = _primary_beam_parity_setup()
+    _, _, _, dataset_sparse_unbeamed, _ = _primary_beam_parity_setup(primary_beam=None)
+
+    mapper = _mapper_from(
+        mask=mask, pixels=9, shape=(3, 3), regularization=aa.reg.Constant(1.0)
+    )
+    settings = aa.Settings(interferometer_route="numba")
+
+    assert isinstance(
+        aa.Inversion(
+            dataset=dataset_sparse_unbeamed, linear_obj_list=[mapper], settings=settings
+        ),
+        InversionInterferometerSparseNumba,
+    )
+    assert (
+        type(
+            aa.Inversion(
+                dataset=dataset_sparse, linear_obj_list=[mapper], settings=settings
+            )
+        )
+        is aa.InversionInterferometerSparse
+    )
+
+    with pytest.raises(aa.exc.InversionException, match="primary beam"):
+        InversionInterferometerSparseNumba(
+            dataset=dataset_sparse, linear_obj_list=[mapper], settings=settings
+        )
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 cd3537b8f..0f2673c34 100644
--- a/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py
+++ b/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py
@@ -2486,3 +2486,69 @@ def test__interferometer__from_stream_and_apply_sparse_operator_from_chunks__pas
             rtol=1.0e-12,
             atol=1.0e-12 * np.abs(reference.dirty_image_fine).max(),
         )
+
+
+def test__sparse_triplets_from__slim_weights_scale_vals_by_their_slim_row():
+    """
+    `slim_weights` (the interferometer primary beam) multiplies each triplet's value by the weight of its
+    slim image pixel, giving the triplets of `diag(w) A`; `None` and ones leave the triplets bit-for-bit
+    unchanged. Checked on both the NumPy and (when available) JAX branches.
+    """
+    mask = aa.Mask2D.circular(shape_native=(10, 10), pixel_scales=1.0, radius=3.5)
+
+    grid = aa.Grid2D.from_mask(mask=mask, over_sample_size=2)
+
+    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)
+
+    weights = np.random.default_rng(seed=3).uniform(0.2, 1.0, size=mask.pixels_in_mask)
+
+    xps = [np]
+
+    try:
+        import jax.numpy as jnp
+
+        xps.append(jnp)
+    except ModuleNotFoundError:
+        pass
+
+    for xp in xps:
+
+        def triplets_from(slim_weights, return_rows_slim):
+            return [
+                np.asarray(array)
+                for array in 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=return_rows_slim,
+                    slim_weights=slim_weights,
+                    xp=xp,
+                )
+            ]
+
+        rows_slim, cols, vals = triplets_from(None, return_rows_slim=True)
+
+        _, cols_weighted, vals_weighted = triplets_from(weights, return_rows_slim=True)
+
+        assert np.array_equal(cols_weighted, cols)
+        np.testing.assert_allclose(
+            vals_weighted, vals * weights[rows_slim], rtol=1.0e-15, atol=0.0
+        )
+
+        for return_rows_slim in (True, False):
+            reference = triplets_from(None, return_rows_slim=return_rows_slim)
+            ones = triplets_from(
+                np.ones(mask.pixels_in_mask), return_rows_slim=return_rows_slim
+            )
+
+            for array, array_reference in zip(ones, reference):
+                assert np.array_equal(array, array_reference)
diff --git a/test_autoarray/operators/test_transformer.py b/test_autoarray/operators/test_transformer.py
index 5567edfab..67ea899e8 100644
--- a/test_autoarray/operators/test_transformer.py
+++ b/test_autoarray/operators/test_transformer.py
@@ -368,3 +368,212 @@ def test__nufftax_exception__keeps_the_plain_install_instruction():
     message = str(exc.value)
     assert "Install it via the command `pip install nufftax`" in message
     assert "Do NOT" not in message
+
+
+def _primary_beam_setup(seed=4):
+    """
+    A non-square, partially masked image grid with random uv-coverage, a random positive image and
+    mapping matrix, and a non-trivial primary beam with values spread over (0.2, 1], for the
+    primary-beam transformer tests below.
+    """
+    rng = np.random.default_rng(seed=seed)
+
+    mask = aa.Mask2D.circular(shape_native=(8, 9), pixel_scales=0.2, radius=0.7)
+    uv_wavelengths = rng.normal(size=(23, 2)) * 2.0e5
+
+    image = aa.Array2D(values=rng.normal(size=mask.pixels_in_mask), mask=mask)
+    mapping_matrix = rng.normal(size=(mask.pixels_in_mask, 4))
+    primary_beam = aa.Array2D(
+        values=rng.uniform(0.2, 1.0, size=mask.pixels_in_mask), mask=mask
+    )
+
+    return mask, uv_wavelengths, image, mapping_matrix, primary_beam
+
+
+def _transformer_classes():
+    classes = [aa.TransformerDFT]
+
+    try:
+        import nufftax  # noqa: F401
+
+        classes.append(aa.TransformerNUFFT)
+    except ModuleNotFoundError:
+        pass
+
+    return classes
+
+
+def test__primary_beam__none_and_ones__bit_for_bit_unchanged():
+    """
+    `primary_beam=None` takes exactly the code path of a transformer built without one, and
+    `primary_beam=ones` multiplies every pixel by exactly 1.0, so both leave the forward transforms and
+    the adjoint bit-for-bit unchanged.
+    """
+    mask, uv_wavelengths, image, mapping_matrix, _ = _primary_beam_setup()
+
+    ones = aa.Array2D.ones(shape_native=mask.shape_native, pixel_scales=0.2).native
+
+    for transformer_class in _transformer_classes():
+        reference = transformer_class(
+            uv_wavelengths=uv_wavelengths, real_space_mask=mask
+        )
+
+        visibilities = reference.visibilities_from(image=image)
+
+        for primary_beam in (None, ones):
+            transformer = transformer_class(
+                uv_wavelengths=uv_wavelengths,
+                real_space_mask=mask,
+                primary_beam=primary_beam,
+            )
+
+            assert np.array_equal(
+                transformer.visibilities_from(image=image).array,
+                visibilities.array,
+            )
+            assert np.array_equal(
+                transformer.transform_mapping_matrix(mapping_matrix=mapping_matrix),
+                reference.transform_mapping_matrix(mapping_matrix=mapping_matrix),
+            )
+            assert np.array_equal(
+                transformer.image_from(visibilities=visibilities).array,
+                reference.image_from(visibilities=visibilities).array,
+            )
+
+
+def test__primary_beam__forward_transforms_equal_unbeamed_transforms_of_beamed_input():
+    """
+    With a primary beam `P`, `visibilities_from(i)` is the un-beamed `visibilities_from(P i)` and
+    `transform_mapping_matrix(M)` is the un-beamed `transform_mapping_matrix(diag(P) M)`, while
+    `image_from` (the adjoint, used for the dirty image and the w-tilde terms) is unchanged.
+    """
+    mask, uv_wavelengths, image, mapping_matrix, primary_beam = _primary_beam_setup()
+
+    beamed_image = aa.Array2D(values=image.array * primary_beam.array, mask=mask)
+
+    for transformer_class in _transformer_classes():
+        unbeamed = transformer_class(uv_wavelengths=uv_wavelengths, real_space_mask=mask)
+        beamed = unbeamed.with_primary_beam(primary_beam)
+
+        assert unbeamed.primary_beam_slim is None
+        assert beamed.primary_beam_slim == pytest.approx(primary_beam.array, 1.0e-14)
+
+        np.testing.assert_allclose(
+            beamed.visibilities_from(image=image).array,
+            unbeamed.visibilities_from(image=beamed_image).array,
+            rtol=1.0e-12,
+            atol=1.0e-12,
+        )
+        np.testing.assert_allclose(
+            beamed.transform_mapping_matrix(mapping_matrix=mapping_matrix),
+            unbeamed.transform_mapping_matrix(
+                mapping_matrix=mapping_matrix * primary_beam.array[:, None]
+            ),
+            rtol=1.0e-12,
+            atol=1.0e-12,
+        )
+
+        visibilities = unbeamed.visibilities_from(image=image)
+
+        assert np.array_equal(
+            beamed.image_from(visibilities=visibilities).array,
+            unbeamed.image_from(visibilities=visibilities).array,
+        )
+
+        # The constructor kwarg and `with_primary_beam` agree, and a native-shape beam is accepted.
+        constructed = transformer_class(
+            uv_wavelengths=uv_wavelengths,
+            real_space_mask=mask,
+            primary_beam=primary_beam.native.array,
+        )
+
+        assert np.array_equal(constructed.primary_beam_slim, beamed.primary_beam_slim)
+
+
+def test__primary_beam__with_primary_beam__shallow_copy_leaves_original_untouched():
+    mask, uv_wavelengths, _, _, primary_beam = _primary_beam_setup()
+
+    transformer = aa.TransformerDFT(uv_wavelengths=uv_wavelengths, real_space_mask=mask)
+
+    assert transformer.with_primary_beam(None) is transformer
+
+    beamed = transformer.with_primary_beam(primary_beam)
+
+    assert beamed is not transformer
+    assert transformer.primary_beam_slim is None
+    assert beamed.grid is transformer.grid
+
+    assert beamed.with_primary_beam(None).primary_beam_slim is None
+    assert beamed.primary_beam_slim is not None
+
+
+def test__primary_beam__wrong_shape_raises():
+    mask, uv_wavelengths, _, _, _ = _primary_beam_setup()
+
+    with pytest.raises(aa.exc.DatasetException):
+        aa.TransformerDFT(
+            uv_wavelengths=uv_wavelengths,
+            real_space_mask=mask,
+            primary_beam=np.ones((5, 5)),
+        )
+
+
+def test__primary_beam__jax_jit_matches_numpy():
+    """
+    The beamed forward transforms trace under `jax.jit` and match NumPy, and with `P = None` / `P = ones`
+    the jitted outputs are bit-for-bit those of the transformer without a primary beam.
+    """
+    jax = pytest.importorskip("jax")
+    import jax.numpy as jnp
+
+    mask, uv_wavelengths, image, mapping_matrix, primary_beam = _primary_beam_setup()
+
+    ones = np.ones(mask.pixels_in_mask)
+
+    for transformer_class in _transformer_classes():
+        unbeamed = transformer_class(uv_wavelengths=uv_wavelengths, real_space_mask=mask)
+
+        def jitted(transformer):
+            @jax.jit
+            def visibilities_from(image_slim):
+                return transformer.visibilities_from(
+                    image=aa.Array2D(values=image_slim, mask=mask), xp=jnp
+                ).array
+
+            @jax.jit
+            def transform_mapping_matrix(mapping_matrix):
+                return transformer.transform_mapping_matrix(
+                    mapping_matrix=mapping_matrix, xp=jnp
+                )
+
+            return (
+                np.asarray(visibilities_from(jnp.asarray(image.array))),
+                np.asarray(transform_mapping_matrix(jnp.asarray(mapping_matrix))),
+            )
+
+        visibilities_reference, mapping_reference = jitted(unbeamed)
+
+        for primary_beam_unit in (None, ones):
+            visibilities, transformed = jitted(
+                unbeamed.with_primary_beam(primary_beam_unit)
+            )
+
+            assert np.array_equal(visibilities, visibilities_reference)
+            assert np.array_equal(transformed, mapping_reference)
+
+        beamed = unbeamed.with_primary_beam(primary_beam)
+
+        visibilities, transformed = jitted(beamed)
+
+        np.testing.assert_allclose(
+            visibilities,
+            beamed.visibilities_from(image=image).array,
+            rtol=1.0e-10,
+            atol=1.0e-10,
+        )
+        np.testing.assert_allclose(
+            transformed,
+            beamed.transform_mapping_matrix(mapping_matrix=mapping_matrix),
+            rtol=1.0e-10,
+            atol=1.0e-10,
+        )