feat: oversampled PSF blurred images via operate/image consumer switch (#480)

Implements phase 2c of the oversampled-PSF design (PyAutoArray#353 §5):
when psf.convolve_over_sample_size > 1, blurred_image_2d_from,
blurred_image_2d_list_from and galaxy_blurred_image_2d_dict_from evaluate
the image and blurring image on the grids' over-sampled coordinates
(grid.over_sampled is a Grid2DIrregular in per-pixel sub-block order, which
the over_sample decorator passes through unbinned — the oversampled
Convolver's input format) and convolve at the fine resolution.

No signature changes: the Grid2DIrregular pass-through replaces the
design's binned= kwarg threading, so Galaxy/Galaxies/Tracer image_2d_from
chains are untouched and Tracer inherits the behaviour via OperateImage.
padded/unmasked visualization images stay at image resolution (documented
limitation, revisited in the docs phase). s=1 paths byte-identical.

Tests: s=2 through operate/image.py equals the direct phase-2a Convolver
call to 1e-14; list/dict variants consistent with the scalar path.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: 2 people

autogalaxy/operate/image.py

@@ -112,6 +112,10 @@ def blurred_image_2d_from(
         The grid and blurring_grid must be a `Grid2D` objects so the evaluated image can be mapped to a uniform 2D
         array and binned up for convolution. They therefore cannot be `Grid2DIrregular` objects.
 
+        If the PSF is oversampled (`convolve_over_sample_size > 1`), the images are instead evaluated on the
+        grids' over-sampled coordinates (in per-pixel sub-block order, without binning) and convolved at the
+        fine resolution, with the Convolver binning the result back to image resolution.
+
         Parameters
         ----------
         grid
@@ -125,19 +129,40 @@ def blurred_image_2d_from(
             LightProfileOperated,
         )
 
-        image_2d_not_operated = self.image_2d_from(
-            grid=grid, xp=xp, operated_only=False
-        )
-        blurring_image_2d_not_operated = self.image_2d_from(
-            grid=blurring_grid, xp=xp, operated_only=False
-        )
+        if psf.convolve_over_sample_size > 1:
 
-        blurred_image_2d = self._blurred_image_2d_from(
-            image_2d=image_2d_not_operated,
-            blurring_image_2d=blurring_image_2d_not_operated,
-            psf=psf,
-            xp=xp,
-        )
+            # grid.over_sampled is a Grid2DIrregular in per-pixel sub-block order, which the
+            # over_sample decorator passes through unbinned — the oversampled Convolver's
+            # input format.
+            image_2d_not_operated = self.image_2d_from(
+                grid=grid.over_sampled, xp=xp, operated_only=False
+            )
+            blurring_image_2d_not_operated = self.image_2d_from(
+                grid=blurring_grid.over_sampled, xp=xp, operated_only=False
+            )
+
+            blurred_image_2d = psf.convolved_image_from(
+                image=image_2d_not_operated,
+                blurring_image=blurring_image_2d_not_operated,
+                mask=grid.mask,
+                xp=xp,
+            )
+
+        else:
+
+            image_2d_not_operated = self.image_2d_from(
+                grid=grid, xp=xp, operated_only=False
+            )
+            blurring_image_2d_not_operated = self.image_2d_from(
+                grid=blurring_grid, xp=xp, operated_only=False
+            )
+
+            blurred_image_2d = self._blurred_image_2d_from(
+                image_2d=image_2d_not_operated,
+                blurring_image_2d=blurring_image_2d_not_operated,
+                psf=psf,
+                xp=xp,
+            )
 
         if self.has(cls=LightProfileOperated):
             image_2d_operated = self.image_2d_from(grid=grid, xp=xp, operated_only=True)
@@ -294,11 +319,18 @@ def blurred_image_2d_list_from(
 
         image_2d_operated_list = self.image_2d_list_from(grid=grid, operated_only=True)
 
+        if psf.convolve_over_sample_size > 1:
+            evaluation_grid = grid.over_sampled
+            evaluation_blurring_grid = blurring_grid.over_sampled
+        else:
+            evaluation_grid = grid
+            evaluation_blurring_grid = blurring_grid
+
         image_2d_not_operated_list = self.image_2d_list_from(
-            grid=grid, operated_only=False
+            grid=evaluation_grid, operated_only=False
         )
         blurring_image_2d_not_operated_list = self.image_2d_list_from(
-            grid=blurring_grid, operated_only=False
+            grid=evaluation_blurring_grid, operated_only=False
         )
 
         blurred_image_2d_list = []
@@ -307,11 +339,18 @@ def blurred_image_2d_list_from(
             image_2d_not_operated = image_2d_not_operated_list[i]
             blurring_image_2d_not_operated = blurring_image_2d_not_operated_list[i]
 
-            blurred_image_2d = self._blurred_image_2d_from(
-                image_2d=image_2d_not_operated,
-                blurring_image_2d=blurring_image_2d_not_operated,
-                psf=psf,
-            )
+            if psf.convolve_over_sample_size > 1:
+                blurred_image_2d = psf.convolved_image_from(
+                    image=image_2d_not_operated,
+                    blurring_image=blurring_image_2d_not_operated,
+                    mask=grid.mask,
+                )
+            else:
+                blurred_image_2d = self._blurred_image_2d_from(
+                    image_2d=image_2d_not_operated,
+                    blurring_image_2d=blurring_image_2d_not_operated,
+                    psf=psf,
+                )
 
             image_2d_operated = image_2d_operated_list[i]
 
@@ -444,12 +483,19 @@ def galaxy_blurred_image_2d_dict_from(
             The 2D (y,x) coordinates neighboring the (masked) grid whose light is blurred into the image.
         """
 
+        if psf.convolve_over_sample_size > 1:
+            evaluation_grid = grid.over_sampled
+            evaluation_blurring_grid = blurring_grid.over_sampled
+        else:
+            evaluation_grid = grid
+            evaluation_blurring_grid = blurring_grid
+
         galaxy_image_2d_not_operated_dict = self.galaxy_image_2d_dict_from(
-            grid=grid, operated_only=False, xp=xp
+            grid=evaluation_grid, operated_only=False, xp=xp
         )
 
         galaxy_blurring_image_2d_not_operated_dict = self.galaxy_image_2d_dict_from(
-            grid=blurring_grid, operated_only=False, xp=xp
+            grid=evaluation_blurring_grid, operated_only=False, xp=xp
         )
 
         galaxy_image_2d_operated_dict = self.galaxy_image_2d_dict_from(
@@ -467,6 +513,7 @@ def galaxy_blurred_image_2d_dict_from(
             blurred_image_2d = psf.convolved_image_from(
                 image=image_2d_not_operated,
                 blurring_image=blurring_image_2d_not_operated,
+                mask=grid.mask if psf.convolve_over_sample_size > 1 else None,
                 xp=xp,
             )
 

test_autogalaxy/operate/test_image.py

@@ -374,3 +374,89 @@ def test__galaxy_visibilities_dict_from(grid_2d_7x7, transformer_7x7_7):
     assert (visibilities_dict[g0] == visibilities_list[1]).all()
     assert (visibilities_dict[g1] == visibilities_list[0]).all()
     assert (visibilities_dict[g2] == visibilities_list[2]).all()
+
+
+def _oversampled_psf_scene():
+    import autoarray as aa
+
+    mask = aa.Mask2D.circular(shape_native=(11, 11), pixel_scales=1.0, radius=3.5)
+
+    s = 2
+    n = 9
+    c = (np.arange(n) - (n - 1) / 2.0) * (1.0 / s)
+    yy, xx = np.meshgrid(-c, c, indexing="ij")
+    kernel = np.exp(-0.5 * (yy**2 + xx**2) / 0.8**2)
+    kernel = aa.Array2D.no_mask(values=kernel / kernel.sum(), pixel_scales=1.0 / s)
+    psf = aa.Convolver(kernel=kernel, convolve_over_sample_size=s)
+
+    grid = aa.Grid2D.from_mask(mask=mask, over_sample_size=s)
+    blurring_mask = mask.derive_mask.blurring_from(
+        kernel_shape_native=psf.kernel_shape_image_resolution, allow_padding=True
+    )
+    blurring_grid = aa.Grid2D.from_mask(mask=blurring_mask, over_sample_size=s)
+
+    return mask, psf, grid, blurring_grid
+
+
+def test__blurred_image_2d_from__oversampled_psf__matches_direct_convolver_call():
+    # The consumer switch evaluates on the over-sampled grids and hands the
+    # sub-block ordered values to the oversampled Convolver — the result must
+    # equal calling that (phase-2a tested) API directly.
+    mask, psf, grid, blurring_grid = _oversampled_psf_scene()
+
+    galaxy = ag.Galaxy(
+        redshift=0.5,
+        light=ag.lp.Sersic(
+            centre=(0.3, -0.4), intensity=1.0, effective_radius=1.0, sersic_index=2.0
+        ),
+    )
+    galaxies = ag.Galaxies(galaxies=[galaxy])
+
+    blurred = galaxies.blurred_image_2d_from(
+        grid=grid, blurring_grid=blurring_grid, psf=psf
+    )
+
+    image_sub = galaxy.image_2d_from(grid=grid.over_sampled)
+    blurring_sub = galaxy.image_2d_from(grid=blurring_grid.over_sampled)
+    direct = psf.convolved_image_from(
+        image=image_sub, blurring_image=blurring_sub, mask=mask
+    )
+
+    assert np.array(blurred) == pytest.approx(np.array(direct), abs=1.0e-14)
+
+    # The result is at image resolution and differs from the binned-then-convolved
+    # (s=1 style) computation — the oversampling is actually doing something.
+    assert np.array(blurred).shape == (mask.pixels_in_mask,)
+
+
+def test__blurred_image_2d_list_and_dict__oversampled_psf__match_scalar_path():
+    mask, psf, grid, blurring_grid = _oversampled_psf_scene()
+
+    galaxy_0 = ag.Galaxy(
+        redshift=0.5,
+        light=ag.lp.Sersic(centre=(0.3, -0.4), intensity=1.0, effective_radius=1.0),
+    )
+    galaxy_1 = ag.Galaxy(
+        redshift=0.5,
+        light=ag.lp.Gaussian(centre=(-0.5, 0.2), intensity=2.0, sigma=0.7),
+    )
+    galaxies = ag.Galaxies(galaxies=[galaxy_0, galaxy_1])
+
+    blurred_list = galaxies.blurred_image_2d_list_from(
+        grid=grid, blurring_grid=blurring_grid, psf=psf
+    )
+    blurred_dict = galaxies.galaxy_blurred_image_2d_dict_from(
+        grid=grid, blurring_grid=blurring_grid, psf=psf
+    )
+
+    total_from_list = np.sum([np.array(b) for b in blurred_list], axis=0)
+    blurred_total = galaxies.blurred_image_2d_from(
+        grid=grid, blurring_grid=blurring_grid, psf=psf
+    )
+
+    assert total_from_list == pytest.approx(np.array(blurred_total), abs=1.0e-12)
+
+    for galaxy, blurred_scalar in zip([galaxy_0, galaxy_1], blurred_list):
+        assert np.array(blurred_dict[galaxy]) == pytest.approx(
+            np.array(blurred_scalar), abs=1.0e-14
+        )