diff --git a/firedrake/adjoint_utils/__init__.py b/firedrake/adjoint_utils/__init__.py index 7d78643036..80fe9386fa 100644 --- a/firedrake/adjoint_utils/__init__.py +++ b/firedrake/adjoint_utils/__init__.py @@ -9,7 +9,9 @@ FunctionMixin, CofunctionMixin ) from firedrake.adjoint_utils.assembly import annotate_assemble # noqa F401 -from firedrake.adjoint_utils.projection import annotate_project, ProjectorMixin # noqa F401 +from firedrake.adjoint_utils.projection import ( # noqa F401 + annotate_project, annotate_super_project, ProjectorMixin +) from firedrake.adjoint_utils.variational_solver import ( # noqa F401 NonlinearVariationalProblemMixin, NonlinearVariationalSolverMixin ) diff --git a/firedrake/adjoint_utils/blocks/__init__.py b/firedrake/adjoint_utils/blocks/__init__.py index fa10ff29fa..9701c9545d 100644 --- a/firedrake/adjoint_utils/blocks/__init__.py +++ b/firedrake/adjoint_utils/blocks/__init__.py @@ -1,8 +1,6 @@ from firedrake.adjoint_utils.blocks.assembly import AssembleBlock # noqa F401 from firedrake.adjoint_utils.blocks.solving import ( # noqa F401 - CachedSolverBlock, GenericSolveBlock, - ProjectBlock, SupermeshProjectBlock, SolveVarFormBlock, - NonlinearVariationalSolveBlock + NonlinearVariationalSolveBlock, ProjectBlock, SupermeshProjectBlock ) from firedrake.adjoint_utils.blocks.function import ( # noqa F401 FunctionAssignBlock, FunctionMergeBlock, SubfunctionBlock diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py index b2f22f5d9c..b5683e57ab 100644 --- a/firedrake/adjoint_utils/blocks/solving.py +++ b/firedrake/adjoint_utils/blocks/solving.py @@ -1,12 +1,6 @@ -import numpy -import ufl -from ufl.domain import extract_domains, extract_unique_domain -from ufl import replace -from ufl.formatting.ufl2unicode import ufl2unicode from enum import Enum -from pyadjoint import Block, stop_annotating, get_working_tape -from pyadjoint.enlisting import Enlist +from pyadjoint import Block, stop_annotating import firedrake from firedrake.adjoint_utils.checkpointing import maybe_disk_checkpoint @@ -39,7 +33,7 @@ class SolverType(Enum): HESSIAN = SolverType.HESSIAN -class CachedSolverBlock(Block): +class NonlinearVariationalSolveBlock(Block): def __init__( self, forward_cache, tangent_cache, adjoint_cache, hessian_cache, ad_block_tag=None ): @@ -305,520 +299,6 @@ def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs, block_v return hessian_output -class GenericSolveBlock(Block): - pop_kwargs_keys = ["adj_cb", "adj_bdy_cb", "adj2_cb", "adj2_bdy_cb", - "forward_args", "forward_kwargs", "adj_args", - "adj_kwargs"] - - def __init__(self, lhs, rhs, func, bcs, *args, **kwargs): - super().__init__(ad_block_tag=kwargs.pop('ad_block_tag', None)) - self.adj_cb = kwargs.pop("adj_cb", None) - self.adj_bdy_cb = kwargs.pop("adj_bdy_cb", None) - self.adj2_cb = kwargs.pop("adj2_cb", None) - self.adj2_bdy_cb = kwargs.pop("adj2_bdy_cb", None) - - self.forward_args = [] - self.forward_kwargs = {} - self.adj_args = [] - self.adj_kwargs = {} - self.assemble_kwargs = {} - - # Equation LHS - self.lhs = lhs - # Equation RHS - self.rhs = rhs - # Solution function - self.func = func - self.function_space = self.func.function_space() - # Boundary conditions - self.bcs = [] - if bcs is not None: - self.bcs = Enlist(bcs) - - if isinstance(self.lhs, ufl.Form) and isinstance(self.rhs, (ufl.Form, ufl.Cofunction)): - self.linear = True - for c in self.rhs.coefficients(): - self.add_dependency(c, no_duplicates=True) - else: - self.linear = False - - for c in self.lhs.coefficients(): - self.add_dependency(c, no_duplicates=True) - - for bc in self.bcs: - self.add_dependency(bc, no_duplicates=True) - - try: # add all meshes as dependency - for mesh in extract_domains(self.lhs): - self.add_dependency(mesh, no_duplicates=True) - except AttributeError: - pass - - if isinstance(self.rhs, ufl.BaseForm): - # add all meshes as dependency - for mesh in extract_domains(self.rhs): - self.add_dependency(mesh, no_duplicates=True) - - self._init_solver_parameters(args, kwargs) - - def _init_solver_parameters(self, args, kwargs): - self.forward_args = kwargs.pop("forward_args", []) - self.forward_kwargs = kwargs.pop("forward_kwargs", {}) - self.adj_args = kwargs.pop("adj_args", []) - self.adj_kwargs = kwargs.pop("adj_kwargs", {}) - self.assemble_kwargs = {} - - def __str__(self): - try: - lhs_string = ufl2unicode(self.lhs) - except AttributeError: - lhs_string = str(self.lhs) - try: - rhs_string = ufl2unicode(self.rhs) - except AttributeError: - rhs_string = str(self.rhs) - return "solve({} = {})".format(lhs_string, rhs_string) - - def _create_F_form(self): - # Process the equation forms, replacing values with checkpoints, - # and gathering lhs and rhs in one single form. - if self.linear: - tmp_u = firedrake.Function(self.function_space) - F_form = firedrake.action(self.lhs, tmp_u) - self.rhs - else: - tmp_u = self.func - F_form = self.lhs - - replace_map = self._replace_map(F_form) - replace_map[tmp_u] = self.get_outputs()[0].saved_output - return ufl.replace(F_form, replace_map) - - def _homogenize_bcs(self): - bcs = [] - for bc in self.bcs: - if isinstance(bc, firedrake.DirichletBC): - bc = bc.reconstruct(g=0) - bcs.append(bc) - return bcs - - def _create_initial_guess(self): - return firedrake.Function(self.function_space) - - def _recover_bcs(self): - bcs = [] - for block_variable in self.get_dependencies(): - c = block_variable.output - c_rep = block_variable.saved_output - - if isinstance(c, firedrake.DirichletBC): - bcs.append(c_rep) - return bcs - - def _replace_map(self, form): - replace_coeffs = {} - for block_variable in self.get_dependencies(): - coeff = block_variable.output - if coeff in form.coefficients(): - replace_coeffs[coeff] = block_variable.saved_output - return replace_coeffs - - def _replace_form(self, form, func=None): - """Replace the form coefficients with checkpointed values - - func represents the initial guess if relevant. - """ - replace_map = self._replace_map(form) - if func is not None and self.func in replace_map: - firedrake.Function.assign(func, replace_map[self.func]) - replace_map[self.func] = func - return ufl.replace(form, replace_map) - - def _should_compute_boundary_adjoint(self, relevant_dependencies): - # Check if DirichletBC derivative is relevant - bdy = False - for _, dep in relevant_dependencies: - if isinstance(dep.output, firedrake.DirichletBC): - bdy = True - break - return bdy - - @property - def adj_sol(self): - return self.adj_state - - @adj_sol.setter - def adj_sol(self, value): - if self.adj_state is None: - self.adj_state = value.copy(deepcopy=True) - else: - self.adj_state.assign(value) - - def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies): - fwd_block_variable = self.get_outputs()[0] - u = fwd_block_variable.output - - dJdu = adj_inputs[0] - - F_form = self._create_F_form() - - dFdu = firedrake.derivative( - F_form, - fwd_block_variable.saved_output, - firedrake.TrialFunction( - u.function_space() - ) - ) - dFdu_form = firedrake.adjoint(dFdu) - dJdu = dJdu.copy() - - compute_bdy = self._should_compute_boundary_adjoint( - relevant_dependencies - ) - adj_sol, adj_sol_bdy = self._assemble_and_solve_adj_eq( - dFdu_form, dJdu, compute_bdy - ) - self.adj_sol = 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: - self.adj_bdy_cb(adj_sol_bdy) - - r = {} - r["form"] = F_form - r["adj_sol"] = adj_sol - r["adj_sol_bdy"] = adj_sol_bdy - return r - - def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, - prepared=None): - if not self.linear and self.func == block_variable.output: - # We are not able to calculate derivatives wrt initial guess. - return None - F_form = prepared["form"] - adj_sol = prepared["adj_sol"] - adj_sol_bdy = prepared["adj_sol_bdy"] - c = block_variable.output - c_rep = block_variable.saved_output - - if isinstance(c, (firedrake.Function, firedrake.Cofunction)): - trial_function = firedrake.TrialFunction(c.function_space()) - elif isinstance(c, firedrake.DirichletBC): - tmp_bc = c.reconstruct( - g=extract_subfunction(adj_sol_bdy, c.function_space()) - ) - return [tmp_bc] - elif isinstance(c, firedrake.MeshGeometry): - # Using CoordinateDerivative requires us to do action before - # differentiating, might change in the future. - F_form_tmp = firedrake.action(F_form, adj_sol) - X = firedrake.SpatialCoordinate(c_rep) - dFdm = firedrake.derivative( - -F_form_tmp, X, - firedrake.TestFunction(c._ad_function_space()) - ) - - if dFdm == 0: - return firedrake.Function(c._ad_function_space().dual()) - - dFdm = firedrake.assemble(dFdm, **self.assemble_kwargs) - return dFdm - - dFdm = -firedrake.derivative(F_form, c_rep, trial_function) - dFdm = firedrake.adjoint(dFdm) - dFdm = firedrake.action(dFdm, adj_sol) - 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 - - F_form = self._create_F_form() - - # Obtain dFdu. - dFdu = firedrake.derivative( - F_form, - fwd_block_variable.saved_output, - firedrake.TrialFunction(u.function_space()) - ) - V = self.get_outputs()[idx].output.function_space() - - bcs = [] - dFdm = 0. - for block_variable in self.get_dependencies(): - tlm_value = block_variable.tlm_value - c = block_variable.output - c_rep = block_variable.saved_output - - if isinstance(c, firedrake.DirichletBC): - if tlm_value is None: - bcs.append(c.reconstruct(g=0)) - else: - bcs.append(tlm_value) - continue - elif isinstance(c, firedrake.MeshGeometry): - X = firedrake.SpatialCoordinate(c) - c_rep = X - - if tlm_value is None: - continue - - if c == self.func and not self.linear: - continue - - dFdm += firedrake.derivative(-F_form, c_rep, tlm_value) - - if isinstance(dFdm, float): - v = dFdu.arguments()[0] - dFdm = firedrake.inner( - firedrake.Constant(numpy.zeros(v.ufl_shape)), v - ) * firedrake.dx - - dFdm = ufl.algorithms.expand_derivatives(dFdm) - dFdm = firedrake.assemble(dFdm) - dudm = firedrake.Function(V) - 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) - - def _assemble_soa_eq_rhs(self, dFdu_form, adj_sol, hessian_input, d2Fdu2): - # Start piecing together the rhs of the soa equation - b = hessian_input.copy() - if not d2Fdu2.empty(): - b_form = firedrake.action(firedrake.adjoint(d2Fdu2), adj_sol) - else: - b_form = d2Fdu2 - - for bo in self.get_dependencies(): - c = bo.output - c_rep = bo.saved_output - tlm_input = bo.tlm_value - - if (c == self.func and not self.linear) or tlm_input is None: - continue - - if isinstance(c, firedrake.MeshGeometry): - X = firedrake.SpatialCoordinate(c) - dFdu_adj = firedrake.action(firedrake.adjoint(dFdu_form), - adj_sol) - d2Fdudm = ufl.algorithms.expand_derivatives( - firedrake.derivative(dFdu_adj, X, tlm_input)) - if not d2Fdudm.empty(): - b_form += d2Fdudm - 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) - 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 not b_form.empty(): - b -= firedrake.assemble(b_form) - - return b - - def _assemble_and_solve_soa_eq(self, dFdu_form, adj_sol, hessian_input, - d2Fdu2, compute_bdy): - b = self._assemble_soa_eq_rhs(dFdu_form, adj_sol, hessian_input, - d2Fdu2) - dFdu_form = firedrake.adjoint(dFdu_form) - adj_sol2, adj_sol2_bdy = self._assemble_and_solve_adj_eq(dFdu_form, b, - compute_bdy) - if self.adj2_cb is not None: - self.adj2_cb(adj_sol2) - if self.adj2_bdy_cb is not None and compute_bdy: - self.adj2_bdy_cb(adj_sol2_bdy) - return adj_sol2, adj_sol2_bdy - - def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, - relevant_dependencies): - # First fetch all relevant values - fwd_block_variable = self.get_outputs()[0] - hessian_input = hessian_inputs[0] - tlm_output = fwd_block_variable.tlm_value - - if hessian_input is None: - return - - if tlm_output is None: - return - - F_form = self._create_F_form() - - # Using the equation Form derive dF/du, d^2F/du^2 * du/dm * direction. - dFdu_form = firedrake.derivative(F_form, - fwd_block_variable.saved_output) - d2Fdu2 = ufl.algorithms.expand_derivatives( - firedrake.derivative(dFdu_form, fwd_block_variable.saved_output, - tlm_output)) - - adj_sol = self.adj_sol - if adj_sol is None: - raise RuntimeError("Hessian computation was run before adjoint.") - bdy = self._should_compute_boundary_adjoint(relevant_dependencies) - adj_sol2, adj_sol2_bdy = self._assemble_and_solve_soa_eq( - dFdu_form, adj_sol, hessian_input, d2Fdu2, bdy - ) - - r = {} - r["adj_sol2"] = adj_sol2 - 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, - block_variable, idx, relevant_dependencies, - prepared=None): - c = block_variable.output - if c == self.func and not self.linear: - return None - - adj_sol2 = prepared["adj_sol2"] - adj_sol2_bdy = prepared["adj_sol2_bdy"] - F_form = prepared["form"] - adj_sol = prepared["adj_sol"] - fwd_block_variable = self.get_outputs()[0] - tlm_output = fwd_block_variable.tlm_value - - c_rep = block_variable.saved_output - - # If m = DirichletBC then d^2F(u,m)/dm^2 = 0 and d^2F(u,m)/dudm = 0, - # so we only have the term dF(u,m)/dm * adj_sol2 - if isinstance(c, firedrake.DirichletBC): - tmp_bc = c.reconstruct( - g=extract_subfunction(adj_sol2_bdy, c.function_space()) - ) - return [tmp_bc] - - if isinstance(c, firedrake.MeshGeometry): - X = firedrake.SpatialCoordinate(c) - W = c._ad_function_space() - else: - W = c.function_space() - - dc = firedrake.TestFunction(W) - form_adj = firedrake.action(F_form, adj_sol) - form_adj2 = firedrake.action(F_form, adj_sol2) - if isinstance(c, firedrake.MeshGeometry): - dFdm_adj = firedrake.derivative(form_adj, X, dc) - dFdm_adj2 = firedrake.derivative(form_adj2, X, dc) - else: - dFdm_adj = firedrake.derivative(form_adj, c_rep, dc) - dFdm_adj2 = firedrake.derivative(form_adj2, c_rep, dc) - - # TODO: Old comment claims this might break on split. Confirm if true - # or not. - d2Fdudm = ufl.algorithms.expand_derivatives( - firedrake.derivative(dFdm_adj, fwd_block_variable.saved_output, - tlm_output)) - - d2Fdm2 = 0 - # We need to add terms from every other dependency - # i.e. the terms d^2F/dm_1dm_2 - for _, bv in relevant_dependencies: - c2 = bv.output - c2_rep = bv.saved_output - - if isinstance(c2, firedrake.DirichletBC): - continue - - tlm_input = bv.tlm_value - if tlm_input is None: - continue - - if c2 == self.func and not self.linear: - continue - - # TODO: If tlm_input is a Sum, this crashes in some instances? - if isinstance(c2_rep, firedrake.MeshGeometry): - X = firedrake.SpatialCoordinate(c2_rep) - d2Fdm2 += ufl.algorithms.expand_derivatives( - firedrake.derivative(dFdm_adj, X, tlm_input) - ) - else: - d2Fdm2 += ufl.algorithms.expand_derivatives( - firedrake.derivative(dFdm_adj, c2_rep, tlm_input) - ) - - hessian_form = ufl.algorithms.expand_derivatives( - d2Fdm2 + dFdm_adj2 + d2Fdudm - ) - hessian_output = 0 - if not hessian_form.empty(): - hessian_output = firedrake.assemble(hessian_form) - hessian_output *= -1. - - return hessian_output - - def prepare_recompute_component(self, inputs, relevant_outputs): - return self._replace_recompute_form() - - def _replace_recompute_form(self): - func = self._create_initial_guess() - - bcs = self._recover_bcs() - lhs = self._replace_form(self.lhs, func=func) - rhs = 0 - if self.linear: - rhs = self._replace_form(self.rhs) - - return lhs, rhs, func, bcs - - def _forward_solve(self, lhs, rhs, func, bcs): - firedrake.solve(lhs == rhs, func, bcs, *self.forward_args, - **self.forward_kwargs) - return func - - def _assembled_solve(self, lhs, rhs, func, bcs, **kwargs): - firedrake.solve(lhs, func, rhs, **kwargs) - return func - - def recompute_component(self, inputs, block_variable, idx, prepared): - lhs, rhs, func, bcs = prepared - result = self._forward_solve(lhs, rhs, func, bcs) - if isinstance(block_variable.checkpoint, firedrake.Function): - result = block_variable.checkpoint.assign(result) - return maybe_disk_checkpoint(result) - - def solve_init_params(self, args, kwargs, varform): if len(self.forward_args) <= 0: self.forward_args = args @@ -862,233 +342,8 @@ def solve_init_params(self, args, kwargs, varform): self.assemble_kwargs["appctx"] = kwargs["appctx"] -class SolveVarFormBlock(GenericSolveBlock): - def __init__(self, equation, func, bcs=[], *args, **kwargs): - lhs = equation.lhs - rhs = equation.rhs - super().__init__(lhs, rhs, func, bcs, *args, **kwargs) - - def _init_solver_parameters(self, args, kwargs): - super()._init_solver_parameters(args, kwargs) - solve_init_params(self, args, kwargs, varform=True) - - -class NonlinearVariationalSolveBlock(GenericSolveBlock): - def __init__(self, equation, func, bcs, adj_cache, problem_J, - solver_kwargs, **kwargs): - lhs = equation.lhs - rhs = equation.rhs - - self._adj_cache = adj_cache - self._dFdm_cache = adj_cache.setdefault("dFdm_cache", {}) - self.problem_J = problem_J - self.solver_kwargs = solver_kwargs - - super().__init__(lhs, rhs, func, bcs, **{**solver_kwargs, **kwargs}) - - if self.problem_J is not None: - for coeff in self.problem_J.coefficients(): - self.add_dependency(coeff, no_duplicates=True) - - def _init_solver_parameters(self, args, kwargs): - super()._init_solver_parameters(args, kwargs) - solve_init_params(self, args, kwargs, varform=True) - - def recompute_component(self, inputs, block_variable, idx, prepared): - tape = get_working_tape() - if self._ad_solvers["recompute_count"] == tape.recompute_count - 1: - # Update how many times the block has been recomputed. - self._ad_solvers["recompute_count"] = tape.recompute_count - if self._ad_solvers["forward_nlvs"]._problem._constant_jacobian: - self._ad_solvers["forward_nlvs"].invalidate_jacobian() - self._ad_solvers["update_adjoint"] = True - return super().recompute_component(inputs, block_variable, idx, prepared) - - def _forward_solve(self, lhs, rhs, func, bcs, **kwargs): - self._ad_solver_replace_forms() - self._ad_solvers["forward_nlvs"].solve() - func.assign(self._ad_solvers["forward_nlvs"]._problem.u) - return func - - def _adjoint_solve(self, dJdu, compute_bdy): - dJdu_copy = dJdu.copy() - # Homogenize and apply boundary conditions on adj_dFdu and dJdu. - for bc in self.bcs: - bc.zero(dJdu) - - if ( - self._ad_solvers["forward_nlvs"]._problem._constant_jacobian - and self._ad_solvers["update_adjoint"] - ): - # Update left hand side of the adjoint equation. - 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(SolverType.ADJOINT) - - # Update the right hand side of the adjoint equation. - # problem.F._component[1] is the right hand side of the adjoint. - self._ad_solvers["adjoint_lvs"]._problem.F._components[1].assign(dJdu) - - # Solve the adjoint linear variational solver. - self._ad_solvers["adjoint_lvs"].solve() - u_sol = self._ad_solvers["adjoint_lvs"]._problem.u - - adj_sol_bdy = None - if compute_bdy: - jac_adj = self._ad_solvers["adjoint_lvs"]._problem.J - adj_sol_bdy = self._compute_adj_bdy( - u_sol, adj_sol_bdy, jac_adj, dJdu_copy) - return u_sol, adj_sol_bdy - - def _ad_assign_map(self, form, solver): - 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 - assign_map = {} - form_ad_count_map = dict((count_map[coeff], coeff) - for coeff in form.coefficients()) - for block_variable in self.get_dependencies(): - coeff = block_variable.output - if isinstance(coeff, - (firedrake.Coefficient, firedrake.Constant, - firedrake.Cofunction)): - coeff_count = coeff.count() - if coeff_count in form_ad_count_map: - assign_map[form_ad_count_map[coeff_count]] = \ - block_variable.saved_output - - if ( - solver == SolverType.ADJOINT - and not self._ad_solvers["forward_nlvs"]._problem._constant_jacobian - ): - block_variable = self.get_outputs()[0] - coeff_count = block_variable.output.count() - if coeff_count in form_ad_count_map: - assign_map[form_ad_count_map[coeff_count]] = \ - block_variable.saved_output - return assign_map - - def _ad_assign_coefficients(self, form, solver): - assign_map = self._ad_assign_map(form, solver) - for coeff, value in assign_map.items(): - coeff.assign(value) - - def _ad_assign_bcs(self, problem): - for block_variable in self.get_dependencies(): - bc = block_variable.output - if bc in problem._ad_bc_map: - problem._ad_bc_map[bc].function_arg.assign( - block_variable.saved_output.function_arg) - - 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) - self._ad_assign_bcs(problem) - else: - self._ad_assign_coefficients( - self._ad_solvers["adjoint_lvs"]._problem.J, solver) - - def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies): - compute_bdy = self._should_compute_boundary_adjoint( - relevant_dependencies - ) - adj_sol, adj_sol_bdy = self._adjoint_solve(adj_inputs[0], compute_bdy) - self.adj_sol = 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: - self.adj_bdy_cb(adj_sol_bdy) - - r = {} - r["form"] = self._create_F_form() - r["adj_sol"] = self.adj_state - r["adj_sol_bdy"] = adj_sol_bdy - return r - - def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, - prepared=None): - if not self.linear and self.func == block_variable.output: - # We are not able to calculate derivatives wrt initial guess. - return None - F_form = prepared["form"] - adj_sol = prepared["adj_sol"] - adj_sol_bdy = prepared["adj_sol_bdy"] - c = block_variable.output - c_rep = block_variable.saved_output - - if isinstance(c, (firedrake.Function, firedrake.Cofunction)): - trial_function = firedrake.TrialFunction(c.function_space()) - elif isinstance(c, firedrake.Constant): - try: - mesh = extract_unique_domain(F_form) - except ValueError: - raise ValueError("Expecting a single mesh") - trial_function = firedrake.TrialFunction( - c._ad_function_space(mesh) - ) - elif isinstance(c, firedrake.DirichletBC): - tmp_bc = c.reconstruct( - g=extract_subfunction(adj_sol_bdy, c.function_space()) - ) - return [tmp_bc] - elif isinstance(c, firedrake.MeshGeometry): - # Using CoordianteDerivative requires us to do action before - # differentiating, might change in the future. - F_form_tmp = firedrake.action(F_form, adj_sol) - X = firedrake.SpatialCoordinate(c_rep) - dFdm = firedrake.derivative(-F_form_tmp, X, firedrake.TestFunction( - c._ad_function_space()) - ) - - dFdm = firedrake.assemble(dFdm, **self.assemble_kwargs) - return dFdm - - # dFdm_cache works with original variables, not block saved outputs. - if c in self._dFdm_cache: - dFdm = self._dFdm_cache[c] - else: - dFdm = -firedrake.derivative(self.lhs, c, trial_function) - dFdm = firedrake.adjoint(dFdm) - self._dFdm_cache[c] = dFdm - - # Replace the form coefficients with checkpointed values. - replace_map = self._replace_map(dFdm) - replace_map[self.func] = self.get_outputs()[0].saved_output - dFdm = replace(dFdm, replace_map) - - if isinstance(dFdm, firedrake.Argument): - # Corner case. Should be fixed more permanently upstream in UFL. - # See: https://github.com/FEniCS/ufl/issues/395 - dFdm = ufl.Action(dFdm, adj_sol) - else: - dFdm = dFdm * adj_sol - dFdm = firedrake.assemble(dFdm, **self.assemble_kwargs) - - return dFdm - - -class ProjectBlock(SolveVarFormBlock): - def __init__(self, v, V, output, bcs=[], *args, **kwargs): - mesh = kwargs.pop("mesh", None) - if mesh is None: - mesh = V.mesh() - dx = firedrake.dx(mesh) - w = firedrake.TestFunction(V) - Pv = firedrake.TrialFunction(V) - a = firedrake.inner(Pv, w) * dx - L = firedrake.inner(v, w) * dx - - super().__init__(a == L, output, bcs, *args, **kwargs) - - def _init_solver_parameters(self, args, kwargs): - super()._init_solver_parameters(args, kwargs) - solve_init_params(self, args, kwargs, varform=True) +class ProjectBlock(NonlinearVariationalSolveBlock): + pass class SupermeshProjectBlock(Block): @@ -1114,6 +369,11 @@ class SupermeshProjectBlock(Block): Step 2. solve linear system. """ + + pop_kwargs_keys = ["adj_cb", "adj_bdy_cb", "adj2_cb", "adj2_bdy_cb", + "forward_args", "forward_kwargs", "adj_args", + "adj_kwargs"] + def __init__(self, source, target_space, target, bcs=[], **kwargs): super(SupermeshProjectBlock, self).__init__( ad_block_tag=kwargs.pop("ad_block_tag", None) diff --git a/firedrake/adjoint_utils/function.py b/firedrake/adjoint_utils/function.py index 032a289b3d..cdd20cc101 100644 --- a/firedrake/adjoint_utils/function.py +++ b/firedrake/adjoint_utils/function.py @@ -1,9 +1,8 @@ from functools import wraps import ufl -from ufl.domain import extract_unique_domain from pyadjoint.overloaded_type import create_overloaded_object, FloatingType from pyadjoint.tape import annotate_tape, stop_annotating, get_working_tape -from firedrake.adjoint_utils.blocks import FunctionAssignBlock, ProjectBlock, SubfunctionBlock, FunctionMergeBlock, SupermeshProjectBlock +from firedrake.adjoint_utils.blocks import FunctionAssignBlock, SubfunctionBlock, FunctionMergeBlock import firedrake from .checkpointing import disk_checkpointing, CheckpointFunction, \ CheckpointBase, checkpoint_init_data, DelegatedFunctionCheckpoint @@ -26,32 +25,6 @@ def wrapper(self, *args, **kwargs): init(self, *args, **kwargs) return wrapper - @staticmethod - def _ad_annotate_project(project): - @wraps(project) - def wrapper(self, b, *args, **kwargs): - ad_block_tag = kwargs.pop("ad_block_tag", None) - annotate = annotate_tape(kwargs) - - if annotate: - bcs = kwargs.get("bcs", []) - if isinstance(b, firedrake.Function) and extract_unique_domain(b) != self.function_space().mesh(): - block = SupermeshProjectBlock(b, self.function_space(), self, bcs, ad_block_tag=ad_block_tag) - else: - block = ProjectBlock(b, self.function_space(), self, bcs, ad_block_tag=ad_block_tag) - - tape = get_working_tape() - tape.add_block(block) - - with stop_annotating(): - output = project(self, b, *args, **kwargs) - - if annotate: - block.add_output(output.create_block_variable()) - - return output - return wrapper - @staticmethod def _ad_annotate_subfunctions(subfunctions): @wraps(subfunctions) diff --git a/firedrake/adjoint_utils/projection.py b/firedrake/adjoint_utils/projection.py index e847c74258..08944b53c2 100644 --- a/firedrake/adjoint_utils/projection.py +++ b/firedrake/adjoint_utils/projection.py @@ -1,45 +1,71 @@ +import firedrake from functools import wraps from pyadjoint.tape import annotate_tape, stop_annotating, get_working_tape from firedrake.adjoint_utils.blocks import ProjectBlock, SupermeshProjectBlock -from firedrake import function -from ufl.domain import extract_unique_domain +from firedrake.adjoint_utils.variational_solver import NonlinearVariationalSolverMixin def annotate_project(project): + """A wrapper for the bare project() call to handle switching annotation + on/off without dispatching into all the types of projectors. + """ + @wraps(project) def wrapper(*args, **kwargs): - """The project call performs an equation solve, and so it too must be annotated so that the - adjoint and tangent linear models may be constructed automatically by pyadjoint. + annotate = annotate_tape(kwargs) + if annotate: + # annotation enabled: let the Projector or Assigner record + # its own block through the ProjectorMixin path + return project(*args, **kwargs) + + with stop_annotating(): + # if global annotation was enabled, but annotation=False + # was passed to this call: take the same path, but ensuring + # all nested annotation is disabled + return project(*args, **kwargs) - To disable the annotation of this function, just pass :py:data:`annotate=False`. 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).""" + return wrapper + + +def annotate_super_project(project): + """A wrapper for SupermeshProjector.project(), only handles + the code path that leads to the creation of a SupermeshProjectBlock. + """ - ad_block_tag = kwargs.pop("ad_block_tag", None) + @wraps(project) + def wrapper(self, **kwargs): annotate = annotate_tape(kwargs) + V = self.target.function_space() if annotate: bcs = kwargs.get("bcs", []) - sb_kwargs = ProjectBlock.pop_kwargs(kwargs) - if isinstance(args[1], function.Function): - # block should be created before project because output might also be an input that needs checkpointing - output = args[1] - V = output.function_space() - if isinstance(args[0], function.Function) and extract_unique_domain(args[0]) != V.mesh(): - block = SupermeshProjectBlock(args[0], V, output, bcs, ad_block_tag=ad_block_tag, **sb_kwargs) - else: - block = ProjectBlock(args[0], V, output, bcs, ad_block_tag=ad_block_tag, **sb_kwargs) + sb_kwargs = SupermeshProjectBlock.pop_kwargs(kwargs) + if self._target_is_function: + # block should be created before project because output + # might also be an input that needs checkpointing + block = SupermeshProjectBlock( + self.source, + V, + self.target, + bcs, + ad_block_tag=self.ad_block_tag, + **sb_kwargs, + ) with stop_annotating(): - output = project(*args, **kwargs) + output = project(self, **kwargs) if annotate: tape = get_working_tape() - if not isinstance(args[1], function.Function): - if isinstance(args[0], function.Function) and extract_unique_domain(args[0]) != args[1].mesh(): - block = SupermeshProjectBlock(args[0], args[1], output, bcs, ad_block_tag=ad_block_tag, **sb_kwargs) - else: - block = ProjectBlock(args[0], args[1], output, bcs, ad_block_tag=ad_block_tag, **sb_kwargs) + if not self._target_is_function: + block = SupermeshProjectBlock( + self.source, + V, + output, + bcs, + ad_block_tag=self.ad_block_tag, + **sb_kwargs, + ) + tape.add_block(block) block.add_output(output.create_block_variable()) @@ -48,37 +74,51 @@ def wrapper(*args, **kwargs): return wrapper -class ProjectorMixin: +class ProjectorMixin(NonlinearVariationalSolverMixin): + """A wrapper over the nonlinear variational solver mixin for projectors. + + This establishes the form used for adjoint evaluation of the projector.""" + + def __init__(self, ad_block_tag=None): + V = self.target.function_space() + mesh = V.mesh() + ad_target = firedrake.Function(V) + + dx = firedrake.dx(mesh) + w = firedrake.TestFunction(V) + Pv = firedrake.TrialFunction(V) + a = firedrake.inner(Pv, w) * dx + L = firedrake.inner(self.source, w) * dx + p = firedrake.LinearVariationalProblem( + a, L, ad_target, bcs=self.bcs, + ) + + # If we project from an expression containing u (e.g. u + c) onto u, the + # resulting form looks nonlinear. However, the projection is in effect a + # solve onto a temporary function (here ad_target), followed by an implicit + # assignment. Instead of representing this as an explicit solve and assign + # block, we solve for a temporary function and modify _ad_u on the problem. + # This is used for creating the output variable (and thus the dependency + # resolution), in effect merging the implicit assignment into the block + # output phase. + p._ad_u = self.target + + solver_params = firedrake.LinearVariationalSolver.DEFAULT_SNES_PARAMETERS | self.solver_parameters + + self._init_as_solver( + p, + forward_kwargs={"solver_parameters": solver_params}, + ad_block_tag=ad_block_tag, + ) + + @NonlinearVariationalSolverMixin._ad_annotate_init + def _init_as_solver(self, problem, **kwargs): + return + + @staticmethod + def _ad_block_type(): + return ProjectBlock + @staticmethod def _ad_annotate_project(project): - @wraps(project) - def wrapper(self, *args, **kwargs): - """Annotate the projection performed by a (possibly reused) - Projector, so that repeatedly calling ``project`` tapes one - ProjectBlock per call against the current value of the source. - - To disable the annotation, just pass :py:data:`annotate=False`.""" - ad_block_tag = kwargs.pop("ad_block_tag", None) - annotate = annotate_tape(kwargs) - if annotate: - bcs = self.bcs or [] - output = self.target - V = output.function_space() - # The block is created before projecting because the output - # might also be an input that needs checkpointing. - if isinstance(self.source, function.Function) and extract_unique_domain(self.source) != V.mesh(): - block = SupermeshProjectBlock(self.source, V, output, bcs, ad_block_tag=ad_block_tag) - else: - block = ProjectBlock(self.source, V, output, bcs, ad_block_tag=ad_block_tag) - tape = get_working_tape() - tape.add_block(block) - - with stop_annotating(): - output = project(self, *args, **kwargs) - - if annotate: - block.add_output(output.create_block_variable()) - - return output - - return wrapper + return NonlinearVariationalSolverMixin._ad_annotate_solve(project) diff --git a/firedrake/adjoint_utils/solving.py b/firedrake/adjoint_utils/solving.py index c7301288f4..5c055aa4f9 100644 --- a/firedrake/adjoint_utils/solving.py +++ b/firedrake/adjoint_utils/solving.py @@ -1,5 +1,5 @@ from pyadjoint.tape import get_working_tape -from firedrake.adjoint_utils.blocks import CachedSolverBlock, GenericSolveBlock, ProjectBlock +from firedrake.adjoint_utils.blocks import NonlinearVariationalSolveBlock, ProjectBlock def get_solve_blocks(): @@ -11,6 +11,6 @@ def get_solve_blocks(): return [ block for block in get_working_tape().get_blocks() - if issubclass(type(block), (CachedSolverBlock, GenericSolveBlock)) + if issubclass(type(block), NonlinearVariationalSolveBlock) and not issubclass(type(block), ProjectBlock) ] diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py index 4de5d04777..184da642de 100644 --- a/firedrake/adjoint_utils/variational_solver.py +++ b/firedrake/adjoint_utils/variational_solver.py @@ -1,8 +1,7 @@ import copy 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, CachedSolverBlock) +from firedrake.adjoint_utils.blocks import NonlinearVariationalSolveBlock from firedrake.adjoint_utils.blocks.solving import solve_init_params from firedrake.ufl_expr import derivative, adjoint, action from ufl import replace, Action @@ -483,6 +482,10 @@ def _ad_hessian_cache(self): d2Fdudm_forms=d2Fdudm_forms, ) + @staticmethod + def _ad_block_type(): + return NonlinearVariationalSolveBlock + @staticmethod def _ad_annotate_solve(solve): @wraps(solve) @@ -497,7 +500,7 @@ def wrapper(self, **kwargs): raise ValueError( "MissingMathsError: we do not know how to differentiate through a variational inequality") - block = CachedSolverBlock( + block = self._ad_block_type()( self._ad_forward_cache, self._ad_tangent_cache, self._ad_adjoint_cache, @@ -521,66 +524,6 @@ def wrapper(self, **kwargs): 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 - 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).""" - from firedrake import LinearVariationalSolver - annotate = annotate_tape(kwargs) - if annotate: - bounds = kwargs.pop("bounds", None) - if bounds is not None: - raise ValueError( - "MissingMathsError: we do not know how to differentiate through a variational inequality") - - tape = get_working_tape() - problem = self._ad_problem - sb_kwargs = NonlinearVariationalSolveBlock.pop_kwargs(kwargs) - sb_kwargs.update(kwargs) - - block = NonlinearVariationalSolveBlock(problem._ad_F == 0, - problem._ad_u, - problem._ad_bcs, - adj_cache=self._ad_adj_cache, - problem_J=problem._ad_J, - solver_kwargs=self._ad_kwargs, - ad_block_tag=self.ad_block_tag, - **sb_kwargs) - - # Forward variational solver. - if not self._ad_solvers["forward_nlvs"]: - self._ad_solvers["forward_nlvs"] = type(self)( - self._ad_problem_clone(self._ad_problem, block.get_dependencies()), - **self._ad_kwargs - ) - - # Adjoint variational solver. - if not self._ad_solvers["adjoint_lvs"]: - with stop_annotating(): - self._ad_solvers["adjoint_lvs"] = LinearVariationalSolver( - self._ad_adj_lvs_problem(block, problem._ad_adj_F), - *block.adj_args, **block.adj_kwargs) - if self._ad_problem._constant_jacobian: - self._ad_solvers["update_adjoint"] = False - - block._ad_solvers = self._ad_solvers - - 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 - @no_annotations def _ad_problem_clone(self, problem, dependencies): """Replaces every coefficient in the residual, jacobian and Dirichlet boundary values diff --git a/firedrake/function.py b/firedrake/function.py index eef7ccce97..9783185b1a 100644 --- a/firedrake/function.py +++ b/firedrake/function.py @@ -356,7 +356,6 @@ def sub(self, indices: tuple[int] | int) -> "Function": return self._components[indices] @PETSc.Log.EventDecorator() - @FunctionMixin._ad_annotate_project def project(self, b, *args, **kwargs): r"""Project ``b`` onto ``self``. ``b`` must be a :class:`Function` or a UFL expression. diff --git a/firedrake/projection.py b/firedrake/projection.py index 57c7fa4d8a..a4cb87570b 100644 --- a/firedrake/projection.py +++ b/firedrake/projection.py @@ -12,7 +12,7 @@ from firedrake import functionspaceimpl from firedrake import function -from firedrake.adjoint_utils import annotate_project, ProjectorMixin +from firedrake.adjoint_utils import annotate_project, annotate_super_project, ProjectorMixin __all__ = ['project', 'Projector'] @@ -110,19 +110,21 @@ def project( solver_parameters=solver_parameters, form_compiler_parameters=form_compiler_parameters, use_slate_for_inverse=use_slate_for_inverse, - quadrature_degree=quadrature_degree + quadrature_degree=quadrature_degree, + ad_block_tag=ad_block_tag, ).project() val.rename(name) return val class Assigner: - def __init__(self, source, target): + def __init__(self, source, target, **kwargs): self.source = source self.target = target + self._kwargs = kwargs def project(self): - self.target.assign(self.source) + self.target.assign(self.source, **self._kwargs) return self.target @@ -228,13 +230,12 @@ def residual(self): def rhs(self): pass - @ProjectorMixin._ad_annotate_project def project(self): self.apply_massinv(self.target, self.rhs) return self.target -class BasicProjector(ProjectorBase): +class BasicProjector(ProjectorBase, ProjectorMixin): """ A basic projector projects a UFL expression into a function space and places the result in a function from that function space, @@ -243,6 +244,10 @@ class BasicProjector(ProjectorBase): defined on the same mesh. """ + def __init__(self, *args, ad_block_tag=None, **kwargs): + ProjectorBase.__init__(self, *args, **kwargs) + ProjectorMixin.__init__(self, ad_block_tag=ad_block_tag) + @cached_property def rhs_form(self): v = firedrake.TestFunction(self.target.function_space()) @@ -270,6 +275,10 @@ def rhs_form(self): form = firedrake.inner(self.source, v)*firedrake.dx return form + @ProjectorMixin._ad_annotate_project + def project(self): + return super().project() + @cached_property def assembler(self): from firedrake.assemble import get_assembler @@ -283,6 +292,11 @@ def rhs(self): class SupermeshProjector(ProjectorBase): + def __init__(self, *args, **kwargs): + self._target_is_function = kwargs.pop("target_is_function", True) + self.ad_block_tag = kwargs.pop("ad_block_tag", None) + super().__init__(*args, **kwargs) + @cached_property def mixed_mass(self): from firedrake.supermeshing import assemble_mixed_mass_matrix @@ -295,6 +309,10 @@ def rhs(self): self.mixed_mass.mult(u, v) return self.residual + @annotate_super_project + def project(self): + return super().project() + @PETSc.Log.EventDecorator() def Projector( @@ -305,7 +323,8 @@ def Projector( form_compiler_parameters: Optional[dict] = None, constant_jacobian: Optional[bool] = True, use_slate_for_inverse: Optional[bool] = False, - quadrature_degree: Optional[Union[int, tuple[int]]] = None + quadrature_degree: Optional[Union[int, tuple[int]]] = None, + ad_block_tag: Optional[str] = None, ): """ Projection class. @@ -337,6 +356,8 @@ def Projector( function spaces)(only valid for DG function spaces). quadrature_degree Quadrature degree to use when approximating integrands. + ad_block_tag + String for tagging the resulting block on the Pyadjoint tape. """ target = create_output(v_out) source = sanitise_input(v, target.function_space()) @@ -345,14 +366,15 @@ def Projector( raise ValueError("Shape mismatch between source %s and target %s in project" % (source.ufl_shape, target.ufl_shape)) if isinstance(v, function.Function) and not bcs and v.function_space() == target.function_space(): - return Assigner(source, target) + return Assigner(source, target, ad_block_tag=ad_block_tag) elif source_mesh == target_mesh: return BasicProjector( source, target, bcs=bcs, solver_parameters=solver_parameters, form_compiler_parameters=form_compiler_parameters, constant_jacobian=constant_jacobian, use_slate_for_inverse=use_slate_for_inverse, - quadrature_degree=quadrature_degree + quadrature_degree=quadrature_degree, + ad_block_tag=ad_block_tag, ) else: if bcs is not None: @@ -364,5 +386,7 @@ def Projector( form_compiler_parameters=form_compiler_parameters, constant_jacobian=constant_jacobian, use_slate_for_inverse=use_slate_for_inverse, - quadrature_degree=quadrature_degree + quadrature_degree=quadrature_degree, + target_is_function=isinstance(v_out, firedrake.Function), + ad_block_tag=ad_block_tag, ) diff --git a/tests/firedrake/adjoint/test_projection.py b/tests/firedrake/adjoint/test_projection.py index c059dd3dac..6951184b3a 100644 --- a/tests/firedrake/adjoint/test_projection.py +++ b/tests/firedrake/adjoint/test_projection.py @@ -179,3 +179,40 @@ def test_project_to_function_space(): J = assemble(w**2*dx) rf = ReducedFunctional(J, Control(c)) assert taylor_test(rf, Function(R, val=1.), Constant(0.1)) > 1.9 + + +@pytest.mark.skipcomplex +def test_project_blocks(): + mesh = UnitSquareMesh(12, 1) + V = FunctionSpace(mesh, "CG", 1) + R = FunctionSpace(mesh, "R", 0) + u = Function(V) + c = Function(R, val=1.) + project(u+c, u) + assert len(get_solve_blocks()) == 0 + + +@pytest.mark.skipcomplex +def test_project_bcs(): + mesh = UnitIntervalMesh(8) + V = FunctionSpace(mesh, "CG", 1) + m = Function(V).assign(1.0) + bc = DirichletBC(V, 3.0, "on_boundary") + u = Function(V) + + project(m, u, bcs=bc) + J = assemble(u**2 * dx) + rf = ReducedFunctional(J, Control(m)) + + # create new reference value independent of tape evaluation + with stop_annotating(): + m_new = Function(V).assign(2.0) + u_exp = Function(V) + project(m_new, u_exp, bcs=bc) + J_exp = assemble(u_exp**2 * dx) + + # ensure BCs are passed through the tape + assert rf(m_new) == pytest.approx( + J_exp, rel=1e-12, abs=1e-12, + ) + assert taylor_test(rf, m_new, Constant(0.1)) > 1.9