From 17b2c80d381121fbb70ce412d48044586742869c Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 15 May 2026 15:27:36 +0100 Subject: [PATCH] feat(DatasetModel): add grid_rotation_angle for multi-band rotation 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: Claude Opus 4.7 (1M context) --- autoarray/dataset/dataset_model.py | 75 +++++++++++-------- autoarray/fit/fit_dataset.py | 27 +++---- autoarray/structures/grids/irregular_2d.py | 36 +++++++++ autoarray/structures/grids/uniform_2d.py | 50 +++++++++++++ test_autoarray/fit/test_fit_dataset.py | 61 +++++++++++++++ .../structures/grids/test_irregular_2d.py | 35 +++++++++ .../structures/grids/test_uniform_2d.py | 37 +++++++++ 7 files changed, 273 insertions(+), 48 deletions(-) diff --git a/autoarray/dataset/dataset_model.py b/autoarray/dataset/dataset_model.py index e9c551b90..94e3b0f23 100644 --- a/autoarray/dataset/dataset_model.py +++ b/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 diff --git a/autoarray/fit/fit_dataset.py b/autoarray/fit/fit_dataset.py index 3f392ddcf..fd605883d 100644 --- a/autoarray/fit/fit_dataset.py +++ b/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, diff --git a/autoarray/structures/grids/irregular_2d.py b/autoarray/structures/grids/irregular_2d.py index 7fe398a20..4eafcc995 100644 --- a/autoarray/structures/grids/irregular_2d.py +++ b/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)) diff --git a/autoarray/structures/grids/uniform_2d.py b/autoarray/structures/grids/uniform_2d.py index bc5690726..10b2e34f2 100644 --- a/autoarray/structures/grids/uniform_2d.py +++ b/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": """ diff --git a/test_autoarray/fit/test_fit_dataset.py b/test_autoarray/fit/test_fit_dataset.py index a9f24f490..b83c39c8e 100644 --- a/test_autoarray/fit/test_fit_dataset.py +++ b/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) diff --git a/test_autoarray/structures/grids/test_irregular_2d.py b/test_autoarray/structures/grids/test_irregular_2d.py index e3c6aebbe..b71a5ffd5 100644 --- a/test_autoarray/structures/grids/test_irregular_2d.py +++ b/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) diff --git a/test_autoarray/structures/grids/test_uniform_2d.py b/test_autoarray/structures/grids/test_uniform_2d.py index 1e89ef138..2ed85a06a 100644 --- a/test_autoarray/structures/grids/test_uniform_2d.py +++ b/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)