feat: persist the interferometer primary beam through analysis save and aggregator rebuild

autogalaxy/aggregator/interferometer/interferometer.py

@@ -2,8 +2,8 @@
 Aggregator interface for loading ``Interferometer`` datasets from model-fit results.
 
 After an interferometer model-fit the dataset (complex visibilities, noise-map,
-uv-wavelengths, real-space mask, and transformer class) is written to a FITS file in
-the output directory or SQLite database.  This module reconstructs a fully-configured
+uv-wavelengths, real-space mask, transformer class and, if used, primary beam) is
+written to a FITS file in the output directory or SQLite database.  This module reconstructs a fully-configured
 ``aa.Interferometer`` object from those stored files.
 
 Two public objects are provided:
@@ -18,7 +18,7 @@
 
 from __future__ import annotations
 from functools import partial
-from typing import List
+from typing import List, Optional
 
 import numpy as np
 
@@ -59,6 +59,22 @@ def _has_hdu(hdu_list, name: str) -> bool:
     return True
 
 
+def _primary_beam_from(hdu_list) -> Optional[np.ndarray]:
+    """
+    Returns the primary beam stored in the `PRIMARY_BEAM` HDU of a `dataset.fits` HDU list written by
+    `autogalaxy.interferometer.model.analysis.interferometer_hdu_list_from`, as a float64 array of the native
+    shape of the real-space mask (masked pixels 0), or `None` if the file has no such HDU (the dataset had no
+    primary beam, or the file was written before the beam was persisted).
+
+    The HDU is detected by name only (`_has_hdu`), never via `_hdu`, whose fallback to the legacy positional
+    indices would misread another HDU as the beam.
+    """
+    if not _has_hdu(hdu_list, "primary_beam"):
+        return None
+
+    return np.asarray(hdu_list["PRIMARY_BEAM"].data, dtype="float64")
+
+
 def _sparse_terms_from(hdu_list, real_space_mask: aa.Mask2D) -> aa.SparseTerms:
     """
     Returns the `SparseTerms` of an array-free interferometer dataset from the `dataset.fits` HDU list written
@@ -178,6 +194,11 @@ def _interferometer_from(
     present). It is rebuilt via `Interferometer.from_sparse_terms`, so the returned dataset is array-free with
     its sparse operator re-attached, and a fit of it reproduces the original `log_evidence`.
 
+    If the dataset had a primary beam, `dataset.fits` has a trailing `PRIMARY_BEAM` HDU, which is passed back
+    as `primary_beam=` to either constructor, so the reloaded dataset (its transformer or sparse operator) and
+    any fit of it carry the same beam and reproduce the beamed likelihood. A file without that HDU (no beam, or
+    written before the beam was persisted) reloads with `primary_beam=None`.
+
     Each individual attribute can be loaded from the database via the `fit.value()` method.
 
     This method combines all of these attributes and returns a `Interferometer` object, including having its
@@ -203,11 +224,14 @@ def _interferometer_from(
 
         hdu_list = fit.value(name="dataset")
 
+        primary_beam = _primary_beam_from(hdu_list)
+
         if _has_hdu(hdu_list, "nufft_precision_operator"):
             dataset_list.append(
                 aa.Interferometer.from_sparse_terms(
                     _sparse_terms_from(hdu_list, real_space_mask=real_space_mask),
                     real_space_mask=real_space_mask,
+                    primary_beam=primary_beam,
                 )
             )
             continue
@@ -228,6 +252,7 @@ def _interferometer_from(
             uv_wavelengths=uv_wavelengths,
             real_space_mask=real_space_mask,
             transformer_class=transformer_class,
+            primary_beam=primary_beam,
         )
 
         dataset_list.append(dataset)
@@ -255,6 +280,9 @@ def __init__(self, aggregator: af.Aggregator):
     present). It is rebuilt via `Interferometer.from_sparse_terms`, so the returned dataset is array-free with
     its sparse operator re-attached, and a fit of it reproduces the original `log_evidence`.
 
+    A primary beam, if the dataset had one, is reloaded from the trailing `PRIMARY_BEAM` HDU and passed back
+    as `primary_beam=`, so the reloaded dataset carries the same beam.
+
         The `aggregator` contains the path to each of these files, and they can be loaded individually. This class
         can load them all at once and create an `Interferometer` object via the `_interferometer_from` method.
 

autogalaxy/interferometer/model/analysis.py

@@ -73,6 +73,23 @@
 )
 
 
+def _primary_beam_native_from(dataset: aa.Interferometer) -> Optional[np.ndarray]:
+    """
+    Returns the primary beam of `dataset` as a float64 array of the native shape of its real-space mask
+    (masked pixels 0), as written to the `PRIMARY_BEAM` HDU of `dataset.fits`, or `None` if the dataset
+    has no primary beam.
+    """
+    if dataset.primary_beam_slim is None:
+        return None
+
+    return np.asarray(
+        aa.Array2D(
+            values=dataset.primary_beam_slim, mask=dataset.real_space_mask
+        ).native.array,
+        dtype="float64",
+    )
+
+
 def interferometer_hdu_list_from(
     dataset: aa.Interferometer, include_fine_grids: bool = False
 ):
@@ -105,6 +122,13 @@ def interferometer_hdu_list_from(
     grids; rebuild them from the visibilities, or write this HDU list yourself with
     `include_fine_grids=True`, when they are needed.
 
+    A dataset with a primary beam (`Interferometer(primary_beam=...)`) also writes it, in both layouts, as a
+    trailing HDU with EXTNAME `PRIMARY_BEAM`: a float64 array of the native shape of the real-space mask with
+    masked pixels 0. It is the last HDU, so the positional indices of the legacy in-memory layout (0-3) are
+    unchanged, and a dataset without a primary beam writes no such HDU (the file is identical to one written
+    before the beam was persisted). The aggregator detects it by name and passes it back as `primary_beam=`,
+    so a fit of the reloaded dataset reproduces the beamed likelihood.
+
     Parameters
     ----------
     dataset
@@ -114,15 +138,31 @@ def interferometer_hdu_list_from(
     """
     mask = dataset.real_space_mask
 
+    primary_beam_native = _primary_beam_native_from(dataset)
+
+    if primary_beam_native is None:
+        primary_beam_values = []
+        primary_beam_ext_names = []
+    else:
+        primary_beam_values = [primary_beam_native]
+        primary_beam_ext_names = ["primary_beam"]
+
     if not dataset.is_array_free:
         return hdu_list_for_output_from(
             values_list=[
                 mask.astype("float"),
                 dataset.data.in_array,
                 dataset.noise_map.in_array,
                 dataset.uv_wavelengths,
+                *primary_beam_values,
+            ],
+            ext_name_list=[
+                "mask",
+                "data",
+                "noise_map",
+                "uv_wavelengths",
+                *primary_beam_ext_names,
             ],
-            ext_name_list=["mask", "data", "noise_map", "uv_wavelengths"],
             header_dict=mask.header_dict,
         )
 
@@ -205,6 +245,7 @@ def interferometer_hdu_list_from(
                 dtype="float64",
             ),
             *fine_values,
+            *primary_beam_values,
         ],
         ext_name_list=[
             "mask",
@@ -213,6 +254,7 @@ def interferometer_hdu_list_from(
             "dirty_beam",
             "sparse_terms_scalars",
             *fine_ext_names,
+            *primary_beam_ext_names,
         ],
         header_dict=header_dict,
     )
@@ -418,6 +460,8 @@ def save_attributes(self, paths: af.DirectoryPaths):
         - The real space mask applied to the dataset, in the `PrimaryHDU`.
         - The interferometer dataset (data / noise-map / uv_wavelengths), or for an array-free dataset
           its `SparseTerms` (see `interferometer_hdu_list_from`).
+        - The primary beam on the real-space mask (native, masked pixels 0), as a trailing `PRIMARY_BEAM`
+          HDU, if the dataset has one, so the aggregator-rebuilt dataset and fit carry it.
 
         It is common for these attributes to be loaded by many of the template aggregator functions given in the
         `aggregator` modules. For example, when using the database tools to perform a fit, the default behaviour is for