diff --git a/autogalaxy/profiles/light/linear/abstract.py b/autogalaxy/profiles/light/linear/abstract.py index ddc2428b..5c4c0af9 100644 --- a/autogalaxy/profiles/light/linear/abstract.py +++ b/autogalaxy/profiles/light/linear/abstract.py @@ -346,6 +346,43 @@ def operated_mapping_matrix_override(self) -> Optional[np.ndarray]: if isinstance(self.light_profile_list[0], LightProfileOperated): return self.mapping_matrix + if ( + self._xp is np + and self.psf.convolve_over_sample_size == 1 + and self.blurring_grid is not None + ): + # Numpy fast path: convolve every light profile in one batched call + # instead of looping `convolved_image_from` once per profile. For an + # MGE of N Gaussians the loop rebuilds the `ConvolverState` and runs a + # separate 2D convolution N times; here the images are stacked into a + # single (pixels, N) mapping matrix and convolved as one (ny, nx, N) + # cube against the kernel. + # + # This is numerically the same operation: for `xp is np` and + # `convolve_over_sample_size == 1`, `convolved_image_from` itself + # dispatches to `convolved_image_via_real_space_np_from`, which + # scatters the image and blurring image into the very same + # `ConvolverState` frame (`state.mask` / `state.blurring_mask`) that + # the batched call below uses. + mapping_matrix = self.mapping_matrix + + blurring_mapping_matrix = np.stack( + [ + light_profile.image_2d_from( + grid=self.blurring_grid, xp=np + ).slim.array + for light_profile in self.light_profile_list + ], + axis=1, + ) + + return self.psf.convolved_mapping_matrix_via_real_space_np_from( + mapping_matrix=mapping_matrix, + mask=self.grid.mask, + blurring_mapping_matrix=blurring_mapping_matrix, + blurring_mask=self.blurring_grid.mask, + ) + if self.psf.convolve_over_sample_size > 1: # Retain the parent Grid2D while requesting unbinned values. A # discrete profile needs the parent-pixel geometry to place its diff --git a/test_autogalaxy/profiles/light/linear/test_abstract.py b/test_autogalaxy/profiles/light/linear/test_abstract.py index ff57f997..0a0b7151 100644 --- a/test_autogalaxy/profiles/light/linear/test_abstract.py +++ b/test_autogalaxy/profiles/light/linear/test_abstract.py @@ -278,3 +278,86 @@ def test__operated_mapping_matrix_override__oversampled_psf__matches_direct_conv image=image_sub, blurring_image=blurring_sub, mask=mask ) assert override[:, i] == pytest.approx(np.array(direct), abs=1.0e-14) + + +def test__operated_mapping_matrix_override__batched_numpy_path__matches_per_profile_convolution(): + # The numpy fast path convolves every linear light profile in one batched call + # instead of looping `psf.convolved_image_from` once per profile. Each column + # must still equal the per-profile convolution exactly, including the flux + # blurred in from outside the mask (a bright Gaussian sits at the mask edge) + # and for a non-symmetric kernel. + mask = aa.Mask2D.circular(shape_native=(25, 25), pixel_scales=0.1, radius=0.9) + + kernel_native = np.random.default_rng(7).random((5, 7)) + 0.05 + kernel = aa.Array2D.no_mask( + values=kernel_native / kernel_native.sum(), pixel_scales=0.1 + ) + psf = aa.Convolver(kernel=kernel) + + assert psf.convolve_over_sample_size == 1 + + grid = aa.Grid2D.from_mask(mask=mask) + 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) + + light_profile_list = [ + # Bright and narrow, sat on the mask edge so its flux is dominated by the + # blurring region -- this is what breaks if the blurring mapping matrix is + # scattered in the wrong order. + ag.lp_linear.Gaussian(centre=(0.85, 0.0), sigma=0.05), + ag.lp_linear.Gaussian(centre=(-0.8, 0.6), sigma=0.08), + ag.lp_linear.Sersic(centre=(0.1, -0.2), effective_radius=0.4, sersic_index=3.0), + ag.lp_linear.Gaussian(centre=(0.0, 0.0), sigma=0.3), + ] + + func_list = LightProfileLinearObjFuncList( + grid=grid, + blurring_grid=blurring_grid, + psf=psf, + light_profile_list=light_profile_list, + regularization=None, + ) + + override = np.array(func_list.operated_mapping_matrix_override) + + assert override.shape == (mask.pixels_in_mask, len(light_profile_list)) + + for i, light_profile in enumerate(light_profile_list): + image = light_profile.image_2d_from(grid=grid, xp=np) + blurring_image = light_profile.image_2d_from(grid=blurring_grid, xp=np) + + direct = psf.convolved_image_from( + image=image, blurring_image=blurring_image, xp=np + ) + + assert override[:, i] == pytest.approx(np.array(direct), abs=1.0e-13) + + +def test__operated_mapping_matrix_override__blurring_mask_ordering_matches_convolver_state(): + # The batched call scatters the blurring mapping matrix using the blurring mask + # derived inside the `ConvolverState` (from the resized FFT-frame mask), whereas + # the blurring grid is built upstream from the unresized mask. The two masks live + # on different frames, so this asserts their slim orderings are the same + # permutation of the same pixels (a pure translation of the native indices). + mask = aa.Mask2D.circular(shape_native=(21, 21), pixel_scales=0.1, radius=0.7) + + kernel_native = np.random.default_rng(3).random((5, 7)) + 0.05 + kernel = aa.Array2D.no_mask( + values=kernel_native / kernel_native.sum(), pixel_scales=0.1 + ) + psf = aa.Convolver(kernel=kernel) + + upstream_blurring_mask = mask.derive_mask.blurring_from( + kernel_shape_native=psf.kernel_shape_image_resolution, allow_padding=True + ) + state_blurring_mask = psf.state_from(mask=mask).blurring_mask + + y_up, x_up = (np.asarray(a) for a in upstream_blurring_mask.slim_to_native_tuple) + y_st, x_st = (np.asarray(a) for a in state_blurring_mask.slim_to_native_tuple) + + assert y_up.shape == y_st.shape + + assert (y_st - (y_st[0] - y_up[0]) == y_up).all() + assert (x_st - (x_st[0] - x_up[0]) == x_up).all()