Merge pull request #516 from PyAutoLabs/feature/numpy-deflections-p1

perf(decorators): to_grid propagates only an explicit over_sampled; Grid2D.over_sampled short-circuits at sub-size 1 (#514)

Author: James Nightingale

autoarray/structures/decorators/to_grid.py

@@ -8,23 +8,32 @@
 
 
 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 [
             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
         ]

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,

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

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)
+    )