feat(DatasetModel): add grid_rotation_angle for multi-band rotation (#312)

Mirrors the existing grid_offset pattern so each band in a multi-band fit
can be rotated as well as shifted relative to a reference dataset. Adds
Grid2D / Grid2DIrregular.subtracted_and_rotated_from helpers (shift then
rotate CCW about the offset point) and wires FitDataset.grids through
them so lp / pixelization / blurring grids all carry the transform into
the fit.

Refs PyAutoLens#511, prototype by @qiuhan06 on dev_Q.

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

Author: 3 people

autoarray/dataset/dataset_model.py

@@ -1,33 +1,42 @@
-from typing import Tuple
-
-
-class DatasetModel:
-    def __init__(
-        self,
-        background_sky_level: float = 0.0,
-        grid_offset: Tuple[float, float] = (0.0, 0.0),
-    ):
-        """
-        Attributes which allow for parts of a dataset to be treated as a model, meaning they can be fitted
-        for in the `fit` module.
-
-        The following aspects of a dataset can be treated as a model:
-
-         - `background_sky_level`: The data may have a constant signal in the background which is estimated
-           and subtracted from the data beforehand with a degree of uncertainty. By including it in the model it can be
-           marginalized over. Units are dimensionless and derived from the data.
-
-        - `grid_offset`: Two datasets may be offset from one another, for example if they are taken with different
-          pointing positions. This offset can be included in the model and marginalized over. Units are arc seconds.
-
-        Parameters
-        ----------
-        background_sky_level
-            Overall normalisation of the sky which is added or subtracted from the data. Units are dimensionless and
-            derived from the data, which is expected to be electrons per second in Astronomy analyses.
-        grid_offset
-            Offset between two datasets, in arc seconds. This is used to align datasets which are taken at different
-            pointing positions.
-        """
-        self.background_sky_level = background_sky_level
-        self.grid_offset = grid_offset
+from typing import Tuple
+
+
+class DatasetModel:
+    def __init__(
+        self,
+        background_sky_level: float = 0.0,
+        grid_offset: Tuple[float, float] = (0.0, 0.0),
+        grid_rotation_angle: float = 0.0,
+    ):
+        """
+        Attributes which allow for parts of a dataset to be treated as a model, meaning they can be fitted
+        for in the `fit` module.
+
+        The following aspects of a dataset can be treated as a model:
+
+         - `background_sky_level`: The data may have a constant signal in the background which is estimated
+           and subtracted from the data beforehand with a degree of uncertainty. By including it in the model it can be
+           marginalized over. Units are dimensionless and derived from the data.
+
+        - `grid_offset`: Two datasets may be offset from one another, for example if they are taken with different
+          pointing positions. This offset can be included in the model and marginalized over. Units are arc seconds.
+
+        - `grid_rotation_angle`: Two datasets may also be rotated relative to one another (e.g. a different
+          telescope roll angle). This rotation can be included in the model and marginalized over. Units are degrees,
+          with positive angles applying a counter-clockwise rotation about the offset point.
+
+        Parameters
+        ----------
+        background_sky_level
+            Overall normalisation of the sky which is added or subtracted from the data. Units are dimensionless and
+            derived from the data, which is expected to be electrons per second in Astronomy analyses.
+        grid_offset
+            Offset between two datasets, in arc seconds. This is used to align datasets which are taken at different
+            pointing positions.
+        grid_rotation_angle
+            Rotation between two datasets, in degrees. Applied counter-clockwise about the offset point after the
+            offset is subtracted. Used to align datasets which share a pointing centre but differ in roll angle.
+        """
+        self.background_sky_level = background_sky_level
+        self.grid_offset = grid_offset
+        self.grid_rotation_angle = grid_rotation_angle

autoarray/fit/fit_dataset.py

@@ -185,26 +185,23 @@ def mask(self) -> Mask2D:
     @property
     def grids(self) -> GridsInterface:
         """
-        The grids of (y,x) coordinates associated with the dataset, adjusted by any `grid_offset` specified in
-        the `dataset_model`. Each grid (`lp`, `pixelization`, `blurring`) has the offset subtracted from it
-        before being returned.
+        The grids of (y,x) coordinates associated with the dataset, adjusted by any `grid_offset` and
+        `grid_rotation_angle` specified in the `dataset_model`. Each grid (`lp`, `pixelization`, `blurring`)
+        has the offset subtracted from it and is then rotated counter-clockwise by `grid_rotation_angle`
+        about the offset point before being returned.
         """
 
-        def subtracted_from(grid, offset):
+        offset = self.dataset_model.grid_offset
+        angle = self.dataset_model.grid_rotation_angle
+
+        def shift_and_rotate(grid):
             if grid is None:
                 return None
+            return grid.subtracted_and_rotated_from(offset=offset, angle=angle, xp=self._xp)
 
-            return grid.subtracted_from(offset=offset, xp=self._xp)
-
-        lp = subtracted_from(
-            grid=self.dataset.grids.lp, offset=self.dataset_model.grid_offset
-        )
-        pixelization = subtracted_from(
-            grid=self.dataset.grids.pixelization, offset=self.dataset_model.grid_offset
-        )
-        blurring = subtracted_from(
-            grid=self.dataset.grids.blurring, offset=self.dataset_model.grid_offset
-        )
+        lp = shift_and_rotate(self.dataset.grids.lp)
+        pixelization = shift_and_rotate(self.dataset.grids.pixelization)
+        blurring = shift_and_rotate(self.dataset.grids.blurring)
 
         return GridsInterface(
             lp=lp,

autoarray/structures/grids/irregular_2d.py

@@ -279,3 +279,39 @@ def grid_of_closest_from(self, grid_pair: "Grid2DIrregular") -> "Grid2DIrregular
         closest_points = self.array[closest_idx]
 
         return Grid2DIrregular(closest_points)
+
+    def subtracted_from(self, offset, xp=np) -> "Grid2DIrregular":
+        """
+        Return a new Grid2DIrregular with ``offset`` subtracted from every (y, x) coordinate.
+        """
+        offset_array = xp.array(offset)
+        return Grid2DIrregular(self.array - offset_array)
+
+    def subtracted_and_rotated_from(
+        self, offset, angle: float, xp=np
+    ) -> "Grid2DIrregular":
+        """
+        Return a new Grid2DIrregular where the (y, x) coordinates of this grid have an offset
+        subtracted and are then rotated counter-clockwise by ``angle`` (in degrees) about the
+        offset point.
+
+        Order matches :meth:`Grid2D.subtracted_and_rotated_from`: shift, then rotate.
+
+        Parameters
+        ----------
+        offset
+            The (y, x) offset subtracted from every grid coordinate before rotation.
+        angle
+            The rotation angle in degrees. Positive values rotate counter-clockwise.
+        """
+        offset_array = xp.array(offset)
+        angle_rad = xp.deg2rad(angle)
+        cos_a = xp.cos(angle_rad)
+        sin_a = xp.sin(angle_rad)
+
+        shifted = self.array - offset_array
+        sy = shifted[:, 0]
+        sx = shifted[:, 1]
+        ry = sx * sin_a + sy * cos_a
+        rx = sx * cos_a - sy * sin_a
+        return Grid2DIrregular(xp.stack((ry, rx), axis=-1))

autoarray/structures/grids/uniform_2d.py

@@ -736,6 +736,56 @@ def subtracted_from(
             over_sampler=self.over_sampler,
         )
 
+    def subtracted_and_rotated_from(
+        self, offset: Tuple[float, float], angle: float, xp=np
+    ) -> "Grid2D":
+        """
+        Return a new Grid2D where the (y, x) coordinates of this grid have an offset subtracted
+        and are then rotated counter-clockwise by ``angle`` (in degrees) about the offset point.
+
+        Order: shift first, then rotate. With ``offset = (oy, ox)`` and ``angle = theta`` (degrees):
+
+          (y', x') = (y - oy, x - ox)
+          y'' = y' cos(theta) + x' sin(theta)
+          x'' = x' cos(theta) - y' sin(theta)
+
+        Parameters
+        ----------
+        offset
+            The (y, x) offset subtracted from every grid coordinate before rotation.
+        angle
+            The rotation angle in degrees. Positive values rotate counter-clockwise.
+        """
+        offset_array = xp.array(offset)
+        angle_rad = xp.deg2rad(angle)
+        cos_a = xp.cos(angle_rad)
+        sin_a = xp.sin(angle_rad)
+
+        def _shift_and_rotate(grid_array):
+            shifted = grid_array - offset_array
+            sy = shifted[:, 0]
+            sx = shifted[:, 1]
+            ry = sx * sin_a + sy * cos_a
+            rx = sx * cos_a - sy * sin_a
+            return xp.stack((ry, rx), axis=-1)
+
+        grid_rotated = _shift_and_rotate(self.array)
+        over_sampled_rotated = _shift_and_rotate(self.over_sampled.array)
+
+        mask = Mask2D(
+            mask=self.mask,
+            pixel_scales=self.pixel_scales,
+            origin=(self.origin[0] - offset[0], self.origin[1] - offset[1]),
+        )
+
+        return Grid2D(
+            values=grid_rotated,
+            mask=mask,
+            over_sample_size=self.over_sample_size,
+            over_sampled=Grid2DIrregular(over_sampled_rotated),
+            over_sampler=self.over_sampler,
+        )
+
     @property
     def slim(self) -> "Grid2D":
         """

test_autoarray/fit/test_fit_dataset.py

@@ -112,3 +112,64 @@ def test__grids__with_dataset_model_grid_offset__lp_and_pixelization_grids_offse
     assert fit.dataset_model.grid_offset == (1.0, 2.0)
     assert fit.grids.lp[0] == pytest.approx((0.0, -3.0), 1.0e-4)
     assert fit.grids.pixelization[0] == pytest.approx((0.0, -3.0), 1.0e-4)
+
+
+def test__grids__with_dataset_model_grid_rotation_angle__lp_grid_rotated_correctly(
+    imaging_7x7, mask_2d_7x7, model_image_7x7
+):
+    masked_imaging_7x7 = imaging_7x7.apply_mask(mask=mask_2d_7x7)
+
+    # Rotation by 90 degrees CCW about the origin maps (y, x) -> (x, -y).
+    fit = aa.m.MockFitImaging(
+        dataset=masked_imaging_7x7,
+        use_mask_in_fit=False,
+        model_data=model_image_7x7,
+        dataset_model=aa.DatasetModel(grid_rotation_angle=90.0),
+    )
+
+    assert fit.dataset_model.grid_rotation_angle == 90.0
+    # 90 deg CCW rotation in (y, x) order maps (y, x) -> (x, -y).
+    original = masked_imaging_7x7.grids.lp[0]
+    rotated = fit.grids.lp[0]
+    assert rotated[0] == pytest.approx(original[1], 1.0e-4)
+    assert rotated[1] == pytest.approx(-original[0], 1.0e-4)
+
+
+def test__grids__with_grid_offset_and_grid_rotation_angle__shift_then_rotate(
+    imaging_7x7, mask_2d_7x7, model_image_7x7
+):
+    masked_imaging_7x7 = imaging_7x7.apply_mask(mask=mask_2d_7x7)
+
+    fit = aa.m.MockFitImaging(
+        dataset=masked_imaging_7x7,
+        use_mask_in_fit=False,
+        model_data=model_image_7x7,
+        dataset_model=aa.DatasetModel(
+            grid_offset=(1.0, 2.0), grid_rotation_angle=90.0
+        ),
+    )
+
+    # First subtract the offset, then rotate 90deg CCW: (y, x) -> (x, -y).
+    original = masked_imaging_7x7.grids.lp[0]
+    shifted_y = original[0] - 1.0
+    shifted_x = original[1] - 2.0
+    expected = (shifted_x, -shifted_y)
+
+    assert fit.grids.lp[0] == pytest.approx(expected, 1.0e-4)
+
+
+def test__grids__with_grid_rotation_angle_zero__matches_subtracted_from(
+    imaging_7x7, mask_2d_7x7, model_image_7x7
+):
+    masked_imaging_7x7 = imaging_7x7.apply_mask(mask=mask_2d_7x7)
+
+    fit_rotated = aa.m.MockFitImaging(
+        dataset=masked_imaging_7x7,
+        use_mask_in_fit=False,
+        model_data=model_image_7x7,
+        dataset_model=aa.DatasetModel(grid_offset=(1.0, 2.0), grid_rotation_angle=0.0),
+    )
+
+    # angle=0 is identity rotation, so the result must equal the offset-only path.
+    assert fit_rotated.grids.lp[0] == pytest.approx((0.0, -3.0), 1.0e-4)
+    assert fit_rotated.grids.pixelization[0] == pytest.approx((0.0, -3.0), 1.0e-4)

test_autoarray/structures/grids/test_irregular_2d.py

@@ -1,6 +1,7 @@
 import os
 import shutil
 import numpy as np
+import pytest
 
 import autoarray as aa
 
@@ -125,3 +126,37 @@ def test__grid_of_closest_from():
     assert (
         grid_of_closest == np.array([[0.0, 0.0], [0.0, 0.0], [0.0, 1.0], [0.0, 0.0]])
     ).all()
+
+
+def test__subtracted_from():
+    grid = aa.Grid2DIrregular(np.array([[1.0, 2.0], [3.0, 4.0]]))
+
+    shifted = grid.subtracted_from(offset=(0.5, -0.5))
+
+    assert shifted.array == pytest.approx(np.array([[0.5, 2.5], [2.5, 4.5]]), 1.0e-4)
+
+
+def test__subtracted_and_rotated_from__zero_angle_is_pure_shift():
+    grid = aa.Grid2DIrregular(np.array([[1.0, 2.0], [3.0, 4.0]]))
+
+    shifted = grid.subtracted_and_rotated_from(offset=(0.5, -0.5), angle=0.0)
+
+    assert shifted.array == pytest.approx(np.array([[0.5, 2.5], [2.5, 4.5]]), 1.0e-4)
+
+
+def test__subtracted_and_rotated_from__90_degrees_about_origin():
+    grid = aa.Grid2DIrregular(np.array([[1.0, 0.0], [0.0, 1.0]]))
+
+    rotated = grid.subtracted_and_rotated_from(offset=(0.0, 0.0), angle=90.0)
+
+    # 90 deg CCW in (y, x) order maps (y, x) -> (x, -y).
+    assert rotated.array == pytest.approx(np.array([[0.0, -1.0], [1.0, 0.0]]), 1.0e-4)
+
+
+def test__subtracted_and_rotated_from__shift_first_then_rotate():
+    grid = aa.Grid2DIrregular(np.array([[2.0, 3.0]]))
+
+    rotated = grid.subtracted_and_rotated_from(offset=(1.0, 1.0), angle=90.0)
+
+    # Shifted -> (1.0, 2.0); 90 deg CCW -> (2.0, -1.0).
+    assert rotated.array == pytest.approx(np.array([[2.0, -1.0]]), 1.0e-4)

test_autoarray/structures/grids/test_uniform_2d.py

@@ -830,3 +830,40 @@ def test__apply_over_sampling():
     grid = grid.apply_over_sampling(over_sample_size=2)
 
     assert grid.over_sampled.shape[0] == 16
+
+
+def test__subtracted_and_rotated_from__zero_angle_is_pure_shift():
+    grid = aa.Grid2D.uniform(shape_native=(3, 3), pixel_scales=1.0, over_sample_size=1)
+
+    shifted = grid.subtracted_and_rotated_from(offset=(0.5, -0.5), angle=0.0)
+
+    expected = grid.array - np.array([0.5, -0.5])
+    assert shifted.array == pytest.approx(expected, 1.0e-4)
+
+
+def test__subtracted_and_rotated_from__90_degrees_about_origin():
+    grid = aa.Grid2D.uniform(shape_native=(3, 3), pixel_scales=1.0, over_sample_size=1)
+
+    rotated = grid.subtracted_and_rotated_from(offset=(0.0, 0.0), angle=90.0)
+
+    # 90 deg CCW in (y, x) order maps (y, x) -> (x, -y).
+    expected = np.stack((grid.array[:, 1], -grid.array[:, 0]), axis=-1)
+    assert rotated.array == pytest.approx(expected, 1.0e-4)
+
+
+def test__subtracted_and_rotated_from__180_degrees_inverts_coordinates():
+    grid = aa.Grid2D.uniform(shape_native=(3, 3), pixel_scales=1.0, over_sample_size=1)
+
+    rotated = grid.subtracted_and_rotated_from(offset=(0.0, 0.0), angle=180.0)
+
+    assert rotated.array == pytest.approx(-grid.array, 1.0e-4)
+
+
+def test__subtracted_and_rotated_from__shift_first_then_rotate():
+    grid = aa.Grid2D.uniform(shape_native=(3, 3), pixel_scales=1.0, over_sample_size=1)
+
+    shifted = grid.array - np.array([1.0, 2.0])
+    expected = np.stack((shifted[:, 1], -shifted[:, 0]), axis=-1)
+
+    rotated = grid.subtracted_and_rotated_from(offset=(1.0, 2.0), angle=90.0)
+    assert rotated.array == pytest.approx(expected, 1.0e-4)