feat: primary beam for interferometer datasets, transformers and the sparse inversion

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

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,13 +289,21 @@ 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,
             noise_map=noise_map,
             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(

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, ...)`)."
+    )