From dccff0a8801a147bf537fdc3dcf52174ec1cb2ec Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 10 Jul 2026 15:50:40 +0100 Subject: [PATCH] =?UTF-8?q?fix:=20dataset-scoped=20preloads=20consumption?= =?UTF-8?q?=20=E2=80=94=20cross-type=20shared=20state=20reduces=20to=20mes?= =?UTF-8?q?h=20view?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit In a joint imaging+interferometer FactorGraphModel, the lead factor's shared preloads are forwarded to EVERY factor. A PreloadsInterferometer carries a mapper + curvature matrix that embed ITS dataset's grids — consuming them in an imaging fit (or vice versa) would silently corrupt the likelihood. Both fits now scope their preloads at consumption (_preloads_scoped): same-type passes through by identity; cross-type reduces to the mesh-geometry view, the only part valid across dataset types (PyAutoLens#599 D5). AnalysisInterferometer.shared_state_from also populates the mesh fields so an interferometer lead shares the mesh cross-type. Co-Authored-By: Claude Fable 5 --- autolens/imaging/fit_imaging.py | 21 ++++++++- autolens/interferometer/fit_interferometer.py | 21 ++++++++- autolens/interferometer/model/analysis.py | 6 +++ .../imaging/model/test_analysis_imaging.py | 31 +++++++++++++ .../model/test_analysis_interferometer.py | 44 +++++++++++++++++++ 5 files changed, 121 insertions(+), 2 deletions(-) diff --git a/autolens/imaging/fit_imaging.py b/autolens/imaging/fit_imaging.py index ef9d12eb2..feebe9e25 100644 --- a/autolens/imaging/fit_imaging.py +++ b/autolens/imaging/fit_imaging.py @@ -130,6 +130,25 @@ def profile_subtracted_image(self) -> aa.Array2D: """ return self.data - self.blurred_image + @property + def _preloads_scoped(self): + """ + The preloads as consumed by this fit's inversion, scoped to its dataset type. + + Cross-dataset-type shared state (e.g. an interferometer lead factor in a joint + imaging + interferometer graph) may carry a mapper / curvature matrix that embed the + OTHER dataset's grids — consuming them here would silently corrupt the fit. Only the + source-plane mesh geometry is valid across dataset types, so any non-imaging preloads + are reduced to their mesh-geometry view. + """ + if self.preloads is None or isinstance(self.preloads, aa.PreloadsImaging): + return self.preloads + + return aa.PreloadsImaging( + source_plane_mesh_grid=self.preloads.source_plane_mesh_grid, + image_plane_mesh_grid=self.preloads.image_plane_mesh_grid, + ) + @property def tracer_to_inversion(self) -> TracerToInversion: @@ -147,7 +166,7 @@ def tracer_to_inversion(self) -> TracerToInversion: adapt_images=self.adapt_images, settings=self.settings, xp=self._xp, - preloads=self.preloads, + preloads=self._preloads_scoped, ) @cached_property diff --git a/autolens/interferometer/fit_interferometer.py b/autolens/interferometer/fit_interferometer.py index af5b6a3ff..f61bc5606 100644 --- a/autolens/interferometer/fit_interferometer.py +++ b/autolens/interferometer/fit_interferometer.py @@ -146,7 +146,26 @@ def tracer_to_inversion(self) -> TracerToInversion: adapt_images=self.adapt_images, settings=self.settings, xp=self._xp, - preloads=self.preloads, + preloads=self._preloads_scoped, + ) + + @property + def _preloads_scoped(self): + """ + The preloads as consumed by this fit's inversion, scoped to its dataset type. + + Cross-dataset-type shared state (e.g. an imaging lead factor in a joint + imaging + interferometer graph) is valid here only through its source-plane mesh + geometry — a mapper / curvature matrix from another dataset type would embed that + dataset's grids and silently corrupt the fit, so non-interferometer preloads are + reduced to their mesh-geometry view. + """ + if self.preloads is None or isinstance(self.preloads, aa.PreloadsInterferometer): + return self.preloads + + return aa.PreloadsInterferometer( + source_plane_mesh_grid=self.preloads.source_plane_mesh_grid, + image_plane_mesh_grid=self.preloads.image_plane_mesh_grid, ) @cached_property diff --git a/autolens/interferometer/model/analysis.py b/autolens/interferometer/model/analysis.py index f22c4e527..640be47fc 100644 --- a/autolens/interferometer/model/analysis.py +++ b/autolens/interferometer/model/analysis.py @@ -205,9 +205,15 @@ def shared_state_from(self, instance: af.ModelInstance): tracer_to_inversion = fit.tracer_to_inversion inversion = tracer_to_inversion.inversion + # The mesh-geometry fields are also populated so that cross-dataset-type factors of a + # joint graph (e.g. an imaging factor when this interferometer analysis leads) can share + # the source-plane mesh; they consume ONLY these fields (see `_preloads_scoped` on the + # fits) because the mapper and curvature matrix embed this dataset's grids. return aa.PreloadsInterferometer( curvature_matrix=inversion.curvature_matrix, mapper_galaxy_dict=tracer_to_inversion.mapper_galaxy_dict, + 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( diff --git a/test_autolens/imaging/model/test_analysis_imaging.py b/test_autolens/imaging/model/test_analysis_imaging.py index c1eb5b18a..e5d51ac73 100644 --- a/test_autolens/imaging/model/test_analysis_imaging.py +++ b/test_autolens/imaging/model/test_analysis_imaging.py @@ -266,3 +266,34 @@ def counting(self, instance): _factor_graph_log_likelihood(masked_imaging_7x7, shared_preloads=True) assert calls["n"] == 1 + + +def test__preloads_scoped__cross_type_preloads_reduced_to_mesh_view(masked_imaging_7x7): + import autoarray as aa + + lens = al.Galaxy(redshift=0.5, light=al.lp.Sersic(intensity=0.1)) + tracer = al.Tracer(galaxies=[lens]) + + # Cross-dataset-type preloads (e.g. from an interferometer lead factor in a joint graph): + # the mapper / curvature matrix embed the other dataset's grids and must NOT be consumed + # by an imaging fit — only the mesh-geometry view survives the scoping. + cross_type = aa.PreloadsInterferometer( + curvature_matrix="other-datasets-F", + mapper_galaxy_dict="other-datasets-mapper", + source_plane_mesh_grid=[["mesh"]], + image_plane_mesh_grid=[["image-mesh"]], + ) + + fit = al.FitImaging(dataset=masked_imaging_7x7, tracer=tracer, preloads=cross_type) + + scoped = fit._preloads_scoped + assert isinstance(scoped, aa.PreloadsImaging) + assert scoped.source_plane_mesh_grid == [["mesh"]] + assert scoped.image_plane_mesh_grid == [["image-mesh"]] + assert scoped.curvature_matrix is None + assert scoped.mapper_galaxy_dict is None + + # Same-type preloads pass through untouched. + same_type = aa.PreloadsImaging(source_plane_mesh_grid=[["mesh"]]) + fit = al.FitImaging(dataset=masked_imaging_7x7, tracer=tracer, preloads=same_type) + assert fit._preloads_scoped is same_type diff --git a/test_autolens/interferometer/model/test_analysis_interferometer.py b/test_autolens/interferometer/model/test_analysis_interferometer.py index d8dd65675..5b796e869 100644 --- a/test_autolens/interferometer/model/test_analysis_interferometer.py +++ b/test_autolens/interferometer/model/test_analysis_interferometer.py @@ -130,3 +130,47 @@ def test__shared_state_from__returns_none_when_not_opted_in(interferometer_7): analysis = al.AnalysisInterferometer(dataset=dataset, use_jax=False) assert analysis.shared_state_from(instance=instance) is None + + +def test__preloads_scoped__cross_type_preloads_reduced_to_mesh_view(interferometer_7): + lens = al.Galaxy(redshift=0.5, light=al.lp.Sersic(intensity=0.1)) + tracer = al.Tracer(galaxies=[lens]) + + # Cross-dataset-type preloads (e.g. from an imaging lead factor in a joint graph): only + # the mesh-geometry view is valid for an interferometer fit. + cross_type = aa.PreloadsImaging( + source_plane_mesh_grid=[["mesh"]], image_plane_mesh_grid=[["image-mesh"]] + ) + + fit = al.FitInterferometer( + dataset=interferometer_7, tracer=tracer, preloads=cross_type + ) + + scoped = fit._preloads_scoped + assert isinstance(scoped, aa.PreloadsInterferometer) + assert scoped.source_plane_mesh_grid == [["mesh"]] + assert scoped.curvature_matrix is None + assert scoped.mapper_galaxy_dict is None + + same_type = aa.PreloadsInterferometer(curvature_matrix="F") + fit = al.FitInterferometer( + dataset=interferometer_7, tracer=tracer, preloads=same_type + ) + assert fit._preloads_scoped is same_type + + +def test__shared_state_from__populates_mesh_geometry_fields(interferometer_7): + dataset = interferometer_7.apply_sparse_operator(use_jax=False) + + model = _pixelization_model() + instance = model.instance_from_unit_vector([]) + + analysis = al.AnalysisInterferometer( + dataset=dataset, use_jax=False, shared_preloads=True + ) + + # The mesh-geometry fields ride alongside the curvature matrix + mapper so that + # cross-dataset-type factors of a joint graph can consume the shared mesh. + shared = analysis.shared_state_from(instance=instance) + assert shared.source_plane_mesh_grid is not None + assert shared.image_plane_mesh_grid is not None