feat: Preloads API for reusing channel-invariant inversion quantities

Add an opt-in `Preloads` mechanism: precompute an invariant inversion
quantity once and inject it so repeated evaluations reuse it instead of
rebuilding it.

- New `aa.AbstractPreloads` / `aa.PreloadsInterferometer` containers
  (autoarray/preloads/) with an optional `curvature_matrix` field.
- The inversion stack (`AbstractInversion`, the interferometer inversions,
  and the `inversion_from` / `inversion_interferometer_from` factory) gains
  an optional `preloads=None` argument, threaded as a single object.
- The sparse interferometer inversion returns `preloads.curvature_matrix`
  (`F = LᵀW̃L`) directly when provided, skipping the dominant rebuild.

Default behaviour is unchanged (`preloads=None`). The primary use is a
shared / combined likelihood (e.g. a datacube FactorGraphModel) where
invariance is explicit; an earlier preload system was removed for being
bug-prone, so the docstrings make the opt-in, caller-owns-invariance
contract explicit.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: 2 people

autoarray/__init__.py

@@ -28,6 +28,8 @@
 from .inversion import regularization as reg
 from .settings import Settings
 from .inversion.inversion.abstract import AbstractInversion
+from .preloads import AbstractPreloads
+from .preloads import PreloadsInterferometer
 from .inversion.regularization.abstract import AbstractRegularization
 from .inversion.inversion.factory import inversion_from as Inversion
 from .inversion.inversion.dataset_interface import DatasetInterface

autoarray/inversion/inversion/abstract.py

@@ -27,6 +27,7 @@ def __init__(
         linear_obj_list: List[LinearObj],
         settings: Settings = None,
         xp=np,
+        preloads=None,
     ):
         """
         An `Inversion` reconstructs an input dataset using a list of linear objects (e.g. a list of analytic functions
@@ -65,6 +66,14 @@ def __init__(
             Settings controlling how an inversion is fitted, for example which linear algebra formalism is used.
         xp
             The array module to use (`numpy` by default; pass `jax.numpy` for JAX support).
+        preloads
+            An optional `AbstractPreloads` (e.g. a `PreloadsInterferometer`) carrying pre-computed
+            inversion quantities (e.g. the `curvature_matrix` `F`). When a quantity is invariant
+            across the evaluations reusing this inversion — a fixed pixelization across a search, or
+            the channel-invariant quantities of a datacube `FactorGraphModel` — it is computed once
+            and preloaded here so the inversion reuses it instead of rebuilding the dominant
+            inversion-setup cost. `None` (the default) leaves the standard behaviour unchanged, as
+            does any individual preload field left `None`.
         """
 
         self.dataset = dataset
@@ -75,6 +84,8 @@ def __init__(
 
         self.use_jax = xp is not np
 
+        self._preloads = preloads
+
     @property
     def _xp(self):
         if self.use_jax:

autoarray/inversion/inversion/factory.py

@@ -32,6 +32,7 @@ def inversion_from(
     linear_obj_list: List[LinearObj],
     settings: Settings = None,
     xp=np,
+    preloads=None,
 ):
     """
     Factory which given an input dataset and list of linear objects, creates an `Inversion`.
@@ -70,7 +71,11 @@ def inversion_from(
         )
 
     return inversion_interferometer_from(
-        dataset=dataset, linear_obj_list=linear_obj_list, settings=settings, xp=xp
+        dataset=dataset,
+        linear_obj_list=linear_obj_list,
+        settings=settings,
+        xp=xp,
+        preloads=preloads,
     )
 
 
@@ -152,6 +157,7 @@ def inversion_interferometer_from(
     linear_obj_list: List[LinearObj],
     settings: Settings = None,
     xp=np,
+    preloads=None,
 ):
     """
     Factory which given an input `Interferometer` dataset and list of linear objects, creates
@@ -201,6 +207,7 @@ def inversion_interferometer_from(
             linear_obj_list=linear_obj_list,
             settings=settings,
             xp=xp,
+            preloads=preloads,
         )
 
     return InversionInterferometerMapping(

autoarray/inversion/inversion/interferometer/abstract.py

@@ -19,6 +19,7 @@ def __init__(
         linear_obj_list: List[LinearObj],
         settings: Settings = None,
         xp=np,
+        preloads=None,
     ):
         """
         Constructs linear equations (via vectors and matrices) which allow for sets of simultaneous linear equations
@@ -44,7 +45,11 @@ def __init__(
         """
 
         super().__init__(
-            dataset=dataset, linear_obj_list=linear_obj_list, settings=settings, xp=xp
+            dataset=dataset,
+            linear_obj_list=linear_obj_list,
+            settings=settings,
+            xp=xp,
+            preloads=preloads,
         )
 
     @property

autoarray/inversion/inversion/interferometer/sparse.py

@@ -20,6 +20,7 @@ def __init__(
         linear_obj_list: List[LinearObj],
         settings: Settings = None,
         xp=np,
+        preloads=None,
     ):
         """
         Constructs linear equations (via vectors and matrices) which allow for sets of simultaneous linear equations
@@ -49,6 +50,7 @@ def __init__(
             linear_obj_list=linear_obj_list,
             settings=settings,
             xp=xp,
+            preloads=preloads,
         )
 
         self.settings = settings
@@ -83,7 +85,16 @@ def curvature_matrix(self) -> np.ndarray:
         If there are multiple linear objects their `operated_mapping_matrix` properties will have already been
         concatenated ensuring their `curvature_matrix` values are solved for simultaneously. This includes all
         diagonal and off-diagonal terms describing the covariances between linear objects.
+
+        If a `preloads.curvature_matrix` was injected (e.g. the datacube shared-state path, where `F` is
+        identical for every channel) it is returned directly, so the dominant `F = LᵀW̃L` build is skipped.
+        This is the only invariant inversion-setup quantity that must be preloaded explicitly: the mapper
+        (and therefore `mapping_matrix` and `regularization_matrix`) is already reused by passing the same
+        `linear_obj_list` to every channel's inversion.
         """
+        if self._preloads is not None and self._preloads.curvature_matrix is not None:
+            return self._preloads.curvature_matrix
+
         return self.curvature_matrix_diag
 
     @property

autoarray/preloads/__init__.py

@@ -0,0 +1,2 @@
+from .abstract import AbstractPreloads
+from .interferometer import PreloadsInterferometer

autoarray/preloads/abstract.py

@@ -0,0 +1,40 @@
+class AbstractPreloads:
+    def __init__(self, curvature_matrix=None):
+        """
+        Container for quantities that are *preloaded* into a fit / inversion: computed once and
+        injected so that repeated evaluations reuse them instead of recomputing them.
+
+        Preloading is an **advanced, opt-in** optimisation. A quantity may be preloaded only when it
+        is genuinely invariant across the evaluations that reuse it; preloading a quantity that in
+        fact changes silently corrupts the result. Establishing that invariance is the caller's
+        responsibility — and the difficulty of knowing *when* a quantity is safe to preload is why
+        preloads are **not applied as standard** (an earlier preload system was removed because this
+        was bug-prone and hard to maintain).
+
+        The natural, safe home for preloads is a **shared / combined likelihood**, where invariance
+        is explicit and easy to verify:
+
+        - In a datacube `FactorGraphModel` the lens model is identical for every spectral channel,
+          so the channel-invariant inversion quantities are computed once (by
+          `AnalysisInterferometer.shared_state_from`) and reused by every channel. It is obvious
+          here exactly which quantities are invariant, which is what makes preloading safe.
+
+        More general single-fit preloading (e.g. a source pixelization held fixed across a search)
+        is possible but much harder to reason about, and is deliberately not the default.
+
+        Every field is optional and defaults to `None`, so a caller preloads only the quantities it
+        knows are invariant; each consumer reuses a field only when it is not `None`, otherwise
+        falling back to the standard computation. New preloadable quantities are added as fields
+        without changing the call signature threaded through the fit / inversion stack.
+
+        This abstract base holds the quantities common to every dataset type. Dataset-specific
+        subclasses (`PreloadsInterferometer`, and `PreloadsImaging` when needed) add the preloads
+        unique to their formalism, mirroring the `AnalysisImaging` / `AnalysisInterferometer` split.
+
+        Parameters
+        ----------
+        curvature_matrix
+            The pre-computed curvature matrix `F = LᵀW̃L` — the dominant inversion-setup cost. When
+            provided, the inversion returns it directly instead of rebuilding it.
+        """
+        self.curvature_matrix = curvature_matrix

autoarray/preloads/interferometer.py

@@ -0,0 +1,24 @@
+from autoarray.preloads.abstract import AbstractPreloads
+
+
+class PreloadsInterferometer(AbstractPreloads):
+    def __init__(self, curvature_matrix=None):
+        """
+        Preloaded quantities for an interferometer fit / inversion (see `AbstractPreloads`).
+
+        This is the consumer-facing preloads container for interferometer data — for example the
+        object returned by `AnalysisInterferometer.shared_state_from` for the datacube shared-state
+        path, where the channel-invariant inversion quantities are preloaded once and reused by
+        every spectral channel.
+
+        It currently carries only the universal `curvature_matrix` (`F`) from `AbstractPreloads`,
+        but it is the home for interferometer-specific preloads as they are added (e.g. w-tilde /
+        NUFFT translation-invariant quantities), and it may also carry non-inversion preloads used
+        higher up at the fit level.
+
+        Parameters
+        ----------
+        curvature_matrix
+            The pre-computed curvature matrix `F = LᵀW̃L`. See `AbstractPreloads`.
+        """
+        super().__init__(curvature_matrix=curvature_matrix)

test_autoarray/inversion/inversion/interferometer/test_interferometer.py

@@ -197,3 +197,87 @@ def test__curvature_matrix__interferometer_sparse_operator__delaunay__dft_and_nu
     assert inversion_nufft.data_vector == pytest.approx(
         inversion_dft.data_vector, 1.0e-4
     )
+
+
+def test__preloads_interferometer__curvature_matrix_returned_directly_and_skips_rebuild():
+    """
+    A `PreloadsInterferometer` injects a pre-computed `curvature_matrix` (`F`) — e.g. the datacube
+    shared-state path where `F` is identical across channels. The sparse interferometer inversion
+    must return it verbatim (skipping the dominant `F = LᵀW̃L` build) while leaving the per-channel
+    `data_vector` unchanged.
+    """
+    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,
+    )
+
+    grid = aa.Grid2D.from_mask(mask=mask, over_sample_size=1)
+
+    mesh = aa.mesh.Delaunay(pixels=9)
+    image_mesh = aa.image_mesh.Overlay(shape=(3, 3))
+    image_mesh_grid = image_mesh.image_plane_mesh_grid_from(mask=mask, adapt_data=None)
+
+    interpolator = mesh.interpolator_from(
+        source_plane_data_grid=grid,
+        source_plane_mesh_grid=image_mesh_grid,
+    )
+    mapper = aa.Mapper(interpolator=interpolator)
+
+    n_visibilities = 5
+    rng = np.random.default_rng(seed=0)
+    data = aa.Visibilities(
+        visibilities=rng.normal(size=(n_visibilities, 2)).astype(np.float64)
+    )
+    noise_map = aa.VisibilitiesNoiseMap(
+        visibilities=np.ones((n_visibilities, 2), dtype=np.float64)
+    )
+    uv_wavelengths = rng.normal(size=(n_visibilities, 2)).astype(np.float64)
+
+    dataset = aa.Interferometer(
+        data=data,
+        noise_map=noise_map,
+        uv_wavelengths=uv_wavelengths,
+        real_space_mask=mask,
+        transformer_class=aa.TransformerDFT,
+    )
+
+    dataset_sparse = dataset.apply_sparse_operator(use_jax=False)
+
+    inversion = aa.Inversion(dataset=dataset_sparse, linear_obj_list=[mapper])
+    curvature_matrix = inversion.curvature_matrix
+
+    # A sentinel injected via the shared `curvature_matrix` is returned verbatim, proving the
+    # expensive `curvature_matrix_diag` build is skipped (else the real F would be returned).
+    sentinel = np.full_like(curvature_matrix, 7.0)
+    inversion_sentinel = aa.Inversion(
+        dataset=dataset_sparse,
+        linear_obj_list=[mapper],
+        preloads=aa.PreloadsInterferometer(curvature_matrix=sentinel),
+    )
+    assert inversion_sentinel.curvature_matrix is sentinel
+
+    # An empty `PreloadsInterferometer` (curvature_matrix left None) falls back to the standard build.
+    inversion_empty_preloads = aa.Inversion(
+        dataset=dataset_sparse,
+        linear_obj_list=[mapper],
+        preloads=aa.PreloadsInterferometer(),
+    )
+    assert inversion_empty_preloads.curvature_matrix == pytest.approx(curvature_matrix)
+
+    # Preloading the real F reproduces the un-preloaded curvature matrix and leaves the per-channel
+    # data_vector (which does not depend on the preloaded F) unchanged.
+    inversion_preloaded = aa.Inversion(
+        dataset=dataset_sparse,
+        linear_obj_list=[mapper],
+        preloads=aa.PreloadsInterferometer(curvature_matrix=curvature_matrix),
+    )
+    assert inversion_preloaded.curvature_matrix is curvature_matrix
+    assert inversion_preloaded.data_vector == pytest.approx(inversion.data_vector)