Merge pull request #600 from PyAutoLabs/feature/multi-shared-state-core-api

feat: imaging shared-state consumer — AnalysisImaging.shared_state_from + mesh preload reuse (phase 2/4)

Author: Jammy2211

autolens/imaging/fit_imaging.py

@@ -44,7 +44,8 @@ def __init__(
         dataset_model : Optional[aa.DatasetModel] = None,
         adapt_images: Optional[ag.AdaptImages] = None,
         settings: aa.Settings = None,
-        xp=np
+        xp=np,
+        preloads=None,
     ):
         """
         Fits an imaging dataset using a `Tracer` object.
@@ -83,6 +84,11 @@ def __init__(
             reconstructed galaxy's morphology.
         settings
             Settings controlling how an inversion is fitted for example which linear algebra formalism is used.
+        preloads
+            An optional `PreloadsImaging` carrying exposure-invariant quantities (the shared
+            source-plane mesh geometry) computed once and reused by the fit instead of being
+            rebuilt. Supplied by the multi-exposure shared-state path (see
+            `AnalysisImaging.shared_state_from`); `None` (the default) fits as normal.
         """
 
         super().__init__(dataset=dataset, dataset_model=dataset_model, xp=xp)
@@ -94,6 +100,7 @@ def __init__(
 
         self.adapt_images = adapt_images
         self.settings = settings or aa.Settings()
+        self.preloads = preloads
 
     @functools.cached_property
     def blurred_image(self) -> aa.Array2D:
@@ -140,6 +147,7 @@ def tracer_to_inversion(self) -> TracerToInversion:
             adapt_images=self.adapt_images,
             settings=self.settings,
             xp=self._xp,
+            preloads=self.preloads,
         )
 
     @cached_property

autolens/imaging/model/analysis.py

@@ -14,6 +14,7 @@
 """
 import logging
 
+import autoarray as aa
 import autofit as af
 import autogalaxy as ag
 
@@ -39,7 +40,48 @@ class AnalysisImaging(AnalysisDataset):
     Visualizer = VisualizerImaging
     Latent = LatentLens
 
-    def log_likelihood_function(self, instance: af.ModelInstance) -> float:
+    def __init__(
+        self,
+        dataset,
+        positions_likelihood_list=None,
+        adapt_images: ag.AdaptImages = None,
+        cosmology: ag.cosmo.LensingCosmology = None,
+        settings=None,
+        raise_inversion_positions_likelihood_exception: bool = True,
+        title_prefix: str = None,
+        use_jax: bool = True,
+        shared_preloads: bool = False,
+        **kwargs,
+    ):
+        """
+        Fits a lens model to an imaging dataset via a non-linear search (see `AnalysisDataset` for
+        the full docstring of the shared parameters).
+
+        Parameters
+        ----------
+        shared_preloads
+            Opts this analysis into the cross-factor shared-state mechanism of a `FactorGraphModel`
+            (see `shared_state_from`). Set this to `True` only when this analysis is one of many
+            exposures of the same lens (e.g. multi-exposure imaging with per-exposure pixel offsets)
+            sharing an identical lens model, so the exposure-invariant source-plane mesh geometry
+            can be computed once and reused by every exposure. `False` by default, leaving the
+            standard per-analysis behaviour unchanged.
+        """
+        super().__init__(
+            dataset=dataset,
+            positions_likelihood_list=positions_likelihood_list,
+            adapt_images=adapt_images,
+            cosmology=cosmology,
+            settings=settings,
+            raise_inversion_positions_likelihood_exception=raise_inversion_positions_likelihood_exception,
+            title_prefix=title_prefix,
+            use_jax=use_jax,
+            **kwargs,
+        )
+
+        self.shared_preloads = shared_preloads
+
+    def log_likelihood_function(self, instance: af.ModelInstance, shared=None) -> float:
         """
         Given an instance of the model, where the model parameters are set via a non-linear search, fit the model
         instance to the imaging dataset.
@@ -71,6 +113,11 @@ def log_likelihood_function(self, instance: af.ModelInstance) -> float:
         instance
             An instance of the model that is being fitted to the data by this analysis (whose parameters have been set
             via a non-linear search).
+        shared
+            The cross-factor shared state of a `FactorGraphModel`, computed once per evaluation by the lead
+            factor's `shared_state_from` (see that method). For this analysis it is a `PreloadsImaging`
+            carrying the exposure-invariant source-plane mesh geometry; when provided it is reused by the fit
+            instead of being recomputed. `None` (the default, e.g. a standalone fit) leaves behaviour unchanged.
 
         Returns
         -------
@@ -83,16 +130,63 @@ def log_likelihood_function(self, instance: af.ModelInstance) -> float:
         )
 
         if self._use_jax:
-            return self.fit_from(instance=instance).figure_of_merit - log_likelihood_penalty
+            return (
+                self.fit_from(instance=instance, preloads=shared).figure_of_merit
+                - log_likelihood_penalty
+            )
 
         try:
-            return self.fit_from(instance=instance).figure_of_merit - log_likelihood_penalty
+            return (
+                self.fit_from(instance=instance, preloads=shared).figure_of_merit
+                - log_likelihood_penalty
+            )
         except Exception as e:
             raise af.exc.FitException
 
+    def shared_state_from(self, instance: af.ModelInstance):
+        """
+        Compute the exposure-invariant source-plane mesh geometry once so it can be shared across the factors
+        of a multi-exposure `FactorGraphModel` (see `autofit.Analysis.shared_state_from`).
+
+        When `shared_preloads` is set, every factor of the graph is an exposure of the same lens sharing an
+        identical lens model, so the source-plane mesh (the image-mesh centres of this lead exposure,
+        ray-traced through the shared lens model) is built once here and returned inside a `PreloadsImaging`,
+        which `FactorGraphModel` forwards as the `shared` argument to every factor's
+        `log_likelihood_function`. Each exposure then maps its own (offset) data grid onto the shared mesh
+        instead of computing its own image-mesh and mesh ray-trace, so every exposure reconstructs on an
+        identical source-pixel grid.
+
+        Unlike the interferometer datacube case, the mapper, mapping matrix, curvature matrix and
+        regularization matrix are NOT shared — per-exposure PSFs and pixel offsets make the first three
+        per-dataset, and regularization may adapt to per-exposure data.
+
+        Returns `None` when the analysis has not opted in (`shared_preloads=False`) or when the model performs
+        no inversion, in which case no state is shared and every factor fits as normal.
+
+        The caller is responsible for the invariance contract: only enable `shared_preloads` when the factors
+        genuinely share the lens model, so the source-plane mesh really is exposure-invariant. The lead
+        factor's own `DatasetModel` offset (if any) is applied when the mesh is traced, so the mesh is defined
+        in the lead exposure's frame.
+        """
+        if not self.shared_preloads:
+            return None
+
+        fit = self.fit_from(instance=instance)
+
+        if not fit.perform_inversion:
+            return None
+
+        tracer_to_inversion = fit.tracer_to_inversion
+
+        return aa.PreloadsImaging(
+            source_plane_mesh_grid=tracer_to_inversion.traced_mesh_grid_pg_list,
+            image_plane_mesh_grid=tracer_to_inversion.image_plane_mesh_grid_pg_list,
+        )
+
     def fit_from(
         self,
         instance: af.ModelInstance,
+        preloads=None,
     ) -> FitImaging:
         """
         Given a model instance create a `FitImaging` object.
@@ -105,9 +199,10 @@ def fit_from(
         instance
             An instance of the model that is being fitted to the data by this analysis (whose parameters have been set
             via a non-linear search).
-        check_positions
-            Whether the multiple image positions of the lensed source should be checked, i.e. whether they trace
-            within the position threshold of one another in the source plane.
+        preloads
+            An optional `PreloadsImaging` carrying the exposure-invariant source-plane mesh geometry,
+            computed once and reused by the fit instead of being rebuilt. Supplied by the multi-exposure
+            shared-state path (see `shared_state_from`); `None` (the default) fits as normal.
 
         Returns
         -------
@@ -137,7 +232,8 @@ def fit_from(
             dataset_model=dataset_model,
             adapt_images=adapt_images,
             settings=self.settings,
-            xp=self._xp
+            xp=self._xp,
+            preloads=preloads,
         )
 
     def save_attributes(self, paths: af.DirectoryPaths):
@@ -221,7 +317,7 @@ def _register_fit_imaging_pytrees() -> None:
 
         register_instance_pytree(
             FitImaging,
-            no_flatten=("dataset", "adapt_images", "settings"),
+            no_flatten=("dataset", "adapt_images", "settings", "preloads"),
         )
         register_instance_pytree(DatasetModel)
         # ``cosmology`` is a fixed physical constant per fit; ride as aux.

autolens/lens/to_inversion.py

@@ -311,6 +311,11 @@ def image_plane_mesh_grid_pg_list(self) -> List[List]:
         -------
             The list of lists of image-plane mesh grids grouped by plane.
         """
+        if (
+            self._preloads is not None
+            and self._preloads.image_plane_mesh_grid is not None
+        ):
+            return self._preloads.image_plane_mesh_grid
 
         image_plane_mesh_grid_list_of_planes = []
 
@@ -351,6 +356,16 @@ def traced_mesh_grid_pg_list(self) -> List[List]:
         -------
             The list of lists of traced mesh grids grouped by plane.
         """
+        if (
+            self._preloads is not None
+            and self._preloads.source_plane_mesh_grid is not None
+        ):
+            # The shared-state path (e.g. `AnalysisImaging.shared_state_from`): the source-plane
+            # mesh geometry was traced once from the lead dataset and is reused here, so this
+            # dataset skips the image-mesh computation and mesh ray-trace. Its own (offset) data
+            # grid is still traced and mapped onto the shared mesh in `mapper_galaxy_dict`.
+            return self._preloads.source_plane_mesh_grid
+
         image_plane_mesh_grid_pg_list = self.image_plane_mesh_grid_pg_list
 
         traced_mesh_grid_pg_list = []

test_autolens/imaging/model/test_analysis_imaging.py

@@ -125,3 +125,144 @@ def test__positions__likelihood_overwrites__changes_likelihood__double_source_pl
 
     assert analysis_log_likelihood == pytest.approx(-44097289521.734665, 1.0e-4)
 
+
+
+def _shared_mesh_analysis(masked_imaging_7x7, shared_preloads):
+    """
+    An `AnalysisImaging` with an image-mesh (`Overlay`) + `Delaunay` pixelization, the regime the
+    multi-exposure shared-state path applies to (the source-plane mesh is traced from image-plane
+    mesh centres, so it can be shared across exposures).
+    """
+    import autoarray as aa
+
+    lens = al.Galaxy(
+        redshift=0.5,
+        mass=al.mp.Isothermal(centre=(0.1, 0.1), einstein_radius=1.0),
+    )
+
+    pixelization = al.Pixelization(
+        mesh=al.mesh.Delaunay(pixels=9, zeroed_pixels=0),
+        regularization=al.reg.Constant(coefficient=0.01),
+    )
+
+    source = al.Galaxy(redshift=1.0, pixelization=pixelization)
+
+    image_mesh = al.image_mesh.Overlay(shape=(3, 3))
+    image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(
+        mask=masked_imaging_7x7.mask,
+    )
+
+    adapt_images = al.AdaptImages(
+        galaxy_name_image_dict={
+            str(("galaxies", "source")): masked_imaging_7x7.data,
+        },
+        galaxy_name_image_plane_mesh_grid_dict={
+            str(("galaxies", "source")): image_plane_mesh_grid,
+        },
+    )
+
+    model = af.Collection(galaxies=af.Collection(lens=lens, source=source))
+
+    analysis = al.AnalysisImaging(
+        dataset=masked_imaging_7x7,
+        adapt_images=adapt_images,
+        use_jax=False,
+        shared_preloads=shared_preloads,
+        raise_inversion_positions_likelihood_exception=False,
+    )
+
+    return model, analysis
+
+
+def test__shared_state_from__mesh_reused__figure_of_merit_unchanged(
+    masked_imaging_7x7,
+):
+    import autoarray as aa
+
+    model, analysis = _shared_mesh_analysis(masked_imaging_7x7, shared_preloads=True)
+    instance = model.instance_from_unit_vector([])
+
+    # `shared_state_from` builds a `PreloadsImaging` carrying the source-plane mesh geometry (the
+    # exposure-invariant quantity) — NOT the mapper / curvature matrix / regularization matrix,
+    # which are per-exposure for imaging (PSFs, offsets, adaptive regularization).
+    shared = analysis.shared_state_from(instance=instance)
+    assert isinstance(shared, aa.PreloadsImaging)
+    assert shared.source_plane_mesh_grid is not None
+    assert shared.image_plane_mesh_grid is not None
+    assert shared.curvature_matrix is None
+    assert shared.mapper_galaxy_dict is None
+
+    # The preloaded mesh is reused by the fit (identity) and leaves the figure of merit unchanged.
+    fit_unshared = analysis.fit_from(instance=instance)
+    fom_unshared = fit_unshared.figure_of_merit
+
+    fit_shared = analysis.fit_from(instance=instance, preloads=shared)
+
+    assert (
+        fit_shared.tracer_to_inversion.traced_mesh_grid_pg_list
+        is shared.source_plane_mesh_grid
+    )
+    assert fit_shared.figure_of_merit == pytest.approx(fom_unshared)
+
+    # The full `log_likelihood_function` with the shared object matches the unshared call.
+    assert analysis.log_likelihood_function(
+        instance=instance, shared=shared
+    ) == pytest.approx(analysis.log_likelihood_function(instance=instance))
+
+
+def test__shared_state_from__returns_none_when_not_opted_in(masked_imaging_7x7):
+    model, analysis = _shared_mesh_analysis(masked_imaging_7x7, shared_preloads=False)
+    instance = model.instance_from_unit_vector([])
+
+    assert analysis.shared_state_from(instance=instance) is None
+
+
+def test__shared_state_from__returns_none_when_no_inversion(masked_imaging_7x7):
+    lens = al.Galaxy(redshift=0.5, light=al.lp.Sersic(intensity=0.1))
+
+    model = af.Collection(galaxies=af.Collection(lens=lens))
+    instance = model.instance_from_unit_vector([])
+
+    analysis = al.AnalysisImaging(
+        dataset=masked_imaging_7x7, use_jax=False, shared_preloads=True
+    )
+
+    assert analysis.shared_state_from(instance=instance) is None
+
+
+def _factor_graph_log_likelihood(masked_imaging_7x7, shared_preloads):
+    factors = []
+    model = None
+    for _ in range(2):
+        model, analysis = _shared_mesh_analysis(masked_imaging_7x7, shared_preloads)
+        factors.append(af.AnalysisFactor(prior_model=model.copy(), analysis=analysis))
+
+    factor_graph = af.FactorGraphModel(*factors)
+    instance = factor_graph.global_prior_model.instance_from_unit_vector([])
+    return factor_graph.log_likelihood_function(instance)
+
+
+def test__factor_graph__shared_vs_unshared_equal(masked_imaging_7x7):
+    ll_unshared = _factor_graph_log_likelihood(masked_imaging_7x7, shared_preloads=False)
+    ll_shared = _factor_graph_log_likelihood(masked_imaging_7x7, shared_preloads=True)
+
+    print(f"\nunshared={ll_unshared} shared={ll_shared}")
+    assert ll_shared == pytest.approx(ll_unshared, rel=1e-10)
+
+
+def test__factor_graph__shared_state_computed_once(masked_imaging_7x7, monkeypatch):
+    calls = {"n": 0}
+
+    original = al.AnalysisImaging.shared_state_from
+
+    def counting(self, instance):
+        result = original(self, instance)
+        if result is not None:
+            calls["n"] += 1
+        return result
+
+    monkeypatch.setattr(al.AnalysisImaging, "shared_state_from", counting)
+
+    _factor_graph_log_likelihood(masked_imaging_7x7, shared_preloads=True)
+
+    assert calls["n"] == 1