feat: oversampled PSF support in SimulatorImaging (#482) (#483)

OperateImage.convolved_padded_image_2d_from (shared by Galaxies and Tracer):
builds the padded all-false frame from the kernel's image-resolution
footprint with a uniform over-sample size (the regular padded grid pads its
border with size-1 entries, which the uniformity guard correctly rejects),
evaluates unbinned on the padded over-sampled coordinates, convolves at the
fine resolution and returns the image-resolution convolved padded image;
operated profiles are added unblurred as elsewhere. via_galaxies_from
switches on psf.convolve_over_sample_size > 1 and passes
image_is_convolved=True to the base simulator, trimming with
kernel_shape_image_resolution. s=1 path byte-identical.

Delta-kernel identity unit test for the padded-frame geometry.

Co-authored-by: Jammy2211 <JNightingale2211@gmail.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

Author: 3 people

autogalaxy/imaging/simulator.py

@@ -61,6 +61,27 @@ def via_galaxies_from(
                 psf=self.psf,
             )
 
+        if self.psf.convolve_over_sample_size > 1:
+
+            image = galaxies.convolved_padded_image_2d_from(
+                grid=grid, psf=self.psf, xp=xp
+            )
+
+            over_sample_size = grid.over_sample_size.resized_from(
+                new_shape=image.shape_native, mask_pad_value=1
+            )
+
+            dataset = self.via_image_from(
+                image=image,
+                over_sample_size=over_sample_size,
+                image_is_convolved=True,
+                xp=xp,
+            )
+
+            return dataset.trimmed_after_convolution_from(
+                kernel_shape=self.psf.kernel_shape_image_resolution
+            )
+
         image = galaxies.padded_image_2d_from(
             grid=grid, psf_shape_2d=self.psf.kernel.shape_native, xp=xp
         )

autogalaxy/operate/image.py

@@ -194,6 +194,70 @@ def padded_image_2d_from(self, grid, psf_shape_2d, xp=np):
 
         return self.image_2d_from(grid=padded_grid, xp=xp)
 
+    def convolved_padded_image_2d_from(self, grid, psf: aa.Convolver, xp=np):
+        """
+        Evaluate the light object's 2D image on a padded grid and convolve it with an
+        oversampled PSF at the fine resolution, returning the convolved padded image
+        at image resolution (still requiring trimming, as with
+        `padded_image_2d_from` + external convolution).
+
+        This is the simulation path for `convolve_over_sample_size > 1`: the padded
+        frame is sized by the kernel's image-resolution footprint and carries a
+        *uniform* over-sample size equal to the PSF's (the regular padded grid pads
+        its border with size-1 entries, which oversampled convolution correctly
+        rejects). The image is evaluated unbinned on the padded grid's over-sampled
+        coordinates and convolved by the oversampled Convolver, which bins back to
+        image resolution. No blurring image is needed — the padding guarantees all
+        flux that blurs into the frame is evaluated, exactly as in the existing
+        padded flow.
+
+        Parameters
+        ----------
+        grid
+            The 2D (y,x) coordinates of the grid the simulation is performed on, in
+            its original geometric reference frame.
+        psf
+            The oversampled PSF (`convolve_over_sample_size > 1`) the padded image
+            is convolved with.
+        """
+        s = psf.convolve_over_sample_size
+
+        kernel_shape_2d = psf.kernel_shape_image_resolution
+
+        padded_shape = (
+            grid.mask.shape_native[0] + kernel_shape_2d[0] - 1,
+            grid.mask.shape_native[1] + kernel_shape_2d[1] - 1,
+        )
+
+        padded_mask = aa.Mask2D.all_false(
+            shape_native=padded_shape,
+            pixel_scales=grid.mask.pixel_scales,
+            origin=grid.origin,
+        )
+
+        padded_grid = aa.Grid2D.from_mask(mask=padded_mask, over_sample_size=s)
+
+        image_over_sampled = self.image_2d_from(
+            grid=padded_grid.over_sampled, xp=xp, operated_only=False
+        )
+
+        convolved = psf.convolved_image_from(
+            image=image_over_sampled,
+            blurring_image=None,
+            mask=padded_mask,
+            xp=xp,
+        )
+
+        from autogalaxy.profiles.light.operated import LightProfileOperated
+
+        if self.has(cls=LightProfileOperated):
+            image_2d_operated = self.image_2d_from(
+                grid=padded_grid, xp=xp, operated_only=True
+            )
+            return convolved + image_2d_operated
+
+        return convolved
+
     def unmasked_blurred_image_2d_from(self, grid, psf):
         """
         Evaluate the light object's 2D image from a input 2D grid of coordinates and convolve it with a PSF, using a

test_autogalaxy/imaging/test_simulate_and_fit_imaging.py

@@ -490,3 +490,60 @@ def test__linear_light_profiles_agree_with_standard__galaxy_model_image_matches_
     assert fit_linear.galaxy_model_image_dict[galaxy_linear] == pytest.approx(
         galaxy_image.array, 1.0e-4
     )
+
+
+def test__perfect_fit__chi_squared_0__oversampled_psf():
+    # Simulate with an oversampled PSF (convolution at 2x the image resolution)
+    # via via_galaxies_from and fit the same galaxies at s=2: exact round trip.
+    s = 2
+    pixel_scales = 0.2
+
+    grid = ag.Grid2D.uniform(
+        shape_native=(21, 21), pixel_scales=pixel_scales, over_sample_size=s
+    )
+
+    psf = ag.Convolver.from_gaussian(
+        shape_native=(11, 11),
+        pixel_scales=pixel_scales / s,
+        sigma=0.15,
+        normalize=True,
+        convolve_over_sample_size=s,
+    )
+
+    galaxy_0 = ag.Galaxy(
+        redshift=0.5,
+        light=ag.lp.Sersic(centre=(0.0, 0.0), intensity=0.5, effective_radius=0.4),
+    )
+    galaxy_1 = ag.Galaxy(
+        redshift=0.5,
+        light=ag.lp.Exponential(centre=(0.05, 0.05), intensity=0.3, effective_radius=0.2),
+    )
+
+    simulator = ag.SimulatorImaging(
+        exposure_time=300.0, psf=psf, add_poisson_noise_to_data=False
+    )
+    dataset = simulator.via_galaxies_from(galaxies=[galaxy_0, galaxy_1], grid=grid)
+
+    dataset.noise_map = ag.Array2D.ones(
+        shape_native=dataset.data.shape_native, pixel_scales=pixel_scales
+    )
+
+    mask = ag.Mask2D.circular(
+        shape_native=dataset.data.shape_native, pixel_scales=pixel_scales, radius=1.8
+    )
+
+    masked = ag.Imaging(
+        data=dataset.data,
+        noise_map=dataset.noise_map,
+        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)
+
+    fit = ag.FitImaging(
+        dataset=masked, galaxies=ag.Galaxies(galaxies=[galaxy_0, galaxy_1])
+    )
+
+    assert fit.chi_squared == pytest.approx(0.0, abs=1.0e-10)

test_autogalaxy/operate/test_image.py

@@ -460,3 +460,44 @@ def test__blurred_image_2d_list_and_dict__oversampled_psf__match_scalar_path():
         assert np.array(blurred_dict[galaxy]) == pytest.approx(
             np.array(blurred_scalar), abs=1.0e-14
         )
+
+
+def test__convolved_padded_image_2d_from__delta_kernel__equals_binned_padded_image():
+    # With a delta fine kernel the fine convolution is the identity, so the
+    # convolved padded image must equal the binned padded evaluation — testing
+    # the padded-frame geometry and bin-down independently of PSF numerics.
+    import autoarray as aa
+
+    s = 2
+    pixel_scales = 1.0
+
+    grid = aa.Grid2D.uniform(
+        shape_native=(7, 7), pixel_scales=pixel_scales, over_sample_size=s
+    )
+
+    delta = np.zeros((5, 5))
+    delta[2, 2] = 1.0
+    psf = aa.Convolver(
+        kernel=aa.Array2D.no_mask(values=delta, pixel_scales=pixel_scales / s),
+        convolve_over_sample_size=s,
+    )
+
+    galaxy = ag.Galaxy(
+        redshift=0.5,
+        light=ag.lp.Sersic(centre=(0.2, -0.1), intensity=1.0, effective_radius=0.8),
+    )
+    galaxies = ag.Galaxies(galaxies=[galaxy])
+
+    convolved_padded = galaxies.convolved_padded_image_2d_from(grid=grid, psf=psf)
+
+    kernel_shape = psf.kernel_shape_image_resolution
+    padded_shape = (7 + kernel_shape[0] - 1, 7 + kernel_shape[1] - 1)
+    padded_mask = aa.Mask2D.all_false(
+        shape_native=padded_shape, pixel_scales=pixel_scales, origin=grid.origin
+    )
+    padded_grid = aa.Grid2D.from_mask(mask=padded_mask, over_sample_size=s)
+
+    image_sub = galaxy.image_2d_from(grid=padded_grid.over_sampled)
+    binned = np.array(image_sub).reshape(-1, s**2).mean(axis=1)
+
+    assert np.array(convolved_padded) == pytest.approx(binned, abs=1.0e-14)