diff --git a/docs/notebooks/06-pde-constrained-optimisation.py b/docs/notebooks/06-pde-constrained-optimisation.py index 2e38da54e6..e57917fb7d 100644 --- a/docs/notebooks/06-pde-constrained-optimisation.py +++ b/docs/notebooks/06-pde-constrained-optimisation.py @@ -203,12 +203,13 @@ print("Jhat(g) = %.8g\nJhat(g_opt) = %.8g" % (Jhat(g), Jhat(g_opt))) # %% [markdown] -# To see the optimised flow field, we solve the same problem again, only with the new (optimised) value for the boundary data on $\Gamma_\text{circ}$. This time we're not interested in annotating the solve, so we tell `firedrake-adjoint` to ignore it by passing `annotate=False`. +# To see the optimised flow field, we solve the same problem again, only with the new (optimised) value for the boundary data on $\Gamma_\text{circ}$. This time we're not interested in annotating the solve, so we tell `firedrake-adjoint` to ignore it by running it in a `stop_annotating()` context. # %% g.assign(g_opt) w_opt = Function(W) -solve(a == L, w_opt, bcs=bcs, annotate=False) +with stop_annotating(): + solve(a == L, w_opt, bcs=bcs) # %% tags=["nbval-ignore-output"] u_opt, p_opt = w_opt.subfunctions diff --git a/firedrake/adjoint_utils/__init__.py b/firedrake/adjoint_utils/__init__.py index e1a31c8f1b..e9096735a4 100644 --- a/firedrake/adjoint_utils/__init__.py +++ b/firedrake/adjoint_utils/__init__.py @@ -13,7 +13,7 @@ from firedrake.adjoint_utils.variational_solver import ( # noqa F401 NonlinearVariationalProblemMixin, NonlinearVariationalSolverMixin ) -from firedrake.adjoint_utils.solving import annotate_solve, get_solve_blocks # noqa F401 +from firedrake.adjoint_utils.solving import get_solve_blocks # noqa F401 from firedrake.adjoint_utils.mesh import MeshGeometryMixin # noqa F401 from firedrake.adjoint_utils.checkpointing import ( # noqa F401 enable_disk_checkpointing, disk_checkpointing, diff --git a/firedrake/adjoint_utils/blocks/__init__.py b/firedrake/adjoint_utils/blocks/__init__.py index 3bbc926ce9..fa10ff29fa 100644 --- a/firedrake/adjoint_utils/blocks/__init__.py +++ b/firedrake/adjoint_utils/blocks/__init__.py @@ -1,7 +1,7 @@ from firedrake.adjoint_utils.blocks.assembly import AssembleBlock # noqa F401 from firedrake.adjoint_utils.blocks.solving import ( # noqa F401 - GenericSolveBlock, SolveLinearSystemBlock, ProjectBlock, - SupermeshProjectBlock, SolveVarFormBlock, + CachedSolverBlock, GenericSolveBlock, + ProjectBlock, SupermeshProjectBlock, SolveVarFormBlock, NonlinearVariationalSolveBlock ) from firedrake.adjoint_utils.blocks.function import ( # noqa F401 diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py index 1d46f8a488..a2d8582aa6 100644 --- a/firedrake/adjoint_utils/blocks/solving.py +++ b/firedrake/adjoint_utils/blocks/solving.py @@ -25,10 +25,290 @@ def extract_subfunction(u, V): return u -class Solver(Enum): +class SolverType(Enum): """Enum for solver types.""" FORWARD = 0 ADJOINT = 1 + TLM = 2 + HESSIAN = 3 + + +FORWARD = SolverType.FORWARD +ADJOINT = SolverType.ADJOINT +TLM = SolverType.TLM +HESSIAN = SolverType.HESSIAN + + +class CachedSolverBlock(Block): + def __init__( + self, forward_cache, tangent_cache, adjoint_cache, hessian_cache, ad_block_tag=None + ): + super().__init__(ad_block_tag=ad_block_tag) + + self.forward_cache = forward_cache + self.tangent_cache = tangent_cache + self.adjoint_cache = adjoint_cache + self.hessian_cache = hessian_cache + + # The adj_sol in the cached forms is shared by all blocks. + # This adj_sol belongs to this block specifically so we can + # stash the adjoint solution for the hessian calculation. + self.adj_sol = adjoint_cache.adj_sol.copy(deepcopy=True, annotate=False) + + def _coefficient_dependencies(self, dependencies=None): + dependencies = dependencies or self.get_dependencies() + return dependencies[:len(self.forward_cache.replaced_deps)] + + def _mesh_dependencies(self, dependencies=None): + dependencies = dependencies or self.get_dependencies() + len_replaced = len(self.forward_cache.replaced_deps) + len_meshes = len(self.forward_cache.meshes) + return dependencies[len_replaced:len_replaced + len_meshes] + + def _bc_dependencies(self, dependencies=None): + dependencies = dependencies or self.get_dependencies() + if len(self.forward_cache.bcs) > 0: + return dependencies[-len(self.forward_cache.bcs):] + else: + return [] + + def update_dependencies(self, use_output=False): + """Update all dependencies of the forward solve. + """ + # Update the coefficients in the form. + # Use the fact that zip will use the shorter length. + for replaced_dep, dep in zip(self.forward_cache.replaced_deps, + self._coefficient_dependencies()): + replaced_dep.assign(dep.saved_output) + + # 1. For forward recomputation the unknown Function should use + # the incoming value of the dependency as the initial guess. + # 2. For the adjoint, TLM, and Hessian, the unknown Function + # should use the computed value so that the linearised + # Jacobian is correct. + if use_output: + output = self.get_outputs()[0].saved_output + self.forward_cache.solver._problem.u.assign(output) + + # Update the boundary conditions + for replaced_dep, dep in zip(self.forward_cache.bcs, self._bc_dependencies()): + replaced_dep.set_value(dep.saved_output.function_arg) + + def update_tlm_dependencies(self): + """Update all dependencies of the tlm solve. + """ + for replaced_dep, dep in zip(self.tangent_cache.replaced_tlms, + self._coefficient_dependencies()): + if dep.output == self.forward_cache.func: + continue + if dep.tlm_value is None: # This dependency doesn't depend on the controls + continue + replaced_dep.assign(dep.tlm_value) + + for replaced_dep, dep in zip(self.tangent_cache.mesh_tlms, + self._mesh_dependencies()): + if dep.tlm_value is None: + continue + replaced_dep.assign(dep.tlm_value) + + for replaced_dep, dep in zip(self.forward_cache.bcs, self._bc_dependencies()): + if dep.tlm_value is None: # This dependency doesn't depend on the controls + bc_val = 0 + else: + bc_val = dep.tlm_value.function_arg + replaced_dep.set_value(bc_val) + + def update_adj_dependencies(self): + # TODO: Anything to do here? + pass + + def update_hessian_dependencies(self): + # TODO: Anything else to do here? + self.update_tlm_dependencies() + # update the adj_sol in the cached forms with + # the adj_sol value owned by this block. + self.hessian_cache.adj_sol.assign(self.adj_sol) + + def _compute_boundary(self, relevant_dependencies): + return any(isinstance(dep.output, firedrake.DirichletBC) + for _, dep in relevant_dependencies) + + def prepare_recompute_component(self, inputs, relevant_outputs): + return + + def recompute_component(self, inputs, block_variable, idx, prepared): + self.update_dependencies(use_output=False) + + solver = self.forward_cache.solver + solver.solve() + result = solver._problem.u.copy(deepcopy=True) + + # Possibly checkpoint the result for the adjoint solve later. + if isinstance(block_variable.checkpoint, firedrake.Function): + result = block_variable.checkpoint.assign(result) + + return maybe_disk_checkpoint(result) + + def prepare_evaluate_tlm(self, inputs, tlm_inputs, relevant_outputs): + return + + def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepared=None): + self.update_dependencies(use_output=True) + self.update_tlm_dependencies() + + # Assemble the rhs of (dF/du)(du/dm) = -dF/dm + tlm_rhs = self.tangent_cache.rhs + tlm_rhs.zero() + for dFdm, dep in zip(self.tangent_cache.dFdm_forms, self.get_dependencies()): + if dep.tlm_value is None: # This dependency doesn't depend on the controls + continue + if dep.output is self.forward_cache.func: # Can't compute dependence on initial guess + continue + tlm_rhs += firedrake.assemble(dFdm) + + # Solve for dudm + solver = self.tangent_cache.solver + solver._problem.u.zero() + solver.solve() + result = solver._problem.u.copy(deepcopy=True) + return result + + def solve_adj_equation(self, rhs, compute_boundary): + for bc in self.forward_cache.bcs: + bc.homogenize() + + adj_rhs = self.adjoint_cache.rhs + adj_sol = self.adjoint_cache.adj_sol + + adj_rhs.assign(rhs) + adj_sol.zero() + self.adjoint_cache.solver.solve() + + if compute_boundary: + adj_sol_bc = firedrake.assemble(self.adjoint_cache.residual) + adj_sol_bc = adj_sol_bc.riesz_representation("l2") + else: + adj_sol_bc = None + + return adj_sol.copy(deepcopy=True), adj_sol_bc + + def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies): + self.update_dependencies(use_output=True) + self.update_adj_dependencies() + + dJdu = adj_inputs[0] + + compute_boundary = self._compute_boundary(relevant_dependencies) + + adj_sol, adj_sol_bc = self.solve_adj_equation(dJdu, compute_boundary) + + # store adj_sol for Hessian computation later, or for inspecting + # adjoint sensitivities etc. + # self.adj_sol is owned by this block, whereas self.hessian_cache.adj_sol + # is shared between all blocks that the NLVS generates because it is the + # one in the cached forms, so we will update it as necessary if/when each + # block calculates the Hessian action. + self.adj_sol.assign(adj_sol) + + prepared = { + "adj_sol": adj_sol, + "adj_sol_bc": adj_sol_bc + } + return prepared + + def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepared=None): + if block_variable.output == self.forward_cache.func: + return None + + if isinstance(block_variable.output, firedrake.DirichletBC): + bc = block_variable.output + adj_sol_bc = prepared["adj_sol_bc"] + return bc.reconstruct( + g=extract_subfunction(adj_sol_bc, bc.function_space()) + ) + + # assemble sensititivy comment + dFdm = firedrake.assemble(self.adjoint_cache.dFdm_forms[idx]) + + return dFdm + + def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_dependencies): + self.update_dependencies(use_output=True) + self.update_hessian_dependencies() + + hessian_input = hessian_inputs[0] + tlm_output = self.get_outputs()[0].tlm_value + + if hessian_input is None: + return + if tlm_output is None: + return + + # 1. Assemble rhs + + # hessian input contribution + hessian_rhs = hessian_input.copy(deepcopy=True) + + # tlm_output contribution + self.hessian_cache.tlm_output.assign(tlm_output) + if not self.hessian_cache.d2Fdu2_form.empty(): + hessian_rhs -= firedrake.assemble(self.hessian_cache.d2Fdu2_form) + + # tlm_input contribution + for d2Fdmdu, dep in zip(self.hessian_cache.d2Fdmdu_forms, + self._coefficient_dependencies() + self._mesh_dependencies()): + if dep.tlm_value is None: # This dependency doesn't depend on the controls + continue + if dep.output is self.forward_cache.func: # Can't compute dependence on initial guess + continue + if len(d2Fdmdu.integrals()) > 0: + hessian_rhs -= firedrake.assemble(d2Fdmdu) + + # 2. Solve adjoint system + compute_boundary = self._compute_boundary(relevant_dependencies) + adj2_sol, adj2_sol_bc = self.solve_adj_equation(hessian_rhs, compute_boundary) + + self.hessian_cache.adj2_sol.assign(adj2_sol) + + prepared = { + "adj2_sol": adj2_sol, + "adj2_sol_bc": adj2_sol_bc, + } + + return prepared + + def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs, block_variable, idx, relevant_dependencies, prepared=None): + m = block_variable.output + + if m is self.forward_cache.func: + return None + + if isinstance(m, firedrake.DirichletBC): + bc = block_variable.output + adj2_sol_bc = prepared["adj2_sol_bc"] + return bc.reconstruct( + g=extract_subfunction(adj2_sol_bc, bc.function_space()) + ) + + relevant_d2Fdm2_forms = [] + for i, dep in relevant_dependencies: + if i >= len(self._coefficient_dependencies() + self._mesh_dependencies()): + continue + if dep.tlm_value is None: + continue + if dep.output is self.forward_cache.func: + continue + relevant_d2Fdm2_forms.append(self.hessian_cache.d2Fdm2_adj_forms[idx][i]) + + hessian_output = 0 + + for form in (self.hessian_cache.d2Fdudm_forms[idx], + self.hessian_cache.dFdm_adj2_forms[idx], + *relevant_d2Fdm2_forms): + if not form.empty(): + hessian_output += firedrake.assemble(form) + + return hessian_output class GenericSolveBlock(Block): @@ -56,6 +336,8 @@ def __init__(self, lhs, rhs, func, bcs, *args, **kwargs): # Solution function self.func = func self.function_space = self.func.function_space() + # Storage for adjoint solution of this block + self.adj_state_buf = func.copy(deepcopy=True) # Boundary conditions self.bcs = [] if bcs is not None: @@ -204,6 +486,8 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies): dFdu_form, dJdu, compute_bdy ) self.adj_sol = adj_sol + self.adj_state = adj_sol + self.adj_state_buf.assign(adj_sol) if self.adj_cb is not None: self.adj_cb(adj_sol) if self.adj_bdy_cb is not None and compute_bdy: @@ -215,31 +499,6 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies): r["adj_sol_bdy"] = adj_sol_bdy return r - def _assemble_dFdu_adj(self, dFdu_adj_form, **kwargs): - return firedrake.assemble(dFdu_adj_form, **kwargs) - - def _assemble_and_solve_adj_eq(self, dFdu_adj_form, dJdu, compute_bdy): - dJdu_copy = dJdu.copy() - # Homogenize and apply boundary conditions on adj_dFdu. - bcs = self._homogenize_bcs() - dFdu = firedrake.assemble(dFdu_adj_form, bcs=bcs, **self.assemble_kwargs) - - adj_sol = firedrake.Function(self.function_space) - firedrake.solve( - dFdu, adj_sol, dJdu, *self.adj_args, **self.adj_kwargs - ) - - adj_sol_bdy = None - if compute_bdy: - adj_sol_bdy = self._compute_adj_bdy( - adj_sol, adj_sol_bdy, dFdu_adj_form, dJdu_copy) - - return adj_sol, adj_sol_bdy - - def _compute_adj_bdy(self, adj_sol, adj_sol_bdy, dFdu_adj_form, dJdu): - adj_sol_bdy = firedrake.assemble(dJdu - firedrake.action(dFdu_adj_form, adj_sol)) - return adj_sol_bdy.riesz_representation("l2") - def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepared=None): if not self.linear and self.func == block_variable.output: @@ -283,7 +542,36 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, dFdm = firedrake.assemble(dFdm, **self.assemble_kwargs) return dFdm + def _assemble_dFdu_adj(self, dFdu_adj_form, **kwargs): + return firedrake.assemble(dFdu_adj_form, **kwargs) + + def _assemble_and_solve_adj_eq(self, dFdu_adj_form, dJdu, compute_bdy): + dJdu_copy = dJdu.copy() + # Homogenize and apply boundary conditions on adj_dFdu. + bcs = self._homogenize_bcs() + dFdu = firedrake.assemble(dFdu_adj_form, bcs=bcs, **self.assemble_kwargs) + + adj_sol = firedrake.Function(self.function_space) + firedrake.solve( + dFdu, adj_sol, dJdu, *self.adj_args, **self.adj_kwargs + ) + + adj_sol_bdy = None + if compute_bdy: + adj_sol_bdy = self._compute_adj_bdy( + adj_sol, adj_sol_bdy, dFdu_adj_form, dJdu_copy) + + return adj_sol, adj_sol_bdy + + def _compute_adj_bdy(self, adj_sol, adj_sol_bdy, dFdu_adj_form, dJdu): + adj_sol_bdy = firedrake.assemble(dJdu - firedrake.action(dFdu_adj_form, adj_sol)) + return adj_sol_bdy.riesz_representation("l2") + def prepare_evaluate_tlm(self, inputs, tlm_inputs, relevant_outputs): + pass + + def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, + prepared=None): fwd_block_variable = self.get_outputs()[0] u = fwd_block_variable.output @@ -295,16 +583,6 @@ def prepare_evaluate_tlm(self, inputs, tlm_inputs, relevant_outputs): fwd_block_variable.saved_output, firedrake.TrialFunction(u.function_space()) ) - - return { - "form": F_form, - "dFdu": dFdu - } - - def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, - prepared=None): - F_form = prepared["form"] - dFdu = prepared["dFdu"] V = self.get_outputs()[idx].output.function_space() bcs = [] @@ -341,10 +619,11 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, dFdm = ufl.algorithms.expand_derivatives(dFdm) dFdm = firedrake.assemble(dFdm) dudm = firedrake.Function(V) - return self._assemble_and_solve_tlm_eq( + result = self._assemble_and_solve_tlm_eq( firedrake.assemble(dFdu, bcs=bcs, **self.assemble_kwargs), dFdm, dudm, bcs ) + return result def _assemble_and_solve_tlm_eq(self, dFdu, dFdm, dudm, bcs): return self._assembled_solve(dFdu, dFdm, dudm, bcs) @@ -376,7 +655,10 @@ def _assemble_soa_eq_rhs(self, dFdu_form, adj_sol, hessian_input, d2Fdu2): elif not isinstance(c, firedrake.DirichletBC): dFdu_adj = firedrake.action(firedrake.adjoint(dFdu_form), adj_sol) - b_form += firedrake.derivative(dFdu_adj, c_rep, tlm_input) + # b_form += firedrake.derivative(dFdu_adj, c_rep, tlm_input) + bo_form = ufl.algorithms.expand_derivatives( + firedrake.derivative(dFdu_adj, c_rep, tlm_input)) + b_form += bo_form b_form = ufl.algorithms.expand_derivatives(b_form) if len(b_form.integrals()) > 0: @@ -404,6 +686,8 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, hessian_input = hessian_inputs[0] tlm_output = fwd_block_variable.tlm_value + self.adj_state = self.adj_state_buf.copy(deepcopy=True) + if hessian_input is None: return @@ -432,6 +716,7 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, r["adj_sol2_bdy"] = adj_sol2_bdy r["form"] = F_form r["adj_sol"] = adj_sol + return r def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs, @@ -577,6 +862,12 @@ def solve_init_params(self, args, kwargs, varform): ) self.adj_kwargs.pop("appctx", None) + if hasattr(self, "tlm_args") and len(self.tlm_args) <= 0: + self.tlm_args = self.adj_args + + if hasattr(self, "tlm_kwargs") and len(self.tlm_kwargs) <= 0: + self.tlm_kwargs = self.adj_kwargs.copy() + solver_params = kwargs.get("solver_parameters", None) if solver_params is not None and "mat_type" in solver_params: self.assemble_kwargs["mat_type"] = solver_params["mat_type"] @@ -586,25 +877,6 @@ def solve_init_params(self, args, kwargs, varform): self.assemble_kwargs["appctx"] = kwargs["appctx"] -class SolveLinearSystemBlock(GenericSolveBlock): - def __init__(self, A, u, b, *args, **kwargs): - lhs = A.form - func = u.function - rhs = b.form - bcs = A.bcs if hasattr(A, "bcs") else [] - super().__init__(lhs, rhs, func, bcs, *args, **kwargs) - - # Set up parameters initialization - self.ident_zeros_tol = \ - A.ident_zeros_tol if hasattr(A, "ident_zeros_tol") else None - self.assemble_system = \ - A.assemble_system if hasattr(A, "assemble_system") else False - - def _init_solver_parameters(self, args, kwargs): - super()._init_solver_parameters(args, kwargs) - solve_init_params(self, args, kwargs, varform=False) - - class SolveVarFormBlock(GenericSolveBlock): def __init__(self, equation, func, bcs=[], *args, **kwargs): lhs = equation.lhs @@ -627,6 +899,8 @@ def __init__(self, equation, func, bcs, adj_cache, problem_J, self.problem_J = problem_J self.solver_kwargs = solver_kwargs + self.adj_state_buf = func.copy(deepcopy=True) + super().__init__(lhs, rhs, func, bcs, **{**solver_kwargs, **kwargs}) if self.problem_J is not None: @@ -664,12 +938,12 @@ def _adjoint_solve(self, dJdu, compute_bdy): and self._ad_solvers["update_adjoint"] ): # Update left hand side of the adjoint equation. - self._ad_solver_replace_forms(Solver.ADJOINT) + self._ad_solver_replace_forms(SolverType.ADJOINT) self._ad_solvers["adjoint_lvs"].invalidate_jacobian() self._ad_solvers["update_adjoint"] = False elif not self._ad_solvers["forward_nlvs"]._problem._constant_jacobian: # Update left hand side of the adjoint equation. - self._ad_solver_replace_forms(Solver.ADJOINT) + self._ad_solver_replace_forms(SolverType.ADJOINT) # Update the right hand side of the adjoint equation. # problem.F._component[1] is the right hand side of the adjoint. @@ -687,7 +961,7 @@ def _adjoint_solve(self, dJdu, compute_bdy): return u_sol, adj_sol_bdy def _ad_assign_map(self, form, solver): - if solver == Solver.FORWARD: + if solver == SolverType.FORWARD: count_map = self._ad_solvers["forward_nlvs"]._problem._ad_count_map else: count_map = self._ad_solvers["adjoint_lvs"]._problem._ad_count_map @@ -705,7 +979,7 @@ def _ad_assign_map(self, form, solver): block_variable.saved_output if ( - solver == Solver.ADJOINT + solver == SolverType.ADJOINT and not self._ad_solvers["forward_nlvs"]._problem._constant_jacobian ): block_variable = self.get_outputs()[0] @@ -720,8 +994,8 @@ def _ad_assign_coefficients(self, form, solver): for coeff, value in assign_map.items(): coeff.assign(value) - def _ad_solver_replace_forms(self, solver=Solver.FORWARD): - if solver == Solver.FORWARD: + def _ad_solver_replace_forms(self, solver=SolverType.FORWARD): + if solver == SolverType.FORWARD: problem = self._ad_solvers["forward_nlvs"]._problem self._ad_assign_coefficients(problem.F, solver) self._ad_assign_coefficients(problem.J, solver) @@ -734,7 +1008,8 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies): relevant_dependencies ) adj_sol, adj_sol_bdy = self._adjoint_solve(adj_inputs[0], compute_bdy) - self.adj_sol = adj_sol + self.adj_state = adj_sol + self.adj_state_buf.assign(adj_sol) if self.adj_cb is not None: self.adj_cb(adj_sol) if self.adj_bdy_cb is not None and compute_bdy: diff --git a/firedrake/adjoint_utils/solving.py b/firedrake/adjoint_utils/solving.py index e9556e06ea..c7301288f4 100644 --- a/firedrake/adjoint_utils/solving.py +++ b/firedrake/adjoint_utils/solving.py @@ -1,71 +1,5 @@ -from functools import wraps -from pyadjoint.tape import get_working_tape, stop_annotating, annotate_tape - -from firedrake.adjoint_utils.blocks import SolveVarFormBlock, SolveLinearSystemBlock, GenericSolveBlock, ProjectBlock -import ufl - - -def annotate_solve(solve): - """This solve routine wraps the Firedrake :func:`.solve` call. Its purpose is to annotate the model, - recording what solves occur and what forms are involved, so that the adjoint and tangent linear models may be - constructed automatically by pyadjoint. - - To disable the annotation, just pass ``annotate=False`` to this routine, and it acts exactly like the - Firedrake solve call. This is useful in cases where the solve is known to be irrelevant or diagnostic - for the purposes of the adjoint computation (such as projecting fields to other function spaces - for the purposes of visualisation). - - The overloaded solve takes optional callback functions to extract adjoint solutions. - All of the callback functions follow the same signature, taking a single argument of type Function. - - Keyword Args: - adj_cb (:obj:`firedrake.function`, optional): - callback function supplying the adjoint solution in the interior. The boundary values are zero. - adj_bdy_cb (:obj:`firedrake.function`, optional): - callback function supplying the adjoint solution on the boundary. - The interior values are not guaranteed to be zero. - adj2_cb (:obj:`firedrake.function`, optional): - callback function supplying the second-order adjoint solution in the interior. - The boundary values are zero. - adj2_bdy_cb (:obj:`firedrake.function`, optional): - callback function supplying the second-order adjoint solution on - the boundary. The interior values are not guaranteed to be zero. - ad_block_tag (:obj:`string`, optional): - tag used to label the resulting block on the Pyadjoint tape. This - is useful for identifying which block is associated with which equation in the forward model. - - """ - - @wraps(solve) - def wrapper(*args, **kwargs): - - ad_block_tag = kwargs.pop("ad_block_tag", None) - annotate = annotate_tape(kwargs) - - if annotate: - tape = get_working_tape() - solve_block_type = SolveVarFormBlock - if not isinstance(args[0], ufl.equation.Equation): - solve_block_type = SolveLinearSystemBlock - - sb_kwargs = solve_block_type.pop_kwargs(kwargs) - sb_kwargs.update(kwargs) - block = solve_block_type(*args, ad_block_tag=ad_block_tag, **sb_kwargs) - tape.add_block(block) - - with stop_annotating(): - output = solve(*args, **kwargs) - - if annotate: - if hasattr(args[1], "create_block_variable"): - block_variable = args[1].create_block_variable() - else: - block_variable = args[1].function.create_block_variable() - block.add_output(block_variable) - - return output - - return wrapper +from pyadjoint.tape import get_working_tape +from firedrake.adjoint_utils.blocks import CachedSolverBlock, GenericSolveBlock, ProjectBlock def get_solve_blocks(): @@ -77,6 +11,6 @@ def get_solve_blocks(): return [ block for block in get_working_tape().get_blocks() - if issubclass(type(block), GenericSolveBlock) + if issubclass(type(block), (CachedSolverBlock, GenericSolveBlock)) and not issubclass(type(block), ProjectBlock) ] diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py index 21dded9839..3c8f9ce49a 100644 --- a/firedrake/adjoint_utils/variational_solver.py +++ b/firedrake/adjoint_utils/variational_solver.py @@ -1,9 +1,66 @@ import copy -from functools import wraps +from functools import wraps, cached_property from pyadjoint.tape import get_working_tape, stop_annotating, annotate_tape, no_annotations -from firedrake.adjoint_utils.blocks import NonlinearVariationalSolveBlock -from firedrake.ufl_expr import derivative, adjoint -from ufl import replace +from firedrake.adjoint_utils.blocks import ( + 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.algorithms import expand_derivatives +from ufl.domain import extract_domains +from types import SimpleNamespace +from collections import namedtuple + + +ForwardSolveRecomputeCache = namedtuple( + 'ForwardSolveRecomputeCache', + field_names=[ + "func", + "bcs", + "meshes", + "solver", + "replaced_deps", + ] +) + +AdjointSolveRecomputeCache = namedtuple( + 'AdjointSolveRecomputeCache', + field_names=[ + "adj_sol", + "rhs", + "dFdm_forms", + "residual", + "solver", + "dFdu_adj", + ] +) + +TangentSolveRecomputeCache = namedtuple( + "TangentSolveRecomputeCache", + field_names=[ + "tlm_val", + "rhs", + "dFdm_forms", + "solver", + "replaced_tlms", + "mesh_tlms", + "dFdu", + ] +) + +HessianSolveRecomputeCache = namedtuple( + "HessianSolveRecomputeCache", + field_names=[ + "adj_sol", + "adj2_sol", + "tlm_output", + "d2Fdu2_form", + "d2Fdmdu_forms", + "dFdm_adj2_forms", + "d2Fdm2_adj_forms", + "d2Fdudm_forms", + ] +) class NonlinearVariationalProblemMixin: @@ -52,10 +109,419 @@ def wrapper(self, problem, *args, **kwargs): "recompute_count": 0} self._ad_adj_cache = {} + self._ad_solver_cache = {} + + # process args/kwargs for cached solvers + self._ad_args_kwargs = SimpleNamespace( + forward_args=kwargs.pop("forward_args", []), + forward_kwargs=kwargs.pop("forward_kwargs", {}), + adj_args=kwargs.pop("adj_args", []), + adj_kwargs=kwargs.pop("adj_kwargs", {}), + tlm_args=kwargs.pop("tlm_args", []), + tlm_kwargs=kwargs.pop("tlm_kwargs", {}), + assemble_kwargs={} + ) + solve_init_params(self._ad_args_kwargs, + args, kwargs, varform=True) + return wrapper + @cached_property + @no_annotations + def _ad_forward_cache(self): + from firedrake import ( + DirichletBC, + MatrixBase, + NonlinearVariationalProblem, + NonlinearVariationalSolver) + + problem = self._ad_problem + + # Build a new form so that we can update the coefficient + # values without affecting user code. + # We do this by copying all coefficients in the form and + # symbolically replacing the old values with the new. + F = problem.F + replace_map = {} + for old_coeff in F.coefficients(): + new_coeff = old_coeff.copy(deepcopy=True) + replace_map[old_coeff] = new_coeff + + # We need a handle to the new Function being + # solved for so that we can create an NLVS. + Fnew = replace(F, replace_map) + unew = replace_map[problem.u] + + Jnew = replace(problem.J, replace_map) + Jpnew = None + if problem.Jp and problem.Jp is not problem.J: + Jpnew = replace(problem.Jp, replace_map) + + # We also need to "replace" all the bcs in + # the new NLVS so we can modify those values + # without affecting user code. + # Note that ``DirichletBC.reconstruct`` will + # return ``self`` if V, g, and sub_domain are + # all unchanged, so we need to explicitly + # instantiate a new object. + bcs = problem.bcs + bcs_new = [ + DirichletBC(V=bc.function_space(), + g=bc.function_arg, + sub_domain=bc.sub_domain) + for bc in bcs + ] + + # for a pre-assembled Jacobian (i.e. solving an + # assembled equation), we can't pass BCs into the + # solver, but we need to cache them for the TLM + # and adjoint solvers + bcs_fwd = [] + if not isinstance(Jnew, MatrixBase): + bcs_fwd = bcs_new + + # get all the unique meshes for domains on the form + meshes = set() + try: + for mesh in extract_domains(F): + meshes.add(mesh) + except AttributeError: + pass + meshes = list(meshes) + + # This NLVS will be used to recompute the solve. + # TODO: solver_parameters + nlvp = NonlinearVariationalProblem(Fnew, unew, J=Jnew, Jp=Jpnew, bcs=bcs_fwd) + nlvs = NonlinearVariationalSolver( + nlvp, + *self._ad_args_kwargs.forward_args, + **self._ad_args_kwargs.forward_kwargs) + + # The original coefficients will be added as + # dependencies to all solve blocks. + # The block need handles to the newly created + # objects to update their values when recomputing. + self._ad_dependencies_to_add = (*replace_map.keys(), *meshes, *bcs) + + return ForwardSolveRecomputeCache( + func=self._ad_problem.u, + bcs=bcs_new, + meshes=meshes, + solver=nlvs, + replaced_deps=tuple(replace_map.values()), + ) + + @cached_property + @no_annotations + def _ad_tangent_cache(self): + from firedrake import ( + Function, Cofunction, derivative, + LinearVariationalProblem, LinearVariationalSolver, + SpatialCoordinate, + ) + + # If we build the TLM form from the cached + # forward solve form then we can update exactly + # the same coefficients/boundary conditions. + nlvp = self._ad_forward_cache.solver._problem + + F = nlvp.F + u = nlvp.u + V = u.function_space() + + # We need gradient of output/input i.e. du/dm. + # We know F(u; m) = 0 and _total_ dF/dm = 0. + # Then for the _partial_ derivatives: + # (dF/du)*(du/dm) + dF/dm = 0 so we calculate: + # (dF/du)*(du/dm) = -dF/dm + dFdu = expand_derivatives(derivative(F, u)) + dFdm = Cofunction(V.dual()) + dudm = Function(V) + + # Reuse the same bcs as the forward problem. + # TODO: Think about if we should use new bcs. + # TODO: solver_parameters + lvp = LinearVariationalProblem(dFdu, dFdm, dudm, bcs=self._ad_forward_cache.bcs) + lvs = LinearVariationalSolver( + lvp, + *self._ad_args_kwargs.tlm_args, + **self._ad_args_kwargs.tlm_kwargs) + + tlm_rhs = dFdm + + # Do all the symbolic work for calculating dF/dm up front + # so we only pay for the numeric calculations at run time. + replaced_tlms = [] + dFdm_tlm_forms = [] + for m in self._ad_forward_cache.replaced_deps: + mtlm = m.copy(deepcopy=True) + replaced_tlms.append(mtlm) + + dFdm = derivative(-F, m, mtlm) + dFdm_tlm_forms.append(dFdm) + + mesh_tlms = [] + for m in self._ad_forward_cache.meshes: + mtlm = Function(m._ad_function_space()) + mesh_tlms.append(mtlm) + + X = SpatialCoordinate(m) + dFdm = derivative(-F, X, mtlm) + dFdm_tlm_forms.append(dFdm) + + tlm_val = Function(V) + + return TangentSolveRecomputeCache( + tlm_val=tlm_val, + rhs=tlm_rhs, + dFdm_forms=dFdm_tlm_forms, + solver=lvs, + replaced_tlms=replaced_tlms, + mesh_tlms=mesh_tlms, + dFdu=dFdu, + ) + + @cached_property + @no_annotations + def _ad_adjoint_cache(self): + from firedrake import ( + Function, Cofunction, TrialFunction, Argument, + LinearVariationalProblem, LinearVariationalSolver, + SpatialCoordinate, TestFunction, + ) + + # If we build the adjoint form from the cached + # forward solve form then we can update exactly + # the same coefficients/boundary conditions. + nlvp = self._ad_forward_cache.solver._problem + + F = nlvp.F + u = nlvp.u + V = u.function_space() + + # TODO: rewrite for adjoint not TLM + # We need gradient of output/input i.e. du/dm. + # We know F(u; m) = 0 and _total_ dF/dm = 0. + # Then for the _partial_ derivatives: + # (dF/du)*(du/dm) + dF/dm = 0 so we calculate: + # (dF/du)*(du/dm) = -dF/dm + dFdu = self._ad_tangent_cache.dFdu + try: + dFdu_adj = adjoint(dFdu) + except ValueError: + # Try again without expanding derivatives, + # as dFdu might have been simplied to an empty Form + dFdu_adj = adjoint(dFdu, derivatives_expanded=True) + + # This will be the rhs of the adjoint problem + dJdu = Cofunction(V.dual()) + adj_sol = Function(V) + + # Reuse the same bcs as the forward problem. + # TODO: Think about if we should use new bcs. + # TODO: solver_parameters + lvp = LinearVariationalProblem(dFdu_adj, dJdu, adj_sol, bcs=self._ad_forward_cache.bcs) + lvs = LinearVariationalSolver( + lvp, + *self._ad_args_kwargs.adj_args, + **self._ad_args_kwargs.adj_kwargs) + + # We'll need to assemble these forms to calculate + # the adj_component for each dependency. + # Do all the symbolic work for calculating dJ/du up front + # so we only pay for the numeric calculations at run time. + dFdm_adj_forms = [] + for m in self._ad_forward_cache.replaced_deps: + # Action of adjoint solution on dFdm + # TODO: Which of the two implementations should we use? + dFdm = derivative(-F, m, TrialFunction(m.function_space())) + + # 1. from previous cached implementation + dFdm = adjoint(dFdm) + if isinstance(dFdm, Argument): + # Corner case. Should be fixed more permanently upstream in UFL. + # See: https://github.com/FEniCS/ufl/issues/395 + dFdm = Action(dFdm, adj_sol) + else: + dFdm = dFdm * adj_sol + + # 2. from GenericSolveBlock + # if isinstance(dFdm, ufl.Form): + # dFdm = adjoint(dFdm) + # dFdm = action(dFdm, adj_sol) + # else: + # dFdm = dFdm(adj_sol) + + dFdm_adj_forms.append(dFdm) + + for m in self._ad_forward_cache.meshes: + X = SpatialCoordinate(m) + # we can't take the CoordinateDerivative of an Action, so we have + # to invert this form compared to the expression above + dFdm = derivative(action(-F, adj_sol), X, TestFunction(m._ad_function_space())) + dFdm_adj_forms.append(dFdm) + + # To calculate the adjoint component of each DirichletBC + # we'll need the residual of the adjoint equation without + # any DirichletBC using the solution calculated with + # homogeneous DirichletBCs. + adj_residual = dJdu - action(dFdu_adj, adj_sol) + + return AdjointSolveRecomputeCache( + adj_sol=adj_sol, + rhs=dJdu, + dFdm_forms=dFdm_adj_forms, + residual=adj_residual, + solver=lvs, + dFdu_adj=dFdu_adj, + ) + + @cached_property + @no_annotations + def _ad_hessian_cache(self): + from firedrake import ( + Function, TrialFunction, TestFunction, + SpatialCoordinate, MeshGeometry, + ) + + nlvp = self._ad_forward_cache.solver._problem + F = nlvp.F + u = nlvp.u + V = u.function_space() + + # Solution of the adjoint equation, requires evaluate_adj to + # have been called + adj_sol = Function(V) + # Solution of the TLM equation, requires evaluate_tlm to + # have been called (which is handled by the driver) + tlm_output = Function(V) + # Solution of the second-order adjoint equation, + # this is computed during prepare_evaluate_hessian + adj2_sol = Function(V) + + # 1. Forms to calculate rhs of Hessian solve + # Calculate d^2F*/du^2 * du/dm * dm + # where dm is direction for tlm action so du/dm * dm is tlm output + dFdu = self._ad_tangent_cache.dFdu + dFdu_adj = action(adjoint(dFdu), adj_sol) + + d2Fdu2 = derivative(dFdu_adj, u, tlm_output) + # expand_derivatives will simplify down to an empty + # form if required + d2Fdu2_form = expand_derivatives(d2Fdu2) + + # Contributions from each tlm_input + dFdu_adj = action(self._ad_adjoint_cache.dFdu_adj, adj_sol) + d2Fdmdu_forms = [] + for m, dm in zip(self._ad_forward_cache.replaced_deps, + self._ad_tangent_cache.replaced_tlms): + d2Fdmdu = expand_derivatives( + derivative(dFdu_adj, m, dm)) + + d2Fdmdu_forms.append(d2Fdmdu) + + for m, dm in zip(self._ad_forward_cache.meshes, + self._ad_tangent_cache.mesh_tlms): + X = SpatialCoordinate(m) + d2Fdmdu = expand_derivatives( + derivative(dFdu_adj, X, dm) + ) + + d2Fdmdu_forms.append(d2Fdmdu) + + # 2. Forms to calculate contribution from each control + dFdm_adj2_forms = [] + d2Fdm2_adj_forms = [] + d2Fdudm_forms = [] + for m in [*self._ad_forward_cache.replaced_deps, *self._ad_forward_cache.meshes]: + if isinstance(m, MeshGeometry): + X = SpatialCoordinate(m) + dm = TestFunction(m._ad_function_space()) + + F_adj = action(-F, adj_sol) + dFdm_adj = derivative(F_adj, X, dm) + + F_adj2 = action(-F, adj2_sol) + dFdm_adj2 = derivative(F_adj2, X, dm) + 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_adj = -expand_derivatives(action(adjoint(dFdm), adj_sol)) + dFdm_adj2 = -action(adjoint(dFdm), adj2_sol) + + dFdm_adj2_forms.append(dFdm_adj2) + + d2Fdudm = derivative(dFdm_adj, u, tlm_output) + d2Fdudm_forms.append(expand_derivatives(d2Fdudm)) + + 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) + + 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.append(d2Fdm2_adj_forms_k) + + return HessianSolveRecomputeCache( + adj_sol=adj_sol, + adj2_sol=adj2_sol, + tlm_output=tlm_output, + d2Fdu2_form=d2Fdu2_form, + d2Fdmdu_forms=d2Fdmdu_forms, + dFdm_adj2_forms=dFdm_adj2_forms, + d2Fdm2_adj_forms=d2Fdm2_adj_forms, + d2Fdudm_forms=d2Fdudm_forms, + ) + @staticmethod def _ad_annotate_solve(solve): + @wraps(solve) + def wrapper(self, **kwargs): + """To disable the annotation, just pass :py:data:`annotate=False` to this routine, and it acts exactly like the + Firedrake solve call. This is useful in cases where the solve is known to be irrelevant or diagnostic + for the purposes of the adjoint computation (such as projecting fields to other function spaces + for the purposes of visualisation).""" + annotate = annotate_tape(kwargs) + if annotate: + if kwargs.pop("bounds", None) is not None: + raise ValueError( + "MissingMathsError: we do not know how to differentiate through a variational inequality") + + block = CachedSolverBlock( + self._ad_forward_cache, + self._ad_tangent_cache, + self._ad_adjoint_cache, + self._ad_hessian_cache, + ad_block_tag=self.ad_block_tag) + + for dep in self._ad_dependencies_to_add: + block.add_dependency(dep, no_duplicates=True) + # mesh = self._ad_problem.u.function_space().mesh() + # block.add_dependency(mesh, no_duplicates=True) + + get_working_tape().add_block(block) + + with stop_annotating(): + out = solve(self, **kwargs) + + if annotate: + block.add_output(self._ad_problem._ad_u.create_block_variable()) + + return out + + return wrapper + + @staticmethod + def _ad_annotate_solve_old(solve): @wraps(solve) def wrapper(self, **kwargs): """To disable the annotation, just pass :py:data:`annotate=False` to this routine, and it acts exactly like the diff --git a/firedrake/solving.py b/firedrake/solving.py index 414870eba9..c4b80bd4f4 100644 --- a/firedrake/solving.py +++ b/firedrake/solving.py @@ -24,13 +24,11 @@ import firedrake.linear_solver as ls import firedrake.variational_solver as vs from firedrake.function import Function -from firedrake.adjoint_utils import annotate_solve from firedrake.petsc import PETSc from firedrake.utils import ScalarType @PETSc.Log.EventDecorator() -@annotate_solve def solve(*args, **kwargs): r"""Solve linear system Ax = b or variational problem a == L or F == 0. diff --git a/tests/firedrake/adjoint/test_assemble.py b/tests/firedrake/adjoint/test_assemble.py index 22bd7540a1..5e891619a6 100644 --- a/tests/firedrake/adjoint/test_assemble.py +++ b/tests/firedrake/adjoint/test_assemble.py @@ -89,9 +89,11 @@ def test_assemble_1_forms_tlm(rg): h = rg.uniform(V) g = f.copy(deepcopy=True) - f.block_variable.tlm_value = h - tape.evaluate_tlm() - assert (taylor_test(Jhat, g, h, dJdm=J.block_variable.tlm_value) > 1.9) + Jhat(g) + assert (taylor_test(Jhat, g, h, dJdm=Jhat.tlm(h)) > 1.9) + # f.block_variable.tlm_value = h + # tape.evaluate_tlm() + # assert (taylor_test(Jhat, g, h, dJdm=J.block_variable.tlm_value) > 1.9) @pytest.mark.skipcomplex diff --git a/tests/firedrake/adjoint/test_hessian.py b/tests/firedrake/adjoint/test_hessian.py index e3ec975c97..f6942d9cbe 100644 --- a/tests/firedrake/adjoint/test_hessian.py +++ b/tests/firedrake/adjoint/test_hessian.py @@ -124,7 +124,6 @@ def test_function(rg): J = assemble(c ** 2 * u ** 2 * dx) Jhat = ReducedFunctional(J, [control_c, control_f]) - dJdc, dJdf = compute_derivative(J, [control_c, control_f], apply_riesz=True) # Step direction for derivatives and convergence test h_c = Function(R, val=1.0) @@ -161,6 +160,7 @@ def test_nonlinear(rg): Jhat = ReducedFunctional(J, Control(f)) h = rg.uniform(V, 0, 10) + g = f.copy(deepcopy=True) J.block_variable.adj_value = 1.0 f.block_variable.tlm_value = h @@ -171,8 +171,6 @@ def test_nonlinear(rg): J.block_variable.hessian_value = 0 tape.evaluate_hessian() - g = f.copy(deepcopy=True) - dJdm = J.block_variable.tlm_value Hm = f.block_variable.hessian_value.dat.inner(h.dat) assert taylor_test(Jhat, g, h, dJdm=dJdm, Hm=Hm) > 2.8 @@ -223,25 +221,24 @@ def test_dirichlet(rg): def test_burgers(solve_type, rg): tape = Tape() set_working_tape(tape) - n = 100 - mesh = UnitIntervalMesh(n) - V = FunctionSpace(mesh, "CG", 2) + nx = 50 + nt = 5 + mesh = UnitIntervalMesh(nx) + V = FunctionSpace(mesh, "CG", 1) - def Dt(u, u_, timestep): - return (u - u_)/timestep + def Dt(u, u_, dt): + return (u - u_)/dt x, = SpatialCoordinate(mesh) - pr = project(sin(2*pi*x), V, annotate=False) - ic = Function(V).assign(pr) + ic = Function(V).project(sin(2*pi*x)) u_ = Function(V).assign(ic) u = Function(V).assign(ic) v = TestFunction(V) - nu = Constant(0.0001) + nu = Constant(1/100) - dt = 0.01 - nt = 20 + dt = Constant(1/nx) params = { 'snes_rtol': 1e-10, @@ -266,12 +263,6 @@ def Dt(u, u_, timestep): NonlinearVariationalProblem(F, u), solver_parameters=params) - if use_nlvs: - solver.solve() - else: - solve(F == 0, u, bc, solver_parameters=params) - u_.assign(u) - for _ in range(nt): if use_nlvs: solver.solve() @@ -282,16 +273,13 @@ def Dt(u, u_, timestep): J = assemble(u_*u_*dx + ic*ic*dx) Jhat = ReducedFunctional(J, Control(ic)) + h = rg.uniform(V) g = ic.copy(deepcopy=True) - J.block_variable.adj_value = 1.0 - ic.block_variable.tlm_value = h - tape.evaluate_adj() - tape.evaluate_tlm() - - J.block_variable.hessian_value = 0 - tape.evaluate_hessian() - dJdm = J.block_variable.tlm_value - Hm = ic.block_variable.hessian_value.dat.inner(h.dat) - assert taylor_test(Jhat, g, h, dJdm=dJdm, Hm=Hm) > 2.9 + taylor = taylor_to_dict(Jhat, g, h) + from pprint import pprint + pprint(taylor) + 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'] diff --git a/tests/firedrake/adjoint/test_nlvs.py b/tests/firedrake/adjoint/test_nlvs.py new file mode 100644 index 0000000000..9eb7c7a005 --- /dev/null +++ b/tests/firedrake/adjoint/test_nlvs.py @@ -0,0 +1,193 @@ +import pytest + +from firedrake import * +from firedrake.adjoint 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(): + if not annotate_tape(): + continue_annotation() + yield + # Ensure annotation is paused when we finish. + if annotate_tape(): + pause_annotation() + + +def forward(ic, dt, nt, bc_arg=None): + """Burgers equation solver.""" + V = ic.function_space() + + if bc_arg: + bc_val = bc_arg.copy(deepcopy=True) + bcs = [DirichletBC(V, bc_val, 1), + DirichletBC(V, 0, 2)] + else: + bcs = None + + nu = Function(dt.function_space()).assign(0.1) + + u0 = Function(V) + u1 = Function(V) + v = TestFunction(V) + + F = ((u1 - u0)*v + + dt*u1*u1.dx(0)*v + + dt*nu*u1.dx(0)*v.dx(0))*dx + + problem = NonlinearVariationalProblem(F, u1, bcs=bcs) + solver = NonlinearVariationalSolver(problem) + + u1.assign(ic) + + for i in range(nt): + u0.assign(u1) + solver.solve() + nu += dt + # if bc_arg: + # bc_val.assign(bc_val + dt/nt) + + J = assemble(u1*u1*dx) + return J + + +@pytest.mark.skipcomplex +@pytest.mark.parametrize("control_type", ["ic_control", + "dt_control", + "bc_control"]) +@pytest.mark.parametrize("bc_type", ["neumann_bc", + "dirichlet_bc"]) +def test_nlvs_adjoint(control_type, bc_type): + if control_type == 'bc_control' and bc_type == 'neumann_bc': + pytest.skip("Cannot use Neumann BCs as control") + + nx = 100 + nt = 5 + + mesh = UnitIntervalMesh(nx) + x, = SpatialCoordinate(mesh) + + V = FunctionSpace(mesh, "CG", 1) + R = FunctionSpace(mesh, "R", 0) + + dt = Function(R).assign(2/nx) + ic = Function(V).interpolate(cos(2*pi*x)) + + dt0 = dt.copy(deepcopy=True) + ic0 = ic.copy(deepcopy=True) + + if bc_type == 'neumann_bc': + bc_arg = None + bc_arg0 = None + elif bc_type == 'dirichlet_bc': + bc_arg = Function(R).assign(1.) + bc_arg0 = bc_arg.copy(deepcopy=True) + else: + raise ValueError(f"Unrecognised {bc_type=}") + + if control_type == 'ic_control': + control = ic0 + elif control_type == 'dt_control': + control = dt0 + elif control_type == 'bc_control': + control = bc_arg0 + else: + raise ValueError(f"Unrecognised {control_type=}") + + PETSc.Sys.Print("record tape") + continue_annotation() + with set_working_tape() as tape: + J = forward(ic0, dt0, nt, bc_arg=bc_arg0) + Jhat = ReducedFunctional(J, Control(control), tape=tape) + pause_annotation() + + if control_type == 'ic_control': + m = Function(V).assign(0.5*ic) + h = Function(V).interpolate(-0.5*cos(4*pi*x)) + + ic2 = m.copy(deepcopy=True) + dt2 = dt + bc_arg2 = bc_arg + + elif control_type == 'dt_control': + m = Function(R).assign(0.05) + h = Function(R).assign(0.01) + + ic2 = ic + dt2 = m.copy(deepcopy=True) + bc_arg2 = bc_arg + + elif control_type == 'bc_control': + m = Function(R).assign(0.5) + h = Function(R).assign(-0.1) + + ic2 = ic + dt2 = dt + bc_arg2 = m.copy(deepcopy=True) + + from mpi4py import MPI + + # recompute component + PETSc.Sys.Print("recompute test") + assert abs(Jhat(m) - forward(ic2, dt2, nt, bc_arg=bc_arg2)) < 1e-14 + + # tlm + PETSc.Sys.Print("tlm test") + Jhat(m) + assert taylor_test(Jhat, m, h, dJdm=Jhat.tlm(h)) > 1.95 + + # adjoint + PETSc.Sys.Print("adjoint test") + Jhat(m) + assert taylor_test(Jhat, m, h) > 1.95 + + # hessian + PETSc.Sys.Print("hessian test") + Jhat(m) + taylor = taylor_to_dict(Jhat, m, h) + from pprint import pprint + pprint(taylor) + + assert min(taylor['R0']['Rate']) > 0.95 + assert min(taylor['R1']['Rate']) > 1.95 + assert min(taylor['R2']['Rate']) > 2.95 + + return + + for _ in range(3): + stime = MPI.Wtime() + Jhat(m) + etime = MPI.Wtime() + PETSc.Sys.Print(f"Recompute time: {etime - stime:.4f}") + + for _ in range(3): + stime = MPI.Wtime() + Jhat.derivative() + etime = MPI.Wtime() + PETSc.Sys.Print(f"Derivative time: {etime - stime:.4f}") + + for _ in range(3): + stime = MPI.Wtime() + Jhat.tlm(h) + etime = MPI.Wtime() + PETSc.Sys.Print(f"TLM time: {etime - stime:.4f}") + + for _ in range(3): + stime = MPI.Wtime() + Jhat.hessian(h, evaluate_tlm=False) + etime = MPI.Wtime() + PETSc.Sys.Print(f"Hessian time: {etime - stime:.4f}") + + +if __name__ == "__main__": + control_type = "ic_control" + bc_type = "neumann_bc" + PETSc.Sys.Print(f"{control_type=} | {bc_type=}") + test_nlvs_adjoint(control_type, bc_type) diff --git a/tests/firedrake/adjoint/test_solving.py b/tests/firedrake/adjoint/test_solving.py index 05b114ad0e..d3bca15d39 100644 --- a/tests/firedrake/adjoint/test_solving.py +++ b/tests/firedrake/adjoint/test_solving.py @@ -42,6 +42,37 @@ def J(f): _test_adjoint(J, f, rg) +@pytest.mark.skipcomplex +def test_assembled_problem(rg): + assert len(get_working_tape()._blocks) == 0 + mesh = IntervalMesh(10, 0, 1) + V = FunctionSpace(mesh, "Lagrange", 1) + R = FunctionSpace(mesh, "R", 0) + f = Function(V).assign(1.) + + u = TrialFunction(V) + u_ = Function(V) + v = TestFunction(V) + bc = DirichletBC(V, Function(R, val=1), "on_boundary") + + def J(f): + a = inner(grad(u), grad(v))*dx + L = f*v*dx + + A = assemble(a, bcs=bc) + b = assemble(L) + + solve(A, u_, b) + return assemble(u_**2*dx) + + # this tests a bug in recompute() if bcs= keyword argument is provided in solve + J0 = J(f) + rf = ReducedFunctional(J0, Control(f)) + assert_approx_equal(rf(f), J0) + assert rf.tape.recompute_count == 1 + _test_adjoint(J, f, rg) + + @pytest.mark.skipcomplex def test_singular_linear_problem(rg): """This tests whether nullspace and solver_parameters are passed on in adjoint solves"""