feat(AdaptImages): rotate cached mesh grid with DatasetModel transforms (#416)

The image-plane mesh grid cached in galaxy_name_image_plane_mesh_grid_dict
now follows the same shift+rotate transform that FitDataset.grids applies
to the data grid, so adaptive pixelization fits stay aligned in multi-band
work. Also plumbs dataset_model + xp through adapt_images_via_instance_from
and adds a grid_rotation_angle prior to config/priors/dataset_model.yaml.

This is the actual fix for the source-reconstruction misalignment in
@qiuhan06's dev_Q prototype; he had it correct but never re-tested it.

Adds 3 chi^2=0 simulate-and-fit tests in test_simulate_and_fit_imaging.py
covering offset-only, rotation-only, and combined cases.

Refs PyAutoLens#511, PyAutoArray#312.

Co-authored-by: Jammy2211 <JNightingale2211@gmail.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: 3 people

autogalaxy/analysis/adapt_images/adapt_images.py

@@ -157,7 +157,12 @@ def model_image(self) -> aa.Array2D:
         return adapt_model_image
 
     def updated_via_instance_from(
-        self, instance, mask=None, galaxies: Optional[List["Galaxy"]] = None
+        self,
+        instance,
+        dataset_model: Optional["aa.DatasetModel"] = None,
+        mask=None,
+        galaxies: Optional[List["Galaxy"]] = None,
+        xp=np,
     ) -> "AdaptImages":
         """
         Returns adapt-images which have been updated to map galaxy instances instead of galaxy names.
@@ -174,10 +179,19 @@ def updated_via_instance_from(
         galaxy instances are also created on-fly from the database. Database images do not have a mask, so it is
         also applied to the adapt images on-the-fly during database loading.
 
+        When a ``dataset_model`` is supplied with a non-trivial ``grid_offset`` or ``grid_rotation_angle``, the cached
+        ``galaxy_name_image_plane_mesh_grid_dict`` entries are transformed into the same frame as the dataset's
+        image-plane grid (which ``FitDataset.grids`` rotates by the same amount). Without this transform the cached
+        mesh and the data grid would sit in different frames, producing a misaligned source reconstruction.
+
         Parameters
         ----------
         instance
             The instance of the model-fit (e.g. in a non-linear search) which is used to update the adapt images.
+        dataset_model
+            The dataset model whose ``grid_offset`` and ``grid_rotation_angle`` are applied to cached mesh grids so
+            they remain consistent with the rotated/shifted data grid produced by ``FitDataset.grids``. If ``None``,
+            the cached mesh grids are passed through unchanged.
         mask
             A mask which can be applied to the adapt images, which is used when setting up the adaptive images
             via the aggregator and autofit database tools.
@@ -188,6 +202,8 @@ def updated_via_instance_from(
             galaxy instances into fresh objects. When ``None`` the path list is populated in ``path_instance_tuples_for_class``
             order, which matches ``Analysis.galaxies_via_instance_from`` for the common case (no
             ``extra_galaxies`` / ``scaling_galaxies``).
+        xp
+            Array backend (``numpy`` or ``jax.numpy``) used when transforming cached mesh grids.
 
         Returns
         -------
@@ -226,9 +242,14 @@ def updated_via_instance_from(
                 galaxy_name = str(galaxy_name)
 
                 if galaxy_name in self.galaxy_name_image_plane_mesh_grid_dict:
-                    galaxy_image_plane_mesh_grid_dict[galaxy] = (
-                        self.galaxy_name_image_plane_mesh_grid_dict[galaxy_name]
-                    )
+                    cached_mesh = self.galaxy_name_image_plane_mesh_grid_dict[galaxy_name]
+                    if dataset_model is not None:
+                        cached_mesh = cached_mesh.subtracted_and_rotated_from(
+                            offset=dataset_model.grid_offset,
+                            angle=dataset_model.grid_rotation_angle,
+                            xp=xp,
+                        )
+                    galaxy_image_plane_mesh_grid_dict[galaxy] = cached_mesh
 
         if galaxies is not None:
             galaxy_path_list = [path_by_id.get(id(g)) for g in galaxies]

autogalaxy/analysis/analysis/dataset.py

@@ -172,11 +172,16 @@ def save_results(self, paths: af.DirectoryPaths, result: ResultDataset):
     def adapt_images_via_instance_from(
         self,
         instance: af.ModelInstance,
+        dataset_model: Optional[aa.DatasetModel] = None,
         galaxies=None,
+        xp=np,
     ) -> AdaptImages:
         try:
             return self.adapt_images.updated_via_instance_from(
-                instance=instance, galaxies=galaxies
+                instance=instance,
+                dataset_model=dataset_model,
+                galaxies=galaxies,
+                xp=xp,
             )
         except AttributeError:
             pass

autogalaxy/config/priors/dataset_model.yaml

@@ -6,5 +6,8 @@ DatasetModel:
     type: Constant
     value: 0.0
   grid_offset_1:
+    type: Constant
+    value: 0.0
+  grid_rotation_angle:
     type: Constant
     value: 0.0

autogalaxy/imaging/model/analysis.py

@@ -153,7 +153,10 @@ def fit_from(self, instance: af.ModelInstance) -> FitImaging:
         dataset_model = self.dataset_model_via_instance_from(instance=instance)
 
         adapt_images = self.adapt_images_via_instance_from(
-            instance=instance, galaxies=galaxies
+            instance=instance,
+            dataset_model=dataset_model,
+            galaxies=galaxies,
+            xp=self._xp,
         )
 
         return FitImaging(

test_autogalaxy/imaging/test_simulate_and_fit_imaging.py

@@ -81,6 +81,130 @@ def test__perfect_fit__simulate_and_reload__chi_squared_zero():
         shutil.rmtree(file_path)
 
 
+def _perfect_fit_dataset(galaxies, grid):
+    """Helper: simulate noiseless imaging and zero out the noise map for chi^2 tests."""
+    psf = ag.Convolver.from_gaussian(
+        shape_native=(3, 3), pixel_scales=grid.pixel_scales[0], sigma=0.05, normalize=True
+    )
+    simulator = ag.SimulatorImaging(
+        exposure_time=300.0, psf=psf, add_poisson_noise_to_data=False
+    )
+    dataset = simulator.via_galaxies_from(galaxies=galaxies, grid=grid)
+    dataset.noise_map = ag.Array2D.ones(
+        shape_native=dataset.data.shape_native, pixel_scales=grid.pixel_scales
+    )
+    return dataset
+
+
+def test__perfect_fit__sim_offset_centre__fit_with_dataset_model_grid_offset__chi_squared_zero():
+    """Sim a profile with offset centre; fit with origin profile + DatasetModel.grid_offset."""
+    grid = ag.Grid2D.uniform(shape_native=(31, 31), pixel_scales=0.2, over_sample_size=1)
+    centre = (0.3, 0.2)
+
+    sim_galaxy = ag.Galaxy(
+        redshift=0.5,
+        light=ag.lp.Sersic(centre=centre, intensity=0.5, effective_radius=0.5),
+    )
+    dataset = _perfect_fit_dataset([sim_galaxy], grid)
+    mask = ag.Mask2D.circular(
+        shape_native=dataset.data.shape_native, pixel_scales=0.2, radius=2.5
+    )
+    masked = dataset.apply_mask(mask=mask)
+
+    fit_galaxy = ag.Galaxy(
+        redshift=0.5,
+        light=ag.lp.Sersic(centre=(0.0, 0.0), intensity=0.5, effective_radius=0.5),
+    )
+    dataset_model = ag.DatasetModel(grid_offset=centre)
+    fit = ag.FitImaging(
+        dataset=masked, galaxies=[fit_galaxy], dataset_model=dataset_model
+    )
+
+    assert fit.chi_squared == pytest.approx(0.0, abs=1e-4)
+
+
+def test__perfect_fit__sim_rotated_ellipse__fit_with_dataset_model_grid_rotation__chi_squared_zero():
+    """Sim a rotated ellipse; fit with axis-aligned profile + DatasetModel.grid_rotation_angle.
+
+    Convention: profile with ell-angle theta is equivalent to grid rotated by -theta,
+    so fit with grid_rotation_angle=-theta to recover chi^2 = 0.
+    """
+    grid = ag.Grid2D.uniform(shape_native=(31, 31), pixel_scales=0.2, over_sample_size=1)
+    theta = 15.0
+
+    sim_galaxy = ag.Galaxy(
+        redshift=0.5,
+        light=ag.lp.Sersic(
+            centre=(0.0, 0.0),
+            ell_comps=ag.convert.ell_comps_from(axis_ratio=0.6, angle=theta),
+            intensity=0.5,
+            effective_radius=0.5,
+        ),
+    )
+    dataset = _perfect_fit_dataset([sim_galaxy], grid)
+    mask = ag.Mask2D.circular(
+        shape_native=dataset.data.shape_native, pixel_scales=0.2, radius=2.5
+    )
+    masked = dataset.apply_mask(mask=mask)
+
+    fit_galaxy = ag.Galaxy(
+        redshift=0.5,
+        light=ag.lp.Sersic(
+            centre=(0.0, 0.0),
+            ell_comps=ag.convert.ell_comps_from(axis_ratio=0.6, angle=0.0),
+            intensity=0.5,
+            effective_radius=0.5,
+        ),
+    )
+    dataset_model = ag.DatasetModel(grid_rotation_angle=-theta)
+    fit = ag.FitImaging(
+        dataset=masked, galaxies=[fit_galaxy], dataset_model=dataset_model
+    )
+
+    assert fit.chi_squared == pytest.approx(0.0, abs=1e-4)
+
+
+def test__perfect_fit__sim_offset_and_rotated__fit_with_dataset_model_offset_and_rotation__chi_squared_zero():
+    """Combined: sim with offset centre AND rotated ellipse, fit with identity profile +
+    DatasetModel(grid_offset, grid_rotation_angle)."""
+    grid = ag.Grid2D.uniform(shape_native=(31, 31), pixel_scales=0.2, over_sample_size=1)
+    centre = (0.3, 0.2)
+    theta = 12.0
+
+    sim_galaxy = ag.Galaxy(
+        redshift=0.5,
+        light=ag.lp.Sersic(
+            centre=centre,
+            ell_comps=ag.convert.ell_comps_from(axis_ratio=0.6, angle=theta),
+            intensity=0.5,
+            effective_radius=0.5,
+        ),
+    )
+    dataset = _perfect_fit_dataset([sim_galaxy], grid)
+    mask = ag.Mask2D.circular(
+        shape_native=dataset.data.shape_native, pixel_scales=0.2, radius=2.5
+    )
+    masked = dataset.apply_mask(mask=mask)
+
+    fit_galaxy = ag.Galaxy(
+        redshift=0.5,
+        light=ag.lp.Sersic(
+            centre=(0.0, 0.0),
+            ell_comps=ag.convert.ell_comps_from(axis_ratio=0.6, angle=0.0),
+            intensity=0.5,
+            effective_radius=0.5,
+        ),
+    )
+    dataset_model = ag.DatasetModel(
+        grid_offset=centre, grid_rotation_angle=-theta
+    )
+    fit = ag.FitImaging(
+        dataset=masked, galaxies=[fit_galaxy], dataset_model=dataset_model
+    )
+
+    assert fit.chi_squared == pytest.approx(0.0, abs=1e-4)
+
+
 def test__simulate_imaging_data_and_fit__standard_galaxies__known_figure_of_merit():
     grid = ag.Grid2D.uniform(shape_native=(31, 31), pixel_scales=0.2)