Add third order Hessian taylor tests (#721)

Co-authored-by: Thomas Bendall <14180399+tommbendall@users.noreply.github.com>
Co-authored-by: Connor Ward <c.ward20@imperial.ac.uk>
Co-authored-by: jshipton <j.shipton@exeter.ac.uk>

Author: 4 people

gusto/solvers/preconditioners.py

@@ -346,7 +346,7 @@ def apply(self, pc, x, y):
 
             unbroken_res_hdiv = self.unbroken_residual.subfunctions[self.vidx]
             broken_res_hdiv = self.broken_residual.subfunctions[self.vidx]
-            broken_res_hdiv.assign(0)
+            broken_res_hdiv.zero()
             self.average_kernel.apply(broken_res_hdiv, self.weight, unbroken_res_hdiv)
 
             # Compute the rhs for the multiplier system
@@ -374,7 +374,7 @@ def apply(self, pc, x, y):
             # Compute the hdiv projection of the broken hdiv solution
             broken_hdiv = self.broken_solution.subfunctions[self.vidx]
             unbroken_hdiv = self.unbroken_solution.subfunctions[self.vidx]
-            unbroken_hdiv.assign(0)
+            unbroken_hdiv.zero()
 
             self.average_kernel.apply(unbroken_hdiv, self.weight, broken_hdiv)
 
@@ -735,7 +735,7 @@ def apply(self, pc, x, y):
 
         # Recover broken u and rho
         u_broken, rho, l = self.y_hybrid.subfunctions
-        self.u_hdiv.assign(0)
+        self.u_hdiv.zero()
         self._average_kernel.apply(self.u_hdiv, self._weight, u_broken)
         for bc in self.bcs:
             bc.zero(self.u_hdiv)
@@ -745,7 +745,7 @@ def apply(self, pc, x, y):
         self.y.subfunctions[1].assign(rho)
 
         # Recover theta
-        self.theta.assign(0)
+        self.theta.zero()
         self.theta_solver.solve()
         self.y.subfunctions[2].assign(self.theta)
 

gusto/solvers/solver_presets.py

@@ -40,7 +40,7 @@ def hybridised_solver_parameters(equation, solver_prognostics, alpha=0.5, tau_va
 
     # Callback function for the nullspace of the trace system
     def trace_nullsp(T):
-        return VectorSpaceBasis(constant=True)
+        return VectorSpaceBasis(constant=True, comm=equation.function_space.mesh().comm)
 
     # Prepare some generic variables to help with logic below ------------------
 

integration-tests/adjoints/test_diffusion_sensitivity.py

@@ -3,32 +3,16 @@
 
 from firedrake import *
 from firedrake.adjoint import *
-from pyadjoint import get_working_tape
 from gusto import *
 
 
 @pytest.fixture(autouse=True)
-def handle_taping():
-    yield
-    tape = get_working_tape()
-    tape.clear_tape()
-
-
-@pytest.fixture(autouse=True, scope="module")
-def handle_annotation():
-    from firedrake.adjoint import annotate_tape, continue_annotation
-    if not annotate_tape():
-        continue_annotation()
-    yield
-    # Ensure annotation is paused when we finish.
-    annotate = annotate_tape()
-    if annotate:
-        pause_annotation()
+def autouse_set_test_tape(set_test_tape):
+    pass
 
 
 @pytest.mark.parametrize("nu_is_control", [True, False])
 def test_diffusion_sensitivity(nu_is_control, tmpdir):
-    assert get_working_tape()._blocks == []
     n = 30
     mesh = PeriodicUnitSquareMesh(n, n)
     output = OutputParameters(dirname=str(tmpdir))
@@ -41,38 +25,59 @@ def test_diffusion_sensitivity(nu_is_control, tmpdir):
 
     R = FunctionSpace(mesh, "R", 0)
     # We need to define nu as a function in order to have a control variable.
-    nu = Function(R, val=0.0001)
+    nu = Function(R, val=0.2)
     diffusion_params = DiffusionParameters(mesh, kappa=nu)
     eqn = DiffusionEquation(domain, V, "f", diffusion_parameters=diffusion_params)
 
-    diffusion_scheme = BackwardEuler(domain)
+    # Don't let an inexact solve get in the way of a good Taylor test
+    solver_parameters = {
+        'snes_type': 'ksponly',
+        'ksp_type': 'preonly',
+        'pc_type': 'lu',
+        'pc_factor_mat_solver_type': 'mumps',
+    }
+    diffusion_scheme = BackwardEuler(domain, solver_parameters=solver_parameters)
     diffusion_methods = [CGDiffusion(eqn, "f", diffusion_params)]
     timestepper = Timestepper(eqn, diffusion_scheme, io, spatial_methods=diffusion_methods)
 
     x = SpatialCoordinate(mesh)
     fexpr = as_vector((sin(2*pi*x[0]), cos(2*pi*x[1])))
-    timestepper.fields("f").interpolate(fexpr)
+    ic = Function(V).interpolate(fexpr)
 
-    end = 0.1
-    timestepper.run(0., end)
+    # These are the only operations we are interested in rerunning with the tape.
+    with set_working_tape() as tape:
+        # initialise the solution
+        timestepper.fields("f").assign(ic)
 
-    u = timestepper.fields("f")
-    J = assemble(inner(u, u)*dx)
+        # propagate forward
+        end = 0.1
+        timestepper.run(0., end)
 
-    if nu_is_control:
-        control = Control(nu)
-        h = Function(R, val=0.0001)  # the direction of the perturbation
-    else:
-        control = Control(u)
-        # the direction of the perturbation
-        h = Function(V).interpolate(fexpr * np.random.rand())
+        # Make sure we have more than a quadratic nonlinearity so Jhat.hessian isn't exact.
+        u = timestepper.fields("f")
+        J = assemble(inner(u, u)*dx)**2
 
-    # the functional as a pure function of nu
-    Jhat = ReducedFunctional(J, control)
+        if nu_is_control:
+            m = nu
+        else:
+            m = ic
 
+        # the functional as a pure function of the control
+        Jhat = ReducedFunctional(J, Control(m), tape=tape)
+
+    assert np.allclose(J, Jhat(m)), "Functional re-evaluation does not match original evaluation."
+
+    # perturbation direction for taylor test
     if nu_is_control:
-        assert np.allclose(J, Jhat(nu))
-        assert taylor_test(Jhat, nu, h) > 1.95
+        h = Function(R, val=0.1)
     else:
-        assert np.allclose(J, Jhat(u))
-        assert taylor_test(Jhat, u, h) > 1.95
+        h = Function(V).interpolate(as_vector([cos(4*pi*x[1]), sin(4*pi*x[0])]))
+
+    # Check the TLM explicitly before checking the Hessian (which relies on the tlm)
+    assert taylor_test(Jhat, m, h, dJdm=Jhat.tlm(h)) > 1.95, "TLM is not second order accurate."
+
+    # Check the re-evaluation, derivative, and Hessian all converge at the expected rates.
+    taylor = taylor_to_dict(Jhat, m, h)
+    assert min(taylor['R0']['Rate']) > 0.95, taylor['R0']
+    assert min(taylor['R1']['Rate']) > 1.95, taylor['R1']
+    assert min(taylor['R2']['Rate']) > 2.95, taylor['R2']

integration-tests/adjoints/test_moist_thermal_williamson_5_sensitivity.py

@@ -21,7 +21,6 @@
     SpatialCoordinate, as_vector, pi, sqrt, min_value, exp, cos, sin, assemble, dx, inner, Function
 )
 from firedrake.adjoint import *
-from pyadjoint import get_working_tape
 from gusto import (
     Domain, IO, OutputParameters, Timestepper, RK4, DGUpwind,
     ShallowWaterParameters, ThermalShallowWaterEquations, lonlatr_from_xyz,
@@ -31,29 +30,14 @@
 
 
 @pytest.fixture(autouse=True)
-def handle_taping():
-    yield
-    tape = get_working_tape()
-    tape.clear_tape()
-
-
-@pytest.fixture(autouse=True, scope="module")
-def handle_annotation():
-    from firedrake.adjoint import annotate_tape, continue_annotation
-    if not annotate_tape():
-        continue_annotation()
-    yield
-    # Ensure annotation is paused when we finish.
-    annotate = annotate_tape()
-    if annotate:
-        pause_annotation()
+def autouse_set_test_tape(set_test_tape):
+    pass
 
 
 def test_moist_thermal_williamson_5_sensitivity(
         tmpdir, ncells_per_edge=8, dt=600, tmax=50.*24.*60.*60.,
         dumpfreq=2880
 ):
-    assert get_working_tape()._blocks == []
     # ------------------------------------------------------------------------ #
     # Parameters for test case
     # ------------------------------------------------------------------------ #
@@ -150,9 +134,18 @@ def gamma_v(x_in):
         DGUpwind(eqns, field_name) for field_name in eqns.field_names
     ]
 
+    # Don't let an inexact solve get in the way of a good Taylor test
+    solver_parameters = {
+        'snes_type': 'ksponly',
+        'ksp_type': 'preonly',
+        'pc_type': 'lu',
+        'pc_factor_mat_solver_type': 'mumps',
+    }
+
     # Timestepper
     stepper = Timestepper(
-        eqns, RK4(domain), io, spatial_methods=transport_methods
+        eqns, RK4(domain, solver_parameters=solver_parameters),
+        io, spatial_methods=transport_methods
     )
 
     # ------------------------------------------------------------------------ #
@@ -188,29 +181,41 @@ def gamma_v(x_in):
     # Expression for initial vapour depends on initial saturation
     vexpr = mu2 * q_sat(bexpr, Dexpr)
 
-    # Initialise (cloud and rain initially zero)
-    u0.project(uexpr)
-    D0.interpolate(Dexpr)
-    b0.interpolate(bexpr)
-    v0.interpolate(vexpr)
-    c0.assign(0.0)
-    r0.assign(0.0)
+    # Control
+    m_D = Function(D0.function_space()).interpolate(Dexpr)
+
+    with set_working_tape() as tape:
+        # Initialise (cloud and rain initially zero)
+        u0.project(uexpr)
+        D0.interpolate(m_D)
+        b0.interpolate(bexpr)
+        v0.interpolate(vexpr)
+        c0.zero()
+        r0.zero()
+
+        # ----------------------------------------------------------------- #
+        # Run
+        # ----------------------------------------------------------------- #
+        # two timesteps so we hit the codepath for reshuffling data between steps
+        stepper.run(t=0., tmax=2*dt)
+
+        u_tf = stepper.fields('u')  # Final velocity field
+        D_tf = stepper.fields('D')  # Final depth field
+
+        J = assemble(0.5*inner(u_tf, u_tf)*dx + 0.5*g*D_tf**2*dx)
+
+        Jhat = ReducedFunctional(J, Control(m_D), tape=tape)
 
-    # ----------------------------------------------------------------- #
-    # Run
-    # ----------------------------------------------------------------- #
-    stepper.run(t=0, tmax=dt)
+    assert np.allclose(Jhat(m_D), J), "Functional re-evaluation does not match original evaluation."
 
-    u_tf = stepper.fields('u')  # Final velocity field
-    D_tf = stepper.fields('D')  # Final depth field
+    # perturbation direction for taylor test
+    h_D = Function(D0.function_space())
+    h_D.interpolate(0.1*(Dexpr - (mean_depth - tpexpr)))
 
-    J = assemble(0.5*inner(u_tf, u_tf)*dx + 0.5*g*D_tf**2*dx)
+    # Check the TLM explicitly before checking the Hessian (which relies on the tlm)
+    assert taylor_test(Jhat, m_D, h_D, dJdm=Jhat.tlm(h_D)) > 1.95, "TLM is not second order accurate."
 
-    Jhat = ReducedFunctional(J, Control(D0))
+    assert taylor_test(Jhat, m_D, h_D) > 1.95, "Adjoint derivative is not second order accurate."
 
-    assert np.allclose(Jhat(D0), J)
-    with stop_annotating():
-        # Stop annotation to perform the Taylor test
-        h0 = Function(D0.function_space())
-        h0.assign(D0 * np.random.rand())
-        assert taylor_test(Jhat, D0, h0) > 1.95
+    # NOTE: Currently computing the Hessian will raise an exception
+    # see https://github.com/FEniCS/ufl/issues/477