From 2103f17a563b19230d71d705d3c6b214ec5ffc53 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 10 Jul 2026 14:45:32 +0100 Subject: [PATCH] =?UTF-8?q?feat:=20imaging=20shared-state=20consumer=20?= =?UTF-8?q?=E2=80=94=20AnalysisImaging.shared=5Fstate=5Ffrom=20+=20mesh=20?= =?UTF-8?q?preload=20reuse?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Multi-exposure shared-state path (PyAutoArray#379, design PyAutoLens#599 D1-D6 + H amendment): AnalysisImaging(shared_preloads=True) opts in; shared_state_from traces the lead exposure's image-mesh once and returns PreloadsImaging(source_plane_mesh_grid, image_plane_mesh_grid); shared= threads through log_likelihood_function -> fit_from -> FitImaging -> TracerToInversion, whose mesh pg-list properties now consult the preload so every exposure maps its own (offset) grid onto the identical shared mesh. Opt-in default-off; single-dataset path unchanged. Co-Authored-By: Claude Fable 5 --- autolens/imaging/fit_imaging.py | 10 +- autolens/imaging/model/analysis.py | 112 +++++++++++++- autolens/lens/to_inversion.py | 15 ++ .../imaging/model/test_analysis_imaging.py | 141 ++++++++++++++++++ 4 files changed, 269 insertions(+), 9 deletions(-) diff --git a/autolens/imaging/fit_imaging.py b/autolens/imaging/fit_imaging.py index a5652876d..ef9d12eb2 100644 --- a/autolens/imaging/fit_imaging.py +++ b/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 diff --git a/autolens/imaging/model/analysis.py b/autolens/imaging/model/analysis.py index bbc2edfe3..6bcf17d76 100644 --- a/autolens/imaging/model/analysis.py +++ b/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. diff --git a/autolens/lens/to_inversion.py b/autolens/lens/to_inversion.py index 96bb8134e..88054c3b2 100644 --- a/autolens/lens/to_inversion.py +++ b/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 = [] diff --git a/test_autolens/imaging/model/test_analysis_imaging.py b/test_autolens/imaging/model/test_analysis_imaging.py index a401df196..c1eb5b18a 100644 --- a/test_autolens/imaging/model/test_analysis_imaging.py +++ b/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