feat: k x s coupling call sites (#362 phase 2) (#486)

OperateImage._binned_for_convolution applies the partial pre-bin after
evaluation in the three blurred-image variants (no-op when sizes equal s);
the linear light profile override does the same;
convolved_padded_image_2d_from now inherits the input grid's (possibly
adaptive) evaluation sizes on the padded frame, padded with s in the
border, so simulation composes with adaptive evaluation too.

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

Author: 3 people

autogalaxy/operate/image.py

@@ -67,6 +67,26 @@ def has(self, cls) -> bool:
         """
         raise NotImplementedError
 
+    @staticmethod
+    def _binned_for_convolution(values, grid, psf, xp=np):
+        """
+        Partially bin values evaluated on ``grid``'s over-sampled coordinates (per-pixel
+        sizes k_i * s) down to the uniform s the oversampled Convolver requires — the
+        k x s coupling. A no-op when the grid's sizes already equal s.
+        """
+        from autoarray.operators.over_sampling.over_sample_util import (
+            binned_to_convolve_size_from,
+        )
+
+        values = values.array if hasattr(values, "array") else values
+
+        return binned_to_convolve_size_from(
+            values=values,
+            sub_size=np.array(grid.over_sample_size),
+            convolve_over_sample_size=psf.convolve_over_sample_size,
+            xp=xp,
+        )
+
     @staticmethod
     def _psf_evaluation_grids_from(grid, blurring_grid, psf):
         """
@@ -163,8 +183,12 @@ def blurred_image_2d_from(
 
         if convolution_mask is not None:
             blurred_image_2d = psf.convolved_image_from(
-                image=image_2d_not_operated,
-                blurring_image=blurring_image_2d_not_operated,
+                image=self._binned_for_convolution(
+                    image_2d_not_operated, grid, psf, xp=xp
+                ),
+                blurring_image=self._binned_for_convolution(
+                    blurring_image_2d_not_operated, blurring_grid, psf, xp=xp
+                ),
                 mask=convolution_mask,
                 xp=xp,
             )
@@ -247,14 +271,32 @@ def convolved_padded_image_2d_from(self, grid, psf: aa.Convolver, xp=np):
             origin=grid.origin,
         )
 
-        padded_grid = aa.Grid2D.from_mask(mask=padded_mask, over_sample_size=s)
+        # The padded grid inherits the input grid's (possibly adaptive, k x s
+        # coupled) evaluation sizes, padded with s in the border region.
+        pad_width = (
+            (padded_shape[0] - grid.mask.shape_native[0]) // 2,
+            (padded_shape[1] - grid.mask.shape_native[1]) // 2,
+        )
+        over_sample_size = np.pad(
+            np.array(grid.over_sample_size.native),
+            pad_width,
+            mode="constant",
+            constant_values=s,
+        )
+        over_sample_size[over_sample_size == 0] = s
+
+        padded_grid = aa.Grid2D.from_mask(
+            mask=padded_mask, over_sample_size=over_sample_size
+        )
 
         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,
+            image=self._binned_for_convolution(
+                image_over_sampled, padded_grid, psf, xp=xp
+            ),
             blurring_image=None,
             mask=padded_mask,
             xp=xp,
@@ -416,8 +458,12 @@ def blurred_image_2d_list_from(
 
             if convolution_mask is not None:
                 blurred_image_2d = psf.convolved_image_from(
-                    image=image_2d_not_operated,
-                    blurring_image=blurring_image_2d_not_operated,
+                    image=self._binned_for_convolution(
+                        image_2d_not_operated, grid, psf
+                    ),
+                    blurring_image=self._binned_for_convolution(
+                        blurring_image_2d_not_operated, blurring_grid, psf
+                    ),
                     mask=convolution_mask,
                 )
             else:
@@ -584,6 +630,14 @@ def galaxy_blurred_image_2d_dict_from(
                 galaxy_key
             ]
 
+            if convolution_mask is not None:
+                image_2d_not_operated = self._binned_for_convolution(
+                    image_2d_not_operated, grid, psf, xp=xp
+                )
+                blurring_image_2d_not_operated = self._binned_for_convolution(
+                    blurring_image_2d_not_operated, blurring_grid, psf, xp=xp
+                )
+
             blurred_image_2d = psf.convolved_image_from(
                 image=image_2d_not_operated,
                 blurring_image=blurring_image_2d_not_operated,

autogalaxy/profiles/light/linear/abstract.py

@@ -355,13 +355,23 @@ def operated_mapping_matrix_override(self) -> Optional[np.ndarray]:
 
         blurred_image_2d_list = []
 
+        from autogalaxy.operate.image import OperateImage
+
         for pixel, light_profile in enumerate(self.light_profile_list):
             image_2d = light_profile.image_2d_from(grid=evaluation_grid, xp=self._xp)
 
             blurring_image_2d = light_profile.image_2d_from(
                 grid=evaluation_blurring_grid, xp=self._xp
             )
 
+            if convolution_mask is not None:
+                image_2d = OperateImage._binned_for_convolution(
+                    image_2d, self.grid, self.psf, xp=self._xp
+                )
+                blurring_image_2d = OperateImage._binned_for_convolution(
+                    blurring_image_2d, self.blurring_grid, self.psf, xp=self._xp
+                )
+
             blurred_image_2d = self.psf.convolved_image_from(
                 image=image_2d,
                 blurring_image=blurring_image_2d,