Fix NLVS hessian evaluation with Cofunction RHS

firedrake/adjoint_utils/blocks/solving.py

@@ -1,7 +1,7 @@
 import numpy
 import ufl
 from ufl.domain import extract_domains, extract_unique_domain
-from ufl import replace
+from ufl import replace, ZeroBaseForm
 from ufl.formatting.ufl2unicode import ufl2unicode
 from enum import Enum
 
@@ -261,7 +261,7 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_
                 continue
             if dep.output is self.forward_cache.func:  # Can't compute dependence on initial guess
                 continue
-            if len(d2Fdmdu.integrals()) > 0:
+            if not d2Fdmdu.empty() and len(d2Fdmdu.integrals()) > 0:
                 hessian_rhs -= firedrake.assemble(d2Fdmdu)
 
         # 2. Solve adjoint system
@@ -302,10 +302,10 @@ def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs, block_v
 
         hessian_output = 0
 
-        for form in (self.hessian_cache.d2Fdudm_forms[idx],
-                     self.hessian_cache.dFdm_adj2_forms[idx],
+        for form in (self.hessian_cache.dFdm_adj2_forms[idx],
+                     self.hessian_cache.d2Fdudm_forms[idx],
                      *relevant_d2Fdm2_forms):
-            if not form.empty():
+            if not isinstance(form, ZeroBaseForm):
                 hessian_output += firedrake.assemble(form)
 
         return hessian_output

firedrake/adjoint_utils/function.py

@@ -223,9 +223,9 @@ def _ad_convert_riesz(self, value, riesz_map=None):
         return value.riesz_representation(riesz_map=riesz_map or "L2")
 
     def _ad_init_zero(self, dual=False):
-        from firedrake import Function, Cofunction
+        from firedrake import Function
         if dual:
-            return Cofunction(self.function_space().dual())
+            return Function(self.function_space().dual())
         else:
             return Function(self.function_space())
 
@@ -311,7 +311,7 @@ def _ad_iadd(self, other):
         self += other
         return self
 
-    def _ad_function_space(self, mesh):
+    def _ad_function_space(self, mesh=None):
         return self.ufl_function_space()
 
     def _reduce(self, r, r0):

firedrake/adjoint_utils/variational_solver.py

@@ -5,13 +5,28 @@
     NonlinearVariationalSolveBlock, CachedSolverBlock)
 from firedrake.adjoint_utils.blocks.solving import solve_init_params
 from firedrake.ufl_expr import derivative, adjoint, action
-from ufl import replace, Action
+from ufl import replace, Action, ZeroBaseForm, Form
 from ufl.algorithms import expand_derivatives
 from ufl.domain import extract_domains
+from ufl.constantvalue import Zero
 from types import SimpleNamespace
 from collections import namedtuple
 
 
+def _set_arguments_if_zero(form, args):
+    """
+    If form is an empty ufl.Form or ufl.Zero then replace
+    it with a ZeroBaseForm with the provided arguments.
+
+    See: https://github.com/FEniCS/ufl/issues/396
+    """
+    if isinstance(form, Zero):
+        form = ZeroBaseForm(args)
+    if isinstance(form, Form) and form.empty():
+        form = ZeroBaseForm(args)
+    return form
+
+
 ForwardSolveRecomputeCache = namedtuple(
     'ForwardSolveRecomputeCache',
     field_names=[
@@ -380,7 +395,7 @@ def _ad_adjoint_cache(self):
     @no_annotations
     def _ad_hessian_cache(self):
         from firedrake import (
-            Function, TrialFunction, TestFunction,
+            Function, TrialFunction, TestFunction, Cofunction,
             SpatialCoordinate, MeshGeometry,
         )
 
@@ -446,28 +461,38 @@ def _ad_hessian_cache(self):
             else:
                 dm = TrialFunction(m.function_space())
                 # XXX should we try inverting this back to the way it was before?
-                dFdm = derivative(F, m, dm)
+                dFdm = expand_derivatives(derivative(-F, m, dm))
 
-                dFdm_adj = -expand_derivatives(action(adjoint(dFdm), adj_sol))
-                dFdm_adj2 = -action(adjoint(dFdm), adj2_sol)
+                dFdm_star = adjoint(dFdm)
+                # 0. fully special case Cofunctions
+                if isinstance(m, Cofunction):
+                    dFdm_adj = Action(dFdm_star, adj_sol)
+                    dFdm_adj2 = Action(dFdm_star, adj2_sol)
+                else:
+                    dFdm_adj = action(dFdm_star, adj_sol)
+                    dFdm_adj2 = action(dFdm_star, adj2_sol)
+
+            args = (TestFunction(m._ad_function_space()),)
 
-            dFdm_adj2_forms.append(dFdm_adj2)
+            dFdm_adj2_forms.append(_set_arguments_if_zero(dFdm_adj2, args))
 
             d2Fdudm = derivative(dFdm_adj, u, tlm_output)
-            d2Fdudm_forms.append(expand_derivatives(d2Fdudm))
+            d2Fdudm = expand_derivatives(d2Fdudm)
+
+            d2Fdudm_forms.append(_set_arguments_if_zero(d2Fdudm, args))
 
             d2Fdm2_adj_forms_k = []
             for m2, dm2 in zip(self._ad_forward_cache.replaced_deps,
                                self._ad_tangent_cache.replaced_tlms):
                 d2Fdm2_adj = expand_derivatives(
                     derivative(dFdm_adj, m2, dm2))
-                d2Fdm2_adj_forms_k.append(d2Fdm2_adj)
+                d2Fdm2_adj_forms_k.append(_set_arguments_if_zero(d2Fdm2_adj, args))
 
             for m2, dm2 in zip(self._ad_forward_cache.meshes,
                                self._ad_tangent_cache.mesh_tlms):
                 X = SpatialCoordinate(m2)
                 d2Fdm2_adj = expand_derivatives(derivative(dFdm_adj, X, dm2))
-                d2Fdm2_adj_forms_k.append(d2Fdm2_adj)
+                d2Fdm2_adj_forms_k.append(_set_arguments_if_zero(d2Fdm2_adj, args))
 
             d2Fdm2_adj_forms.append(d2Fdm2_adj_forms_k)
 

tests/firedrake/adjoint/test_hessian.py

@@ -15,7 +15,8 @@ def rg():
 
 
 @pytest.mark.skipcomplex
-def test_simple_solve(rg):
+@pytest.mark.parametrize("forcing_type", ["unassembled", "assembled"])
+def test_simple_solve(rg, forcing_type):
     tape = Tape()
     set_working_tape(tape)
 
@@ -30,12 +31,18 @@ def test_simple_solve(rg):
     a = u*v*dx
     L = f*v*dx
 
+    if forcing_type == "assembled":
+        L = assemble(L)
+
     u_ = Function(V)
 
     solve(a == L, u_)
 
     L = u_*v*dx
 
+    if forcing_type == "assembled":
+        L = assemble(L)
+
     u_sol = Function(V)
     solve(a == L, u_sol)
 

tests/firedrake/adjoint/test_nlvs.py

@@ -9,7 +9,7 @@ def autouse_set_test_tape(set_test_tape):
     pass
 
 
-def forward(ic, dt, nt, bc_arg=None):
+def forward(ic, dt, nt, bc_arg=None, forcing_type="unassembled"):
     """Burgers equation solver."""
     V = ic.function_space()
 
@@ -26,21 +26,33 @@ def forward(ic, dt, nt, bc_arg=None):
     u1 = Function(V)
     v = TestFunction(V)
 
-    F = ((u1 - u0)*v
+    F = (u1*v
          + dt*u1*u1.dx(0)*v
          + dt*nu*u1.dx(0)*v.dx(0))*dx
 
+    mass_prev = u0*v*dx
+
+    if forcing_type == "unassembled":
+        F -= mass_prev
+    elif forcing_type == "assembled":
+        L = Cofunction(V.dual())
+        F -= L
+    else:
+        raise ValueError(f"Unrecognised {forcing_type=}")
+
     problem = NonlinearVariationalProblem(F, u1, bcs=bcs)
     solver = NonlinearVariationalSolver(problem)
 
     u1.assign(ic)
 
     for i in range(nt):
         u0.assign(u1)
+        if forcing_type == "assembled":
+            assemble(mass_prev, tensor=L)
         solver.solve()
         nu += dt
-        # if bc_arg:
-        #     bc_val.assign(bc_val + dt/nt)
+        if bc_arg:
+            bc_val.assign(bc_val + dt/nt)
 
     J = assemble(u1*u1*dx)
     return J
@@ -52,7 +64,8 @@ def forward(ic, dt, nt, bc_arg=None):
                                           "bc_control"])
 @pytest.mark.parametrize("bc_type", ["neumann_bc",
                                      "dirichlet_bc"])
-def test_nlvs_adjoint(control_type, bc_type):
+@pytest.mark.parametrize("forcing_type", ["assembled", "unassembled"])
+def test_nlvs_adjoint(control_type, bc_type, forcing_type):
     if control_type == 'bc_control' and bc_type == 'neumann_bc':
         pytest.skip("Cannot use Neumann BCs as control")
 
@@ -92,7 +105,7 @@ def test_nlvs_adjoint(control_type, bc_type):
     PETSc.Sys.Print("record tape")
     continue_annotation()
     with set_working_tape() as tape:
-        J = forward(ic0, dt0, nt, bc_arg=bc_arg0)
+        J = forward(ic0, dt0, nt, bc_arg=bc_arg0, forcing_type=forcing_type)
         Jhat = ReducedFunctional(J, Control(control), tape=tape)
     pause_annotation()
 
@@ -124,7 +137,7 @@ def test_nlvs_adjoint(control_type, bc_type):
 
     # recompute component
     PETSc.Sys.Print("recompute test")
-    assert abs(Jhat(m) - forward(ic2, dt2, nt, bc_arg=bc_arg2)) < 1e-14
+    assert abs(Jhat(m) - forward(ic2, dt2, nt, bc_arg=bc_arg2, forcing_type=forcing_type)) < 1e-14
 
     # tlm
     PETSc.Sys.Print("tlm test")