Merge pull request #601 from PyAutoLabs/feature/multi-shared-state-scoped-preloads

fix: dataset-scoped preloads consumption — cross-type shared state reduces to mesh view

Author: Jammy2211

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

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

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(

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

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