diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py index a2d8582aa6..3c47609157 100644 --- a/firedrake/adjoint_utils/blocks/solving.py +++ b/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 diff --git a/firedrake/adjoint_utils/function.py b/firedrake/adjoint_utils/function.py index a4b83b9828..a0a98b0f1f 100644 --- a/firedrake/adjoint_utils/function.py +++ b/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): diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py index 3c8f9ce49a..054e22d80d 100644 --- a/firedrake/adjoint_utils/variational_solver.py +++ b/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) diff --git a/tests/firedrake/adjoint/test_hessian.py b/tests/firedrake/adjoint/test_hessian.py index f6942d9cbe..698a9ccf1b 100644 --- a/tests/firedrake/adjoint/test_hessian.py +++ b/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) diff --git a/tests/firedrake/adjoint/test_nlvs.py b/tests/firedrake/adjoint/test_nlvs.py index fb6d3f6701..498beff7df 100644 --- a/tests/firedrake/adjoint/test_nlvs.py +++ b/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,10 +26,20 @@ 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) @@ -37,10 +47,12 @@ def forward(ic, dt, nt, bc_arg=None): 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")