perf(galaxy): skip the second deflection call for the over-sampled grid at sub-size 1 (PyAutoArray#514)

autogalaxy/galaxy/galaxy.py

@@ -374,12 +374,29 @@ def traced_grid_2d_from(
             The source-plane (y, x) coordinates after deflection.
         """
         if isinstance(grid, aa.Grid2D):
+            traced_grid = grid - self.deflections_yx_2d_from(grid=grid, xp=xp)
+
+            # At a uniform over sample size of 1 every sub pixel is the pixel itself, so
+            # `grid.over_sampled` is the slim grid in the same order and tracing it again
+            # doubles the deflection angle calculation for a bit-identical result. The
+            # over sample size is host numpy, so this branch is static at JAX trace time.
+
+            sub_size_all_one = bool(np.all(np.asarray(grid.over_sample_size.array) == 1))
+
+            if sub_size_all_one:
+                traced_grid_over_sampled = aa.Grid2DIrregular(
+                    values=traced_grid.array, xp=xp
+                )
+            else:
+                traced_grid_over_sampled = grid.over_sampled - self.deflections_yx_2d_from(
+                    grid=grid.over_sampled, xp=xp
+                )
+
             return aa.Grid2D(
-                values=grid - self.deflections_yx_2d_from(grid=grid, xp=xp),
+                values=traced_grid,
                 mask=grid.mask,
                 over_sample_size=grid.over_sample_size,
-                over_sampled=grid.over_sampled
-                - self.deflections_yx_2d_from(grid=grid.over_sampled, xp=xp),
+                over_sampled=traced_grid_over_sampled,
                 over_sampler=grid.over_sampler,
             )
 

test_autogalaxy/galaxy/test_galaxy.py

@@ -192,6 +192,57 @@ def test__deflections_yx_2d_from__two_mass_profiles__matches_sum_of_individual_d
     assert gal_deflections[1, 0] == mp_deflections[1, 0]
 
 
+@pytest.mark.parametrize("over_sample_size", [1, 4])
+def test__traced_grid_2d_from__over_sampled_grid_is_the_traced_over_sampled_grid(
+    over_sample_size,
+):
+    """
+    The `over_sampled` attribute of the traced grid must be the galaxy's deflections applied to the
+    input grid's `over_sampled` grid. At a uniform over sample size of 1 the second ray-trace is
+    skipped as an optimization, so this is pinned for a sub size of 1 and a sub size above 1.
+    """
+    galaxy = ag.Galaxy(
+        redshift=0.5, mass_profile=ag.mp.IsothermalSph(einstein_radius=1.0)
+    )
+
+    grid = ag.Grid2D.uniform(
+        shape_native=(7, 7), pixel_scales=0.1, over_sample_size=over_sample_size
+    )
+
+    traced_grid = galaxy.traced_grid_2d_from(grid=grid)
+
+    traced_over_sampled = galaxy.traced_grid_2d_from(grid=grid.over_sampled)
+
+    assert np.allclose(
+        np.asarray(traced_grid.over_sampled.array),
+        np.asarray(traced_over_sampled.array),
+        rtol=0.0,
+        atol=1.0e-12,
+    )
+
+    assert not np.allclose(
+        np.asarray(traced_grid.over_sampled.array), np.asarray(grid.over_sampled.array)
+    )
+
+
+def test__traced_grid_2d_from__over_sample_size_1__over_sampled_grid_equals_slim_grid():
+    """
+    At a uniform over sample size of 1 every sub pixel is the pixel itself, so the traced over
+    sampled grid is bit identical to the traced slim grid.
+    """
+    galaxy = ag.Galaxy(
+        redshift=0.5, mass_profile=ag.mp.IsothermalSph(einstein_radius=1.0)
+    )
+
+    grid = ag.Grid2D.uniform(shape_native=(7, 7), pixel_scales=0.1, over_sample_size=1)
+
+    traced_grid = galaxy.traced_grid_2d_from(grid=grid)
+
+    assert np.array_equal(
+        np.asarray(traced_grid.over_sampled.array), np.asarray(traced_grid.array)
+    )
+
+
 def test__no_mass_profile__potential_2d_from__returns_zeros_of_grid_shape(grid_2d_7x7):
     galaxy = ag.Galaxy(redshift=0.5)