diff --git a/autogalaxy/profiles/light/linear/abstract.py b/autogalaxy/profiles/light/linear/abstract.py index 5c4c0af9..d1e4088e 100644 --- a/autogalaxy/profiles/light/linear/abstract.py +++ b/autogalaxy/profiles/light/linear/abstract.py @@ -13,6 +13,7 @@ intensities to be penalized by a smoothness prior. """ +import functools import inspect import itertools import numpy as np @@ -30,6 +31,25 @@ from autogalaxy import exc +@functools.lru_cache(maxsize=1) +def _gaussian_image_2d_from(): + """ + The `Gaussian.image_2d_from` function object, whose body is exactly + `image_2d_via_radii_from(eccentric_radii_grid_from(grid))`. + + `LightProfileLinearObjFuncList._shared_eccentric_radii_index_groups` compares the profile class's + `image_2d_from` against this to decide whether an eccentric-radius grid can be shared across the list. Identity + of the function object is the exact test: every subclass which inherits this implementation (the linear, + operated and spherical Gaussians) matches, and every subclass which overrides it (e.g. `GaussianMultipole`) + does not. + + Imported lazily because `autogalaxy.profiles.light.standard` imports this module's base classes. + """ + from autogalaxy.profiles.light.standard.gaussian import Gaussian + + return Gaussian.image_2d_from + + class LightProfileLinear(LightProfile): """ A linear light profile, which is a light profile whose `intensity` value is solved for via linear algebra @@ -306,15 +326,165 @@ def mapping_matrix(self) -> np.ndarray: The `mapping_matrix` of the linear light profiles. """ - image_2d_list = [] + return self._xp.stack( + self._image_slim_list_from(grid=self.grid, xp=self._xp), axis=1 + ) - for pixel, light_profile in enumerate(self.light_profile_list): + @cached_property + def _shared_eccentric_radii_index_groups(self) -> List[List[int]]: + """ + The light profiles grouped by the geometry they share, as lists of indices into `light_profile_list`. + + Only groups of two or more `Gaussian` profiles with an identical `centre` and `ell_comps` are returned: + those are the groups which can share one reference-frame transform and one eccentric-radius grid. Every + other profile is left out of the groups and evaluated on its own. + + This is the multi-Gaussian expansion (MGE) case: a basis of tens of `Gaussian` profiles which differ only + in their `sigma`. Every one of them would otherwise translate and rotate the same grid into the same + reference frame and then form the same eccentric radii from it. The transform is by far the most + expensive part of evaluating a Gaussian -- it costs an `arctan2`, a `sin` and a `cos` per coordinate, + against the single `exp` of the profile itself -- so hoisting it out of the loop is most of the cost of + an MGE's mapping matrix. + + A basis is deliberately not required to be one group. The workspace's canonical MGE recipe stacks two + sets of 30 Gaussians which share a centre but carry their own `ell_comps`, and lands both sets in a + single `LightProfileLinearObjFuncList`; grouping serves that as two shared transforms rather than 60. + + A `Gaussian`'s `image_2d_from` is exactly `image_2d_via_radii_from(eccentric_radii_grid_from(grid))`, + which is what makes the two shared quantities well defined. The test is the identity of the + `image_2d_from` function object, not `isinstance`: every subclass which inherits that implementation (the + linear, operated and spherical Gaussians) shares the radii, and every subclass which overrides it (e.g. + `GaussianMultipole`, whose multipole perturbation is applied to the radius) does not. + + Other profiles are deliberately excluded even when they share a geometry. `Sersic` and its children + already avoid the polar transform via `_eccentric_radii_grid_from_cartesian`, which is only taken when + the grid handed to them is *not* pre-transformed -- so hoisting a transformed grid into them would be + both slower and a different floating-point expression. + + Returns + ------- + The index groups whose profiles share an eccentric-radius grid. + """ + index_group_dict: Dict[tuple, List[int]] = {} + + for index, light_profile in enumerate(self.light_profile_list): + if type(light_profile).image_2d_from is not _gaussian_image_2d_from(): + continue + + key = ( + type(light_profile), + tuple(light_profile.centre), + tuple(light_profile.ell_comps), + ) + + index_group_dict.setdefault(key, []).append(index) + + return [ + index_group + for index_group in index_group_dict.values() + if len(index_group) > 1 + ] + + def _image_slim_list_from(self, grid, xp=np) -> List[np.ndarray]: + """ + The `slim` image of every light profile in the list, evaluated on the input grid. + + This is the per-profile loop that builds the columns of the `mapping_matrix` (and, for imaging data, of + the blurring mapping matrix). Profiles which form an MGE basis (see + `_shared_eccentric_radii_index_groups`) are evaluated a group at a time by + `_shared_eccentric_radii_image_slim_list_from`, which hoists the work they have in common out of the + loop; every other profile is evaluated independently, exactly as before. + + Parameters + ---------- + grid + The 2D (y,x) coordinates the light profile images are evaluated on. + xp + The array module used (numpy, or `jax.numpy` for the JAX calculation). + + Returns + ------- + A list of the `slim` image of every light profile. + """ + light_profile_list = self.light_profile_list + + index_groups = ( + self._shared_eccentric_radii_index_groups + if xp is np and isinstance(grid, aa.Grid2D) + else [] + ) + + if not index_groups: + return [ + light_profile.image_2d_from(grid=grid, xp=xp).slim.array + for light_profile in light_profile_list + ] + + image_slim_list = [None] * len(light_profile_list) + + for index_group in index_groups: + image_slim_group = self._shared_eccentric_radii_image_slim_list_from( + grid=grid, + light_profile_list=[light_profile_list[i] for i in index_group], + ) + + for index, image_slim in zip(index_group, image_slim_group): + image_slim_list[index] = image_slim + + for index, light_profile in enumerate(light_profile_list): + if image_slim_list[index] is None: + image_slim_list[index] = light_profile.image_2d_from( + grid=grid, xp=np + ).slim.array + + return image_slim_list + + def _shared_eccentric_radii_image_slim_list_from( + self, grid, light_profile_list + ) -> List[np.ndarray]: + """ + The `slim` image of every light profile in one shared-geometry group, computing the two parts of the + calculation which are identical for all of them -- the reference-frame transform of the over-sampled + grid, and the eccentric radii formed from it -- exactly once. - image_2d = light_profile.image_2d_from(grid=self.grid, xp=self._xp).slim + Both are computed by calling the same methods on the first profile of the group, with the same inputs, + that every profile would otherwise call on itself, so each profile is then handed a bit-identical radius + grid and the images returned are bit-identical to the per-profile loop. - image_2d_list.append(image_2d.array) + Binning the over-sampled values back down stays a per-profile call, because + `OverSampler.binned_array_2d_from` takes a 1D array. Evaluating the `exp` stays a per-profile call too: + the arithmetic is transcendental-bound rather than overhead-bound, so stacking the profiles into one + (over-sampled pixels, profiles) block buys nothing and would have to be sliced column-wise to bin. + + Parameters + ---------- + grid + The 2D (y,x) coordinates the light profile images are evaluated on. + light_profile_list + One group of light profiles sharing a class, a `centre` and an `ell_comps`. + + Returns + ------- + A list of the `slim` image of every light profile in the group. + """ + light_profile_0 = light_profile_list[0] + + over_sampler = grid.over_sampler + + transformed_grid = light_profile_0.transformed_to_reference_frame_grid_from( + grid.over_sampled, np + ) + + eccentric_radii = light_profile_0.eccentric_radii_grid_from( + grid=transformed_grid, xp=np + ) - return self._xp.stack(image_2d_list, axis=1) + return [ + over_sampler.binned_array_2d_from( + array=light_profile.image_2d_via_radii_from(eccentric_radii, np), xp=np + ).slim.array + for light_profile in light_profile_list + ] @cached_property def operated_mapping_matrix_override(self) -> Optional[np.ndarray]: @@ -367,12 +537,7 @@ def operated_mapping_matrix_override(self) -> Optional[np.ndarray]: 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 - ], + self._image_slim_list_from(grid=self.blurring_grid, xp=np), axis=1, ) diff --git a/test_autogalaxy/profiles/light/linear/test_abstract.py b/test_autogalaxy/profiles/light/linear/test_abstract.py index 0a0b7151..f8c217ef 100644 --- a/test_autogalaxy/profiles/light/linear/test_abstract.py +++ b/test_autogalaxy/profiles/light/linear/test_abstract.py @@ -361,3 +361,249 @@ def test__operated_mapping_matrix_override__blurring_mask_ordering_matches_convo assert (y_st - (y_st[0] - y_up[0]) == y_up).all() assert (x_st - (x_st[0] - x_up[0]) == x_up).all() + + +def _shared_geometry_grids(): + """ + A masked `Grid2D` with a non-uniform over-sample size (so the binning takes the `bincount` branch) and its + blurring grid, used by the shared-geometry tests below. + """ + mask = aa.Mask2D.circular(shape_native=(25, 25), pixel_scales=0.1, radius=0.9) + + over_sample_size = aa.util.over_sample.over_sample_size_via_radial_bins_from( + grid=aa.Grid2D.from_mask(mask=mask), + sub_size_list=[4, 2, 1], + radial_list=[0.3, 0.6], + centre_list=[(0.0, 0.0)], + ) + + grid = aa.Grid2D.from_mask(mask=mask, over_sample_size=over_sample_size) + + assert not grid.over_sampler.sub_is_uniform + + 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) + + 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) + + return grid, blurring_grid, psf + + +def _per_profile_image_slim_list(light_profile_list, grid): + """ + The per-profile loop the shared-geometry fast path replaces: every profile evaluates the whole of + `image_2d_from`, including its own reference-frame transform and eccentric-radius grid. + """ + return [ + light_profile.image_2d_from(grid=grid, xp=np).slim.array + for light_profile in light_profile_list + ] + + +def test__mapping_matrix__shared_geometry_matches_per_profile_loop(): + # An MGE basis: every Gaussian shares a centre and ell_comps and differs only in sigma, so the + # reference-frame transform and the eccentric-radius grid are computed once and reused. The images + # must be bit-identical to evaluating every profile independently. + grid, blurring_grid, psf = _shared_geometry_grids() + + light_profile_list = [ + ag.lp_linear.Gaussian(centre=(0.1, -0.2), ell_comps=(0.2, 0.3), sigma=sigma) + for sigma in [0.05, 0.1, 0.2, 0.4, 0.8] + ] + + func_list = LightProfileLinearObjFuncList( + grid=grid, + blurring_grid=blurring_grid, + psf=psf, + light_profile_list=light_profile_list, + regularization=None, + ) + + assert func_list._shared_eccentric_radii_index_groups == [[0, 1, 2, 3, 4]] + + for grid_input in [grid, blurring_grid]: + assert np.array_equal( + np.stack(func_list._image_slim_list_from(grid=grid_input, xp=np), axis=1), + np.stack( + _per_profile_image_slim_list(light_profile_list, grid_input), axis=1 + ), + ) + + assert np.array_equal( + np.array(func_list.mapping_matrix), + np.stack(_per_profile_image_slim_list(light_profile_list, grid), axis=1), + ) + + # The override the imaging inversion actually consumes must equal the per-profile convolution. + + override = np.array(func_list.operated_mapping_matrix_override) + + for i, light_profile in enumerate(light_profile_list): + direct = psf.convolved_image_from( + image=light_profile.image_2d_from(grid=grid, xp=np), + blurring_image=light_profile.image_2d_from(grid=blurring_grid, xp=np), + xp=np, + ) + + assert override[:, i] == pytest.approx(np.array(direct), abs=1.0e-13) + + +def test__mapping_matrix__spherical_gaussians_share_radii_and_match_the_loop(): + # `GaussianSph` inherits `Gaussian.image_2d_from`, so it shares radii too -- via the spherical branch of + # `transformed_to_reference_frame_grid_from`, which translates without rotating. + grid, blurring_grid, psf = _shared_geometry_grids() + + light_profile_list = [ + ag.lp_linear.GaussianSph(centre=(0.1, -0.2), sigma=sigma) + for sigma in [0.05, 0.2, 0.8] + ] + + func_list = LightProfileLinearObjFuncList( + grid=grid, + blurring_grid=blurring_grid, + psf=psf, + light_profile_list=light_profile_list, + regularization=None, + ) + + assert func_list._shared_eccentric_radii_index_groups == [[0, 1, 2]] + + assert np.array_equal( + np.stack(func_list._image_slim_list_from(grid=grid, xp=np), axis=1), + np.stack(_per_profile_image_slim_list(light_profile_list, grid), axis=1), + ) + + +def test__mapping_matrix__mixed_geometry_falls_back(): + # Any mismatch -- a different centre, a different ell_comps, or a different profile class -- must take the + # per-profile loop, and match it exactly. + grid, blurring_grid, psf = _shared_geometry_grids() + + shared = ag.lp_linear.Gaussian(centre=(0.1, -0.2), ell_comps=(0.2, 0.3), sigma=0.1) + + mismatch_list = [ + [ + shared, + ag.lp_linear.Gaussian(centre=(0.3, -0.2), ell_comps=(0.2, 0.3), sigma=0.2), + ], + [ + shared, + ag.lp_linear.Gaussian(centre=(0.1, -0.2), ell_comps=(0.1, 0.3), sigma=0.2), + ], + [ + shared, + ag.lp_linear.GaussianSph(centre=(0.1, -0.2), sigma=0.2), + ], + ] + + for light_profile_list in mismatch_list: + func_list = LightProfileLinearObjFuncList( + grid=grid, + blurring_grid=blurring_grid, + psf=psf, + light_profile_list=light_profile_list, + regularization=None, + ) + + assert func_list._shared_eccentric_radii_index_groups == [] + + assert np.array_equal( + np.stack(func_list._image_slim_list_from(grid=grid, xp=np), axis=1), + np.stack(_per_profile_image_slim_list(light_profile_list, grid), axis=1), + ) + + +def test__mapping_matrix__non_gaussian_and_single_profile_take_the_loop(): + # Classes whose `image_2d_from` is not the Gaussian one must take the loop even when they share a geometry: + # `Sersic` reaches its eccentric radii from Cartesian coordinates when the grid is *not* pre-transformed, + # and `GaussianMultipole` perturbs the radius, so neither can be handed a shared radius grid. A one-profile + # list has nothing to share. + grid, blurring_grid, psf = _shared_geometry_grids() + + def func_list_from(light_profile_list): + return LightProfileLinearObjFuncList( + grid=grid, + blurring_grid=blurring_grid, + psf=psf, + light_profile_list=light_profile_list, + regularization=None, + ) + + single = func_list_from([ag.lp_linear.Gaussian(sigma=0.1)]) + + assert single._shared_eccentric_radii_index_groups == [] + + loop_list = [ + [ + ag.lp_linear.Sersic( + centre=(0.1, -0.2), + ell_comps=(0.2, 0.3), + effective_radius=effective_radius, + sersic_index=2.0, + ) + for effective_radius in [0.2, 0.5, 1.0] + ], + [ + ag.lp_linear.GaussianMultipole( + centre=(0.1, -0.2), + ell_comps=(0.2, 0.3), + sigma=sigma, + multipole_3_comps=(0.05, 0.02), + multipole_4_comps=(0.03, 0.01), + ) + for sigma in [0.1, 0.3] + ], + ] + + for light_profile_list in loop_list: + func_list = func_list_from(light_profile_list) + + assert func_list._shared_eccentric_radii_index_groups == [] + + assert np.array_equal( + np.stack(func_list._image_slim_list_from(grid=grid, xp=np), axis=1), + np.stack(_per_profile_image_slim_list(light_profile_list, grid), axis=1), + ) + + +def test__mapping_matrix__two_basis_mge_is_grouped_by_ell_comps(): + # The workspace's canonical MGE recipe stacks two sets of Gaussians which share a centre but carry their + # own ell_comps, and lands both sets in one func list. They must be served as two shared-geometry groups + # rather than falling back to 60 independent evaluations, and the Sersic mixed in with them must still take + # the per-profile loop. + grid, blurring_grid, psf = _shared_geometry_grids() + + light_profile_list = [ + ag.lp_linear.Gaussian(centre=(0.1, -0.2), ell_comps=(0.2, 0.3), sigma=0.05), + ag.lp_linear.Gaussian(centre=(0.1, -0.2), ell_comps=(-0.1, 0.0), sigma=0.05), + ag.lp_linear.Gaussian(centre=(0.1, -0.2), ell_comps=(0.2, 0.3), sigma=0.2), + ag.lp_linear.Sersic( + centre=(0.1, -0.2), ell_comps=(0.2, 0.3), effective_radius=0.5 + ), + ag.lp_linear.Gaussian(centre=(0.1, -0.2), ell_comps=(-0.1, 0.0), sigma=0.2), + ag.lp_linear.Gaussian(centre=(0.1, -0.2), ell_comps=(0.2, 0.3), sigma=0.8), + ] + + func_list = LightProfileLinearObjFuncList( + grid=grid, + blurring_grid=blurring_grid, + psf=psf, + light_profile_list=light_profile_list, + regularization=None, + ) + + assert func_list._shared_eccentric_radii_index_groups == [[0, 2, 5], [1, 4]] + + for grid_input in [grid, blurring_grid]: + assert np.array_equal( + np.stack(func_list._image_slim_list_from(grid=grid_input, xp=np), axis=1), + np.stack( + _per_profile_image_slim_list(light_profile_list, grid_input), axis=1 + ), + )