From 6789f45179358e409d140e87509a1a491367e512 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 8 Jul 2026 18:44:23 +0100 Subject: [PATCH] feat: via_image_from image_is_convolved flag + from_gaussian oversample kwarg (#482 support) SimulatorImaging.via_image_from gains image_is_convolved (default False, behaviour unchanged): skip the internal PSF convolution when handed an already-convolved image, e.g. one convolved at the fine resolution by an oversampled Convolver. Convolver.from_gaussian passes through convolve_over_sample_size. Supports oversampled-PSF simulation (PyAutoLabs/PyAutoGalaxy#482). Co-Authored-By: Claude Fable 5 --- autoarray/dataset/imaging/simulator.py | 8 +- autoarray/operators/convolver.py | 5 + .../dataset/imaging/test_simulator.py | 115 ++++++++++++++++++ 3 files changed, 127 insertions(+), 1 deletion(-) diff --git a/autoarray/dataset/imaging/simulator.py b/autoarray/dataset/imaging/simulator.py index ec76ac9df..0e784cefe 100644 --- a/autoarray/dataset/imaging/simulator.py +++ b/autoarray/dataset/imaging/simulator.py @@ -129,6 +129,7 @@ def via_image_from( self, image: Array2D, over_sample_size: Optional[Union[int, np.ndarray]] = None, + image_is_convolved: bool = False, xp=None, ) -> Imaging: """ @@ -145,6 +146,9 @@ def via_image_from( over_sample_size If provided, the returned dataset has its over-sampling updated via `apply_over_sampling`. Should be an `Array2D` of integer sub-grid sizes with the same shape as the image. + image_is_convolved + If True, the input image has already been convolved with the PSF (e.g. at the fine resolution + by an oversampled Convolver) and the simulator's own convolution step is skipped. xp The array module to use for PSF convolution. When ``None`` (the default), falls back to ``self._xp`` — which is ``jnp`` if the simulator was constructed @@ -170,7 +174,9 @@ def via_image_from( pixel_scales=image.pixel_scales, ) - if self.use_real_space_convolution: + if image_is_convolved: + pass + elif self.use_real_space_convolution: image = self.psf.convolved_image_via_real_space_from( image=image, blurring_image=None, diff --git a/autoarray/operators/convolver.py b/autoarray/operators/convolver.py index 303c70f11..c81defd81 100644 --- a/autoarray/operators/convolver.py +++ b/autoarray/operators/convolver.py @@ -722,6 +722,7 @@ def from_gaussian( axis_ratio: float = 1.0, angle: float = 0.0, normalize: bool = False, + convolve_over_sample_size: int = 1, ) -> "Convolver": """ Setup the Convolver as a 2D symmetric elliptical Gaussian profile, according to the equation: @@ -746,6 +747,9 @@ def from_gaussian( The rotational angle of the Gaussian's ellipse defined counter clockwise from the positive x-axis. normalize If True, the Convolver's array values are normalized such that they sum to 1.0. + convolve_over_sample_size + The over sample size of the PSF (see ``Convolver.__init__``). When above 1 the + ``pixel_scales`` input should be the fine resolution (image pixel scale divided by this size). """ grid = Grid2D.uniform(shape_native=shape_native, pixel_scales=pixel_scales) @@ -780,6 +784,7 @@ def from_gaussian( return Convolver( kernel=gaussian, normalize=normalize, + convolve_over_sample_size=convolve_over_sample_size, ) @classmethod diff --git a/test_autoarray/dataset/imaging/test_simulator.py b/test_autoarray/dataset/imaging/test_simulator.py index 85c14ccfc..c21f5121a 100644 --- a/test_autoarray/dataset/imaging/test_simulator.py +++ b/test_autoarray/dataset/imaging/test_simulator.py @@ -173,3 +173,118 @@ def test__via_image_from__psf_on__psf_and_noise_both_on(image_central_delta_3x3) assert dataset.data.native == pytest.approx( np.array([[3.9, 5.35, 3.55], [5.85, 7.85, 5.5], [3.9, 5.3, 3.75]]), 1e-2 ) + + +def test__via_image_from__image_is_convolved__skips_psf_convolution(): + # An already-convolved image (e.g. from an oversampled Convolver) must pass + # through untouched by the simulator's own convolution step. + image = aa.Array2D.no_mask( + values=np.array([[0.0, 0.0, 0.0], [0.0, 4.0, 0.0], [0.0, 0.0, 0.0]]), + pixel_scales=1.0, + ) + + psf = aa.Convolver.from_gaussian( + shape_native=(3, 3), pixel_scales=1.0, sigma=0.5, normalize=True + ) + + simulator = aa.SimulatorImaging( + exposure_time=1.0, + psf=psf, + add_poisson_noise_to_data=False, + include_poisson_noise_in_noise_map=False, + noise_if_add_noise_false=1.0, + ) + + dataset_convolved = simulator.via_image_from(image=image) + dataset_passthrough = simulator.via_image_from(image=image, image_is_convolved=True) + + # The pass-through equals the input exactly; the convolved one does not. + assert np.array(dataset_passthrough.data.native) == pytest.approx( + np.array(image.native), abs=1.0e-14 + ) + assert not np.allclose( + np.array(dataset_convolved.data.native), np.array(image.native) + ) + + +def test__simulate_and_fit__oversampled_psf__consistent_with_fit_side_convolution(): + # The simulator's oversampled path (fine evaluation of the padded frame via + # image_is_convolved=True) must agree exactly, inside the mask, with the + # fit-side path (mask + blurring-region fine convolution) — the padding + # guarantees all flux within kernel reach of the mask is included in both. + s = 2 + pixel_scales = 1.0 + + def gaussian_on(grid_like, sigma=1.2, centre=(0.3, -0.4)): + arr = np.array(grid_like) + r2 = (arr[:, 0] - centre[0]) ** 2 + (arr[:, 1] - centre[1]) ** 2 + return np.exp(-0.5 * r2 / sigma**2) + + kernel_n = 9 + c = (np.arange(kernel_n) - (kernel_n - 1) / 2.0) * (pixel_scales / s) + yy, xx = np.meshgrid(-c, c, indexing="ij") + kernel = np.exp(-0.5 * (yy**2 + xx**2) / 0.8**2) + psf = aa.Convolver( + kernel=aa.Array2D.no_mask(values=kernel, pixel_scales=pixel_scales / s), + normalize=True, + convolve_over_sample_size=s, + ) + + # Simulator side: evaluate the padded frame fine, convolve, trim. + shape_native = (11, 11) + kernel_shape = psf.kernel_shape_image_resolution + padded_shape = ( + shape_native[0] + kernel_shape[0] - 1, + shape_native[1] + kernel_shape[1] - 1, + ) + padded_mask = aa.Mask2D.all_false( + shape_native=padded_shape, pixel_scales=pixel_scales + ) + padded_grid = aa.Grid2D.from_mask(mask=padded_mask, over_sample_size=s) + + convolved_padded = psf.convolved_image_from( + image=gaussian_on(padded_grid.over_sampled), + blurring_image=None, + mask=padded_mask, + ) + convolved_padded = aa.Array2D(values=convolved_padded, mask=padded_mask) + + simulator = aa.SimulatorImaging( + exposure_time=1.0, + psf=psf, + add_poisson_noise_to_data=False, + include_poisson_noise_in_noise_map=False, + noise_if_add_noise_false=1.0, + ) + dataset = simulator.via_image_from(image=convolved_padded, image_is_convolved=True) + dataset = dataset.trimmed_after_convolution_from(kernel_shape=kernel_shape) + + assert dataset.data.shape_native == shape_native + + # Fit side: mask + blurring-region fine convolution of the same scene. + mask = aa.Mask2D.circular( + shape_native=shape_native, pixel_scales=pixel_scales, radius=3.5 + ) + masked = aa.Imaging( + data=dataset.data, + noise_map=aa.Array2D.no_mask( + values=np.ones(shape_native), pixel_scales=pixel_scales + ), + psf=psf, + over_sample_size_lp=s, + over_sample_size_pixelization=s, + convolve_over_sample_size_lp=s, + convolve_over_sample_size_pixelization=s, + ).apply_mask(mask=mask) + + blurring_grid = masked.grids.blurring + model_data = masked.psf.convolved_image_from( + image=gaussian_on(masked.grids.lp.over_sampled), + blurring_image=gaussian_on(blurring_grid.over_sampled), + ) + + fit = aa.m.MockFitImaging( + dataset=masked, use_mask_in_fit=False, model_data=model_data + ) + + assert fit.chi_squared == pytest.approx(0.0, abs=1.0e-10)