Merge pull request #84 from scientificcomputing/dokken/replace-taylor-third-order

Replace third order taylor test with finite difference check of hessian against gradient

Author: jorgensd

tests/conftest.py

@@ -0,0 +1,85 @@
+import numpy as np
+import pyadjoint
+import pytest
+
+
+@pytest.fixture
+def assert_hessian_matches_finite_difference():
+    """A Hessian-accuracy checker, as a more numerically robust
+    alternative to {py:class}`pyadjoint.taylor_test`'s standard rate-3 Hessian-corrected check.
+
+    That check needs cancelling several O(1) quantities down to an O(eps**3) remainder
+    at eps <= 0.01, which the direct (MUMPS) LU factorization behind the
+    adjoint/TLM/second-order-adjoint solves cannot always resolve to the precision it
+    requires -- observed for saddle-point (e.g. Taylor-Hood velocity/pressure) and other
+    blocked/nonlinear ``NonlinearProblem``/``LinearProblem`` systems in this suite, where
+    ``mat_mumps_icntl_24`` alone does not fully resolve it and PETSc's ``SNESSolve`` can
+    even intermittently fail to converge (error code 91) under repeated nearby re-solves.
+    The returned checker instead only needs the *gradient*'s own precision (already
+    validated wherever a rate-2 ``taylor_test`` passes), comparing
+    ``Jhat.hessian(h)._ad_dot(h)`` directly against a central difference of
+    ``Jhat.derivative()._ad_dot(h)``.
+
+    Returns:
+        A callable ``check(Jhat, m, h, *, fd_eps=1e-3, rtol=1e-2, atol=1e-2)`` -- see
+        its own docstring for details. Exposed as a fixture (rather than a plain
+        module-level function) so every test can use it with no import of its own,
+        matching this project's ``--import-mode=importlib`` pytest configuration.
+    """
+
+    def _check(
+        Jhat: pyadjoint.ReducedFunctional,
+        m: pyadjoint.OverloadedType,
+        h: pyadjoint.OverloadedType,
+        *,
+        fd_eps: float = 1e-3,
+        rtol: float = 1e-2,
+        atol: float = 1e-2,
+    ) -> None:
+        """Verify ``Jhat``'s Hessian-vector product against a central difference of its own gradient.
+
+        Uses ``m``/``h``'s own ``_ad_add``/``_ad_mul`` (the same primitives
+        ``pyadjoint.taylor_test`` perturbs its own evaluation points with) rather than
+        type-specific perturbation code, so this works unchanged for a
+        {py:class}`dolfinx_adjoint.Function`,
+        {py:class}`dolfinx_adjoint.Constant``, or any other
+        {py:class}`pyadjoint.OverloadedType` control.
+
+        Leaves ``Jhat`` evaluated at ``m`` on return.
+
+        ``Hm`` and ``Hm_fd`` are two independent estimates of the same mathematical
+        quantity, so they should agree up to two, unrelated, and much smaller error
+        sources: (1) central-difference truncation, ``O(fd_eps**2)`` relative --
+        `<1e-5` relative at the default ``fd_eps=1e-3``, negligible here; and (2)
+        whatever precision the adjoint/TLM/second-order-adjoint linear solves and the
+        forward (possibly SNES) solve actually achieve at the two perturbed evaluation
+        points. If a particular problem's own solves are markedly less precise (an
+        iterative KSP/SNES rather than a direct LU factorization, say), loosen
+        ``rtol``/``atol`` explicitly for that call rather than lowering the default.
+
+        Args:
+            Jhat: The reduced functional to check.
+            m: The control value to evaluate the Hessian at.
+            h: The direction to evaluate the Hessian-vector product/gradient in.
+            fd_eps: Finite-difference step size, in units of ``h``.
+            rtol: Relative tolerance passed to ``numpy.isclose`` -- see above for why
+                ``1e-2`` is the default.
+            atol: Absolute tolerance passed to ``numpy.isclose`` -- see above for why
+                ``1e-2`` is the default.
+        """
+
+        def dJdm_at(scale: float) -> float:
+            Jhat(m._ad_add(h._ad_mul(scale)))
+            return Jhat.derivative()._ad_dot(h)
+
+        Jhat(m)
+        Jhat.derivative()
+        Hm = Jhat.hessian(h)._ad_dot(h)
+        Hm_fd = (dJdm_at(fd_eps) - dJdm_at(-fd_eps)) / (2 * fd_eps)
+        Jhat(m)
+        assert np.isclose(Hm, Hm_fd, rtol=rtol, atol=atol), (
+            f"Hessian-vector product {Hm} did not match central-difference-of-gradient "
+            f"estimate {Hm_fd} (fd_eps={fd_eps})"
+        )
+
+    return _check

tests/test_blocked_problem.py

@@ -24,7 +24,7 @@ def mesh_2D():
 
 
 @pytest.mark.parametrize("use_mixed_space", [True, False])
-def test_solver(use_mixed_space: bool, mesh_2D):
+def test_solver(use_mixed_space: bool, mesh_2D, assert_hessian_matches_finite_difference):
     pyadjoint.get_working_tape().clear_tape()
     mesh = mesh_2D
     el_u = basix.ufl.element("P", mesh.basix_cell(), 2, shape=(mesh.geometry.dim,))
@@ -116,13 +116,10 @@ def L1(mesh, q):
             f"Expected convergence rate close to 2.0, got {min_rate}"
         )
 
-        # Scale perturbation for hessian
-        Jh(d)
-        dJdm = Jh.derivative()._ad_dot(e)
-        hessian = Jh.hessian(e)
-        dHddu = hessian._ad_dot(e)
-        min_rate = pyadjoint.taylor_test(Jh, d, e, dJdm=dJdm, Hm=dHddu)
-        assert np.isclose(min_rate, 3.0, rtol=0.1, atol=0.1), f"Expected convergence rate close to 3.0, got {min_rate}"
+        # This is a saddle-point (velocity/pressure) system -- the standard rate-3
+        # taylor_test's cubic remainder is not a robust check on it (see
+        # assert_hessian_matches_finite_difference's own docstring in conftest.py).
+        assert_hessian_matches_finite_difference(Jh, d, e)
 
         z = Function(Z)
         z.interpolate(
@@ -133,16 +130,11 @@ def L1(mesh, q):
             lambda x: (0.8 * baseline * x[1] ** 2, 2 * baseline * x[0] ** 2)
         )  # NOTE: Has to be divergence free
         f.x.scatter_forward()
-        Jh(z)
-        dJdm = Jh.derivative()._ad_dot(f)
-        hessian = Jh.hessian(f)
-        dHddu = hessian._ad_dot(f)
-        min_rate = pyadjoint.taylor_test(Jh, z, f, dJdm=dJdm, Hm=dHddu)
-        assert np.isclose(min_rate, 3.0, rtol=0.1, atol=0.1), f"Expected convergence rate close to 3.0, got {min_rate}"
+        assert_hessian_matches_finite_difference(Jh, z, f)
 
 
 @pytest.mark.parametrize("use_mixed_space", [True, False])
-def test_nonlinear_solver(use_mixed_space: bool, mesh_2D):
+def test_nonlinear_solver(use_mixed_space: bool, mesh_2D, assert_hessian_matches_finite_difference):
     """As ``test_solver``, but for a blocked ``NonlinearProblem``: a Navier-Stokes-like
     velocity/pressure system with a viscosity control, genuinely nonlinear in the state
     via the convective term, exercising the same forward/adjoint/TLM/Hessian paths as
@@ -241,12 +233,10 @@ def test_nonlinear_solver(use_mixed_space: bool, mesh_2D):
             f"Expected convergence rate close to 2.0, got {min_rate}"
         )
 
-        Jh(d)
-        dJdm = Jh.derivative()._ad_dot(e)
-        hessian = Jh.hessian(e)
-        dHddu = hessian._ad_dot(e)
-        min_rate = pyadjoint.taylor_test(Jh, d, e, dJdm=dJdm, Hm=dHddu)
-        assert np.isclose(min_rate, 3.0, rtol=0.1, atol=0.1), f"Expected convergence rate close to 3.0, got {min_rate}"
+        # This is a saddle-point (velocity/pressure) system -- the standard rate-3
+        # taylor_test's cubic remainder is not a robust check on it (see
+        # assert_hessian_matches_finite_difference's own docstring in conftest.py).
+        assert_hessian_matches_finite_difference(Jh, d, e)
 
         # A second, independent evaluation point/direction: a cached-but-unrefreshed
         # adjoint/TLM/Hessian operator (see tests/test_tlm_update.py) could pass the
@@ -258,9 +248,4 @@ def test_nonlinear_solver(use_mixed_space: bool, mesh_2D):
         assert np.min(mu2.x.array - h2.x.array) > 0.0, (
             "Taylor test perturbation must not violate positivity of viscosity"
         )
-        Jh(mu2)
-        dJdm = Jh.derivative()._ad_dot(h2)
-        hessian = Jh.hessian(h2)
-        dHddu = hessian._ad_dot(h2)
-        min_rate = pyadjoint.taylor_test(Jh, mu2, h2, dJdm=dJdm, Hm=dHddu)
-        assert np.isclose(min_rate, 3.0, rtol=0.1, atol=0.1), f"Expected convergence rate close to 3.0, got {min_rate}"
+        assert_hessian_matches_finite_difference(Jh, mu2, h2)

tests/test_tlm_update.py

@@ -96,7 +96,9 @@ def _viscous_stokes(mesh):
 
 
 @pytest.mark.parametrize("warm_up_at_another_point", [False, True])
-def test_hessian_is_independent_of_previous_evaluation_points(warm_up_at_another_point, mesh_2D):
+def test_hessian_is_independent_of_previous_evaluation_points(
+    warm_up_at_another_point, mesh_2D, assert_hessian_matches_finite_difference
+):
     """The Hessian at ``m2`` must not depend on whether ``J`` was evaluated at ``m1`` first.
 
     Both parametrizations run the identical second-order Taylor test at ``m2``.  The only
@@ -120,12 +122,10 @@ def test_hessian_is_independent_of_previous_evaluation_points(warm_up_at_another
         Jh.hessian(h)
 
     assert np.min(m2.x.array - h.x.array) > 0.0, "Taylor test perturbation must not violate positivity of viscosity"
-    Jh(m2)
-    dJdm = Jh.derivative()._ad_dot(h)
-    Hm = Jh.hessian(h)._ad_dot(h)
-
-    min_rate = pyadjoint.taylor_test(Jh, m2, h, dJdm=dJdm, Hm=Hm)
-    assert np.isclose(min_rate, 3.0, rtol=0.1, atol=0.1), f"Expected convergence rate close to 3.0, got {min_rate}"
+    # This is a saddle-point (velocity/pressure) system -- the standard rate-3
+    # taylor_test's cubic remainder is not a robust check on it (see
+    # assert_hessian_matches_finite_difference's own docstring in conftest.py).
+    assert_hessian_matches_finite_difference(Jh, m2, h)
 
 
 def _navier_stokes(mesh):
@@ -206,7 +206,9 @@ def _navier_stokes(mesh):
     return pyadjoint.ReducedFunctional(J, pyadjoint.Control(mu)), Z
 
 
-def test_hessian_is_independent_of_previous_evaluation_points_navier_stokes(mesh_2D):
+def test_hessian_is_independent_of_previous_evaluation_points_navier_stokes(
+    mesh_2D, assert_hessian_matches_finite_difference
+):
     """``test_hessian_is_independent_of_previous_evaluation_points``'s ``NonlinearProblem``
     sibling: the same second-order Taylor test, but on a genuinely nonlinear, blocked
     (multi-output) residual, exercising the blocked Hessian path in
@@ -224,12 +226,10 @@ def test_hessian_is_independent_of_previous_evaluation_points_navier_stokes(mesh
     h = Function(Z)
     h.interpolate(lambda x: 0.4 * baseline * np.sin(3 * x[0]))
     assert np.min(m2.x.array - h.x.array) > 0.01, "Taylor test perturbation must not violate positivity of viscosity"
-    Jh(m2)
-    dJdm = Jh.derivative()._ad_dot(h)
-    Hm = Jh.hessian(h)._ad_dot(h)
-
-    min_rate = pyadjoint.taylor_test(Jh, m2, h, dJdm=dJdm, Hm=Hm)
-    assert np.isclose(min_rate, 3.0, rtol=0.1, atol=0.1), f"Expected convergence rate close to 3.0, got {min_rate}"
+    # This is a saddle-point (velocity/pressure) system -- the standard rate-3
+    # taylor_test's cubic remainder is not a robust check on it (see
+    # assert_hessian_matches_finite_difference's own docstring in conftest.py).
+    assert_hessian_matches_finite_difference(Jh, m2, h)
 
 
 def _diffusive_poisson(mesh):
@@ -309,7 +309,7 @@ def test_hessian_is_independent_of_previous_evaluation_points_scalar(warm_up_at_
     assert np.isclose(min_rate, 3.0, rtol=0.1, atol=0.1), f"Expected convergence rate close to 3.0, got {min_rate}"
 
 
-def test_hessian_mpi_breakdown(mesh_2D):
+def test_hessian_mpi_breakdown(mesh_2D, assert_hessian_matches_finite_difference):
     pyadjoint.get_working_tape().clear_tape()
     Jh, Z = _viscous_stokes(mesh_2D)
     baseline = 23.2
@@ -358,8 +358,8 @@ def test_hessian_mpi_breakdown(mesh_2D):
     hess_diff = np.linalg.norm(H_cold_array - H_warm_array)
     assert hess_diff < 1e-10, f"Rank {rank}: Hessian differs after warm up! Diff: {hess_diff}"
 
-    # 4. If arrays match, run Taylor test
-    dJdm = dJ_warm._ad_dot(h)
-    Hm = H_warm._ad_dot(h)
-    min_rate = pyadjoint.taylor_test(Jh, m2, h, dJdm=dJdm, Hm=Hm)
-    assert np.isclose(min_rate, 3.0, rtol=0.1, atol=0.1), f"Expected 3.0, got {min_rate}"
+    # 4. If arrays match, verify the Hessian value itself is correct. This is a
+    # saddle-point (velocity/pressure) system -- the standard rate-3 taylor_test's
+    # cubic remainder is not a robust check on it (see
+    # assert_hessian_matches_finite_difference's own docstring in conftest.py).
+    assert_hessian_matches_finite_difference(Jh, m2, h)