Merge pull request #358 from PyAutoLabs/feature/lens-calc-hessian-richardson

fix: Richardson-extrapolate LensCalc numpy Hessian for JAX-path parity

Author: Jammy2211

autogalaxy/operate/lens_calc.py

@@ -359,9 +359,11 @@ def hessian_from(self, grid, xp=np) -> Tuple:
 
         Two computational paths are available, selected via the `xp` parameter:
 
-        - **NumPy** (``xp=np``, default): finite-difference approximation. Deflection angles are
-          evaluated at four shifted positions around each grid coordinate (±y, ±x) and the
-          central difference is taken. JAX is not imported.
+        - **NumPy** (``xp=np``, default): 2-point central finite-difference approximation,
+          Richardson-extrapolated at step sizes ``h`` and ``h/2`` and combined as
+          ``(4 * H(h/2) - H(h)) / 3``. This cancels the leading ``O(h^2)`` truncation term,
+          giving ``O(h^4)`` accuracy and matching the JAX path to float64 precision. JAX is
+          not imported.
 
         - **JAX** (``xp=jnp``): exact derivatives via ``jax.jacfwd`` applied to
           ``deflections_yx_scalar``, vectorised over the grid with ``jnp.vectorize``.
@@ -377,9 +379,22 @@ def hessian_from(self, grid, xp=np) -> Tuple:
             used and the type of the returned arrays.
         """
         if xp is np:
-            return self._hessian_via_finite_difference(grid=grid)
+            return self._hessian_via_richardson(grid=grid)
         return self._hessian_via_jax(grid=grid, xp=xp)
 
+    def _hessian_via_richardson(self, grid, buffer: float = 0.01) -> Tuple:
+        yy_h, xy_h, yx_h, xx_h = self._hessian_via_finite_difference(
+            grid=grid, buffer=buffer
+        )
+        yy_h2, xy_h2, yx_h2, xx_h2 = self._hessian_via_finite_difference(
+            grid=grid, buffer=buffer / 2.0
+        )
+        hessian_yy = (4.0 * yy_h2 - yy_h) / 3.0
+        hessian_xy = (4.0 * xy_h2 - xy_h) / 3.0
+        hessian_yx = (4.0 * yx_h2 - yx_h) / 3.0
+        hessian_xx = (4.0 * xx_h2 - xx_h) / 3.0
+        return hessian_yy, hessian_xy, hessian_yx, hessian_xx
+
     def _hessian_via_jax(self, grid, xp) -> Tuple:
         import jax
         import jax.numpy as jnp

test_autogalaxy/operate/test_deflections.py

@@ -102,10 +102,10 @@ def test__hessian_from__diagonal_grid__correct_values():
     od = LensCalc.from_mass_obj(mp)
     hessian_yy, hessian_xy, hessian_yx, hessian_xx = od.hessian_from(grid=grid)
 
-    assert hessian_yy == pytest.approx(np.array([1.3883822, 0.694127]), 1.0e-4)
-    assert hessian_xy == pytest.approx(np.array([-1.388124, -0.694094]), 1.0e-4)
-    assert hessian_yx == pytest.approx(np.array([-1.388165, -0.694099]), 1.0e-4)
-    assert hessian_xx == pytest.approx(np.array([1.3883824, 0.694127]), 1.0e-4)
+    assert hessian_yy == pytest.approx(np.array([1.3882113, 0.6941056]), 1.0e-4)
+    assert hessian_xy == pytest.approx(np.array([-1.3882113, -0.6941056]), 1.0e-4)
+    assert hessian_yx == pytest.approx(np.array([-1.3882113, -0.6941056]), 1.0e-4)
+    assert hessian_xx == pytest.approx(np.array([1.3882113, 0.6941056]), 1.0e-4)
 
 
 def test__hessian_from__axis_aligned_grid__correct_values():
@@ -152,8 +152,51 @@ def test__magnification_2d_via_hessian_from():
     od = LensCalc.from_mass_obj(mp)
     magnification = od.magnification_2d_via_hessian_from(grid=grid)
 
-    assert magnification.in_list[0] == pytest.approx(-0.56303, 1.0e-4)
-    assert magnification.in_list[1] == pytest.approx(-2.57591, 1.0e-4)
+    assert magnification.in_list[0] == pytest.approx(-0.5629291, 1.0e-4)
+    assert magnification.in_list[1] == pytest.approx(-2.575917, 1.0e-4)
+
+
+def test__hessian_from__np_richardson_matches_jax_jacfwd_to_float64():
+    import jax.numpy as jnp
+
+    grid = ag.Grid2DIrregular(values=[(0.5, 0.5), (1.0, 1.0), (0.7, -0.3)])
+
+    mp = ag.mp.Isothermal(
+        centre=(0.0, 0.0), ell_comps=(0.05, -0.111111), einstein_radius=1.5
+    )
+
+    od = LensCalc.from_mass_obj(mp)
+
+    np_hess = od.hessian_from(grid=grid, xp=np)
+    jnp_hess = od.hessian_from(grid=grid, xp=jnp)
+
+    for np_component, jnp_component in zip(np_hess, jnp_hess):
+        np.testing.assert_allclose(
+            np.asarray(np_component),
+            np.asarray(jnp_component),
+            rtol=1.0e-8,
+        )
+
+
+def test__magnification_2d_via_hessian_from__np_jnp_agree_to_float64():
+    import jax.numpy as jnp
+
+    grid = ag.Grid2DIrregular(values=[(0.5, 0.5), (1.0, 1.0), (0.7, -0.3)])
+
+    mp = ag.mp.Isothermal(
+        centre=(0.0, 0.0), ell_comps=(0.05, -0.111111), einstein_radius=1.5
+    )
+
+    od = LensCalc.from_mass_obj(mp)
+
+    np_mag = od.magnification_2d_via_hessian_from(grid=grid, xp=np)
+    jnp_mag = od.magnification_2d_via_hessian_from(grid=grid, xp=jnp)
+
+    np.testing.assert_allclose(
+        np.asarray(np_mag.array),
+        np.asarray(jnp_mag),
+        rtol=1.0e-7,
+    )
 
 
 def test__tangential_critical_curve_list_from__radius_matches_einstein_radius():