From 92bb9b2ce29968d0261f611f06cbec7d6d168c51 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 2 Sep 2026 13:20:22 -0400 Subject: [PATCH] perf(decorators): stop to_grid materialising the over sampled grid (#514) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `GridMaker.via_grid_2d` read `result.over_sampled` / `result.over_sampler` through the public properties. When the wrapped result is already a `Grid2D` — which it is for every spherical mass profile, whose `@to_grid` decorated `transformed_to_reference_frame_grid_from` delegates to another `@to_grid` method — that read *materialises* the over sampled grid, running a per-pixel Python loop over the whole mask. On a 15k-pixel HST grid that is ~0.5-1.2 s per deflection-angle call against ~1 ms of actual profile maths, and the value built is mask-derived, so for a translated/rotated grid it is wrong as well as expensive. Nothing reads it. Both call sites now read the private `_over_sampled` / `_over_sampler`, so only a value a caller explicitly passed in propagates (the load-bearing case: `Galaxy.traced_grid_2d_from` and `Tracer.traced_grid_2d_list_from`, which both construct `Grid2D(..., over_sampled=...)`). Anything else is left as `None` for the new grid to compute lazily, if it is ever asked for. `Grid2D.over_sampled` additionally short circuits at a uniform sub size of 1, where every sub pixel is the pixel itself and the over sampled grid is just the slim grid in the same order — equal to the loop's output bit for bit at the default origin, and within 1 ULP of it otherwise. The values are copied so an in-place edit of `over_sampled` cannot write through to the grid. Measured on hst (15361 image pixels, OMP_NUM_THREADS=1), `Grid2D` s/call: IsothermalSph 699.2 ms -> 0.92 ms PowerLawSph 701.0 ms -> 1.37 ms NFWSph 590.9 ms -> 5.79 ms gNFWSph 1.20 s -> 348.6 ms (the remainder is real quadrature) All pinned deflection values still PASS at rtol 1e-6. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01HWjPT94MPbEHT45kJmDpDh --- autoarray/structures/decorators/to_grid.py | 17 +++-- autoarray/structures/grids/uniform_2d.py | 17 +++++ .../structures/decorators/test_to_grid.py | 69 +++++++++++++++++++ .../structures/grids/test_uniform_2d.py | 48 +++++++++++++ 4 files changed, 147 insertions(+), 4 deletions(-) diff --git a/autoarray/structures/decorators/to_grid.py b/autoarray/structures/decorators/to_grid.py index 40c979e48..1b0be8fe9 100644 --- a/autoarray/structures/decorators/to_grid.py +++ b/autoarray/structures/decorators/to_grid.py @@ -8,14 +8,23 @@ class GridMaker(AbstractMaker): + # The private ``_over_sampled`` / ``_over_sampler`` attributes are read here + # instead of the public properties on purpose: when ``result`` is already a + # ``Grid2D`` (e.g. the chained ``@to_grid`` methods of the spherical mass + # profiles) touching the properties *materialises* them, and building the + # over sampled grid runs a per-pixel Python loop that dominates the runtime + # of an otherwise millisecond deflection calculation. Only a value that was + # explicitly passed in by the caller should propagate to the new grid; a + # ``None`` lets the new ``Grid2D`` compute its own lazily, if anything ever + # asks for it. def via_grid_2d(self, result) -> Union[Grid2D, List[Grid2D]]: if not isinstance(result, list): return Grid2D( values=result, mask=self.mask, over_sample_size=self.over_sample_size, - over_sampled=getattr(result, "over_sampled", None), - over_sampler=getattr(result, "over_sampler", None), + over_sampled=getattr(result, "_over_sampled", None), + over_sampler=getattr(result, "_over_sampler", None), ) return [ @@ -23,8 +32,8 @@ def via_grid_2d(self, result) -> Union[Grid2D, List[Grid2D]]: values=res, mask=self.mask, over_sample_size=self.over_sample_size, - over_sampled=getattr(res, "over_sampled", None), - over_sampler=getattr(res, "over_sampler", None), + over_sampled=getattr(res, "_over_sampled", None), + over_sampler=getattr(res, "_over_sampler", None), ) for res in result ] diff --git a/autoarray/structures/grids/uniform_2d.py b/autoarray/structures/grids/uniform_2d.py index b29d3ce7d..bb7132ecb 100644 --- a/autoarray/structures/grids/uniform_2d.py +++ b/autoarray/structures/grids/uniform_2d.py @@ -213,6 +213,23 @@ def over_sampled(self): if self._over_sampled is not None: return self._over_sampled + if np.all(np.asarray(self.over_sample_size.array) == 1): + # At a uniform sub size of 1 every sub pixel is the pixel itself, so the + # over sampled grid is the slim grid in the same order. Short circuiting + # here skips the per-pixel Python loop below, which dominates the runtime + # of grid calculations that never touch the over sampled grid. The values + # are copied so that in-place edits of `over_sampled` do not write through + # to the grid itself, matching the behaviour of the loop below. + grid_slim = grid_2d_util.convert_grid_2d_to_slim( + grid_2d=self.array, mask_2d=self.mask, xp=self._xp + ) + + self._over_sampled = Grid2DIrregular( + values=self._xp.array(grid_slim), xp=self._xp + ) + + return self._over_sampled + over_sampled = over_sample_util.grid_2d_slim_over_sampled_via_mask_from( mask_2d=np.array(self.mask), pixel_scales=self.mask.pixel_scales, diff --git a/test_autoarray/structures/decorators/test_to_grid.py b/test_autoarray/structures/decorators/test_to_grid.py index 00a0089ec..a5599de56 100644 --- a/test_autoarray/structures/decorators/test_to_grid.py +++ b/test_autoarray/structures/decorators/test_to_grid.py @@ -110,3 +110,72 @@ def test__in_ndarray__out_ndarray(): assert isinstance(result, np.ndarray) assert not isinstance(result, aa.Grid2D) assert not isinstance(result, aa.Grid2DIrregular) + + +class MockGrid2DPassThroughObj: + """ + Mimics the profile classes in **PyAutoGalaxy** whose `@to_grid` decorated methods + delegate to another `@to_grid` decorated method, such that the value the decorator + wraps up is already a `Grid2D`. + """ + + def __init__(self, over_sampled=None): + self.centre = (0.0, 0.0) + self.over_sampled = over_sampled + + @aa.decorators.to_grid + def grid_2d_from(self, grid, *args, **kwargs): + return aa.Grid2D( + values=np.multiply(2.0, grid.array), + mask=grid.mask, + over_sample_size=grid.over_sample_size, + over_sampled=self.over_sampled, + ) + + +def _mask_2x2(): + return aa.Mask2D( + mask=[ + [True, True, True, True], + [True, False, False, True], + [True, False, False, True], + [True, True, True, True], + ], + pixel_scales=(1.0, 1.0), + ) + + +def test__to_grid__does_not_materialise_over_sampled_of_wrapped_grid(monkeypatch): + """ + Reading the public `over_sampled` / `over_sampler` properties inside `to_grid` + materialises them via a per-pixel Python loop, which dominates the runtime of + deflection angle calculations. Only an explicitly set value may propagate. + """ + + def _fail(self): + pytest.fail("over_sampled materialised inside to_grid") + + def _fail_sampler(self): + pytest.fail("over_sampler materialised inside to_grid") + + monkeypatch.setattr(aa.Grid2D, "over_sampled", property(_fail)) + monkeypatch.setattr(aa.Grid2D, "over_sampler", property(_fail_sampler)) + + grid_2d = aa.Grid2D.from_mask(mask=_mask_2x2(), over_sample_size=2) + + result = MockGrid2DPassThroughObj().grid_2d_from(grid=grid_2d) + + assert isinstance(result, aa.Grid2D) + assert result._over_sampled is None + assert result._over_sampler is None + + +def test__to_grid__propagates_explicitly_set_over_sampled(): + grid_2d = aa.Grid2D.from_mask(mask=_mask_2x2(), over_sample_size=2) + + sentinel = aa.Grid2DIrregular(values=[(1.0, 2.0), (3.0, 4.0)]) + + result = MockGrid2DPassThroughObj(over_sampled=sentinel).grid_2d_from(grid=grid_2d) + + assert result._over_sampled is sentinel + assert result.over_sampled is sentinel diff --git a/test_autoarray/structures/grids/test_uniform_2d.py b/test_autoarray/structures/grids/test_uniform_2d.py index 5ff4f23e5..25c3dacaf 100644 --- a/test_autoarray/structures/grids/test_uniform_2d.py +++ b/test_autoarray/structures/grids/test_uniform_2d.py @@ -880,3 +880,51 @@ def test__subtracted_and_rotated_from__shift_first_then_rotate(): rotated = grid.subtracted_and_rotated_from(offset=(1.0, 2.0), angle=90.0) assert rotated.array == pytest.approx(expected, 1.0e-4) + + +def test__over_sampled__sub_size_1_is_the_slim_grid(): + mask = aa.Mask2D( + mask=[ + [True, True, True, True, True], + [True, False, False, False, True], + [True, False, True, False, True], + [True, False, False, False, True], + [True, True, True, True, True], + ], + pixel_scales=(1.5, 1.5), + ) + + grid = aa.Grid2D.from_mask(mask=mask, over_sample_size=1) + + over_sampled = grid.over_sampled + + assert isinstance(over_sampled, aa.Grid2DIrregular) + assert np.array_equal(np.array(over_sampled), np.array(grid.array)) + + # The cached value is not a view onto the grid itself. + assert not np.shares_memory(np.array(over_sampled.array), np.array(grid.array)) + + # The property caches, so a second read is the same object. + assert grid.over_sampled is over_sampled + + +def test__over_sampled__sub_size_2_over_samples_every_pixel(): + mask = aa.Mask2D( + mask=[ + [True, True, True, True, True], + [True, False, False, False, True], + [True, False, True, False, True], + [True, False, False, False, True], + [True, True, True, True, True], + ], + pixel_scales=(1.5, 1.5), + ) + + grid = aa.Grid2D.from_mask(mask=mask, over_sample_size=2) + + over_sampled = grid.over_sampled + + assert np.array(over_sampled).shape[0] == 4 * grid.array.shape[0] + assert not np.array_equal( + np.array(over_sampled)[: grid.array.shape[0]], np.array(grid.array) + )