diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index 549c8be..d509fa5 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -11,10 +11,12 @@ from dolfinx.fem.function import Function as _Function from ..compat import compute_form_adjoint -from ..petsc_utils import LinearAdjointProblem, solve_linear_problem from ..types import Function from .assembly import _create_vector, _SpecialVector, assemble_compiled_form +if typing.TYPE_CHECKING: + from ..solvers import LinearProblem, NonlinearProblem + type NestedMutableSequence[T] = T | typing.MutableSequence["NestedMutableSequence[T]"] type NestedSequence[T] = T | typing.Sequence["NestedSequence[T]"] @@ -50,32 +52,69 @@ def assign_mixed_parts( The modified form structures with identical nesting and sequence types, where all unassigned TestFunction and TrialFunction arguments have been mapped. Returns a single structure if one was passed, otherwise returns a tuple. + + Note: + The replacement arguments are drawn from ``ufl.TestFunctions``/``ufl.TrialFunctions`` + of a single ``ufl.MixedFunctionSpace`` built from the row/column function spaces + discovered while walking the structure, rather than hand-constructed + via ``ufl.Argument(..., part=...)``: this is exactly what a user who + builds the block system directly on a ``ufl.MixedFunctionSpace`` (and + then calls ``ufl.extract_blocks``) already gets, so a form assembled + this way and one assembled from a plain ``[[a00, ...], ...]`` nested + list end up as the same UFL objects -- one code path handles both, + rather than two subtly different ones. """ - replace_map = {} + spaces: dict[int, ufl.functionspace.AbstractFunctionSpace] = {} + + def _discover_spaces(obj: NestedSequence[ufl.Form], indices: tuple[int, ...]) -> None: + """Recursively discover, for each row/column index, the function space of the + (as yet unassigned) argument occupying that position. - def _build_map(obj: NestedSequence[ufl.Form], indices: tuple[int, ...]) -> None: - """ - Recursively discover forms, tracking the depth and index of the nesting. `indices` will be `(row,)` for vectors and `(row, col)` for matrices. """ if isinstance(obj, ufl.Form): for arg in obj.arguments(): - # Only map arguments that haven't been assigned a part yet - if arg.part() is None and arg not in replace_map: + if arg.part() is None: num = arg.number() - # Because num is 0 for TestFunctions and 1 for TrialFunctions, # it maps perfectly to our nested dimension indices! # If num < len(indices), we have traversed deep enough to assign it. if num < len(indices): - replace_map[arg] = ufl.Argument(arg.ufl_function_space(), number=num, part=indices[num]) + spaces.setdefault(indices[num], arg.ufl_function_space()) + elif isinstance(obj, typing.Iterable): + for i, item in enumerate(obj): + if item is not None: + _discover_spaces(item, indices + (i,)) + + for struct in form_structs: + _discover_spaces(struct, ()) + + # If no replacements are needed, exit early to save computation + if not spaces: + return form_structs if len(form_structs) > 1 else form_structs[0] + + num_parts = max(spaces) + 1 + mixed_space = ufl.MixedFunctionSpace(*(spaces[i] for i in range(num_parts))) + test_functions = ufl.TestFunctions(mixed_space) + trial_functions = ufl.TrialFunctions(mixed_space) + + replace_map = {} + def _build_map(obj: NestedSequence[ufl.Form], indices: tuple[int, ...]) -> None: + if isinstance(obj, ufl.Form): + for arg in obj.arguments(): + if arg.part() is None and arg not in replace_map: + num = arg.number() + if num < len(indices): + replace_map[arg] = (test_functions if num == 0 else trial_functions)[indices[num]] elif isinstance(obj, typing.Iterable): for i, item in enumerate(obj): if item is not None: - # Append current topological index to the path and recurse _build_map(item, indices + (i,)) + for struct in form_structs: + _build_map(struct, ()) + def _replace(obj: typing.Any) -> typing.Any: """ Recursively rebuild the structure using the populated replace_map, @@ -87,16 +126,7 @@ def _replace(obj: typing.Any) -> typing.Any: return type(obj)(_replace(item) for item in obj) return obj - # 1. Build a shared map across all inputs (e.g., ensuring RHS TestFunctions - # perfectly match LHS TestFunctions) - for struct in form_structs: - _build_map(struct, ()) - - # 2. If no replacements are needed, exit early to save computation - if not replace_map: - return form_structs if len(form_structs) > 1 else form_structs[0] - - # 3. Apply the replacements and unpack if necessary + # Apply the replacements and unpack if necessary replaced = tuple(_replace(struct) for struct in form_structs) return replaced if len(replaced) > 1 else replaced[0] @@ -160,15 +190,11 @@ def __init__( bcs: typing.Sequence[dolfinx.fem.DirichletBC] | None = None, u: _Function | None = None, P: ufl.Form | None = None, - kind: str | None = None, - petsc_options: dict | None = None, form_compiler_options: dict | None = None, jit_options: dict | None = None, entity_maps: typing.Sequence[dolfinx.mesh.EntityMap] | None = None, ad_block_tag: str | None = None, - adjoint_petsc_options: dict | None = None, - tlm_petsc_options: dict | None = None, - petsc_options_prefix: str = "dxa_linear_problem_block_", + problem: "LinearProblem" = ..., ) -> None: ... @typing.overload @@ -180,15 +206,11 @@ def __init__( bcs: typing.Sequence[dolfinx.fem.DirichletBC] | None = None, u: typing.Sequence[_Function] | None = None, P: typing.Sequence[typing.Sequence[ufl.Form]] | None = None, - kind: str | typing.Sequence[typing.Sequence[str]] | None = None, - petsc_options: dict | None = None, form_compiler_options: dict | None = None, jit_options: dict | None = None, entity_maps: typing.Sequence[dolfinx.mesh.EntityMap] | None = None, ad_block_tag: str | None = None, - adjoint_petsc_options: dict | None = None, - tlm_petsc_options: dict | None = None, - petsc_options_prefix: str = "dxa_linear_problem_block_", + problem: "LinearProblem" = ..., ) -> None: ... def __init__( @@ -199,19 +221,29 @@ def __init__( bcs: typing.Sequence[dolfinx.fem.DirichletBC] | None = None, u: _Function | typing.Sequence[_Function] | None = None, P: ufl.Form | typing.Sequence[typing.Sequence[ufl.Form]] | None = None, - kind: str | typing.Sequence[typing.Sequence[str]] | None = None, - petsc_options: dict | None = None, form_compiler_options: dict | None = None, jit_options: dict | None = None, entity_maps: typing.Sequence[dolfinx.mesh.EntityMap] | None = None, ad_block_tag: str | None = None, - adjoint_petsc_options: dict | None = None, - tlm_petsc_options: dict | None = None, - petsc_options_prefix: str = "dxa_linear_problem_block_", + problem: "LinearProblem" = None, # type: ignore[assignment] ) -> None: - self._adjoint_petsc_options = adjoint_petsc_options - self._tlm_petsc_options = tlm_petsc_options + assert problem is not None, "problem must be provided." + # A plain (strong) reference is deliberate: constructing a LinearProblem + # as a throwaway local (solve it, then only ever touch the resulting + # ReducedFunctional) is a common, already-tested pattern, and this + # block -- kept alive by the tape -- must keep the Problem (and hence + # its shared solvers) alive for exactly as long as the block itself is + # reachable, mirroring the lifetime the old per-block-owned solver had. + # A plain LinearProblemBlock is not itself cyclic garbage (verified: + # dropping the Problem and the tape releases it via ordinary + # refcounting, no gc.collect() required), so this does not reintroduce + # the MPI collective-destruction hazard documented in + # dolfinx-adjoint-knowledge's mpi-collective-destruction-hazard note -- + # that hazard is specifically about pyadjoint's checkpoint-schedule + # bookkeeping (an unmerged, not-yet-present feature) making the *tape* + # cyclic, and would need revisiting if/when that lands. + self._problem_obj = problem super().__init__(ad_block_tag=ad_block_tag) # Collect all arguments in variational forms and replace them with similar @@ -272,25 +304,11 @@ def __init__( self.add_dependency(c, no_duplicates=True) else: raise RuntimeError(f"Unknown type for unknown function u={type(self._u)}.") - self._compiled_lhs = dolfinx.fem.form( - self._lhs, - jit_options=jit_options, - form_compiler_options=form_compiler_options, - entity_maps=entity_maps, - ) - self._compiled_rhs = dolfinx.fem.form( - self._rhs, - jit_options=jit_options, - form_compiler_options=form_compiler_options, - entity_maps=entity_maps, - ) # Cache form parameters for later # NOTE: Should probably be in a struct self._jit_options = jit_options self._form_compiler_options = form_compiler_options self._entity_maps = entity_maps - self._petsc_options = petsc_options if petsc_options is not None else {} - self._petsc_options_prefix = petsc_options_prefix self._bcs = bcs if bcs is not None else [] # Add dependencies from the boundary conditions @@ -299,22 +317,10 @@ def __init__( if hasattr(bc, "block_variable"): self.add_dependency(bc, no_duplicates=True) - # Solver for recomputing the linear problem - self._forward_solver = dolfinx.fem.petsc.LinearProblem( - a=self._lhs, # type: ignore[arg-type] - L=self._rhs, # type: ignore[arg-type] - bcs=self._bcs, - u=self._u, # type: ignore[arg-type] - P=self._preconditioner, # type: ignore[arg-type] - petsc_options=self._petsc_options, - petsc_options_prefix=petsc_options_prefix, - form_compiler_options=self._form_compiler_options, - jit_options=self._jit_options, - kind=kind, # type: ignore[arg-type] - entity_maps=self._entity_maps, - ) # type: ignore[misc] - - self._kind = "nest" if self._forward_solver.A.getType() == "nest" else kind + # No forward/adjoint/TLM solver is built here: this block shares the + # ones owned by self._problem() (see LinearProblem in ../solvers.py), + # built once and reused across every block that Problem records + # instead of once per solve() call. if isinstance(self._u, dolfinx.fem.Function): self._adjoint_solutions = self._u.copy() @@ -326,19 +332,9 @@ def __init__( self._second_adjoint_solutions = [u.copy() for u in self._u] self._tlm_solutions = [u.copy() for u in self._u] - self._adjoint_solver = LinearAdjointProblem( - self._compute_adjoint(sum_form(self._lhs)), # type: ignore[arg-type] - self._rhs, # type: ignore[arg-type] - bcs=self._bcs, - P=self._preconditioner, # type: ignore[arg-type] - form_compiler_options=self._form_compiler_options, - jit_options=self._jit_options, - petsc_options=self._adjoint_petsc_options, - petsc_options_prefix=self._petsc_options_prefix, - kind=kind, # type: ignore[arg-type] - entity_maps=self._entity_maps, - ) # type: ignore[misc] - self._tlm_solver = None + def _problem(self) -> "LinearProblem": + """Return this block's owning Problem, which owns the shared solvers.""" + return self._problem_obj def _recover_bcs(self): bcs = [] @@ -350,23 +346,6 @@ def _recover_bcs(self): bcs.append(c_rep) return bcs - def construct_tlm_solver(self): - dFdu_form = self._compute_residual_derivative() - tlm_solver = LinearAdjointProblem( - dFdu_form, # type: ignore[arg-type] - self._rhs, # type: ignore[arg-type] - bcs=self._bcs, - u=self._tlm_solutions, # type: ignore[arg-type] - P=self._preconditioner, # type: ignore[arg-type] - form_compiler_options=self._form_compiler_options, - jit_options=self._jit_options, - petsc_options=self._tlm_petsc_options, - petsc_options_prefix=self._petsc_options_prefix, - kind=self._kind, # type: ignore[arg-type] - entity_maps=self._entity_maps, - ) # type: ignore[misc] - return tlm_solver - def _create_replace_map(self, form: ufl.Form | NestedMutableSequence[ufl.Form] | None) -> dict[Function, Function]: """Replace dependencies with latest checkpoint.""" replace_map = {} @@ -392,82 +371,31 @@ def _create_replace_map(self, form: ufl.Form | NestedMutableSequence[ufl.Form] | replace_map.update(self._create_replace_map(f)) return replace_map - def _create_recompute_replace_map(self, inputs: typing.Sequence[typing.Any]) -> dict: - """Map original dependency coefficients to active recomputation inputs.""" - replace_map = {} - for block_variable, input_val in zip(self.get_dependencies(), inputs, strict=True): - coeff = block_variable.output - val = input_val.output if isinstance(input_val, pyadjoint.block_variable.BlockVariable) else input_val - replace_map[coeff] = val - return replace_map - - @typing.overload - def _replace_form_coefficients_recompute( - self, form: ufl.Form | NestedSequence[ufl.Form], replace_map: dict - ) -> ufl.Form | NestedMutableSequence[ufl.Form | None]: ... - @typing.overload - def _replace_form_coefficients_recompute(self, form: None, replace_map: dict) -> None: ... - def _replace_form_coefficients_recompute( - self, form: ufl.Form | NestedSequence[ufl.Form] | None, replace_map: dict - ) -> ufl.Form | NestedMutableSequence[ufl.Form | None] | None: - """Recursively replace form coefficients for scalar or blocked form structures.""" - if form is None: - return None - if isinstance(form, ufl.Form): - coeffs_in_form = form.coefficients() - sub_map = {k: v for k, v in replace_map.items() if k in coeffs_in_form} - return ufl.replace(form, sub_map) if sub_map else form - elif isinstance(form, typing.Sequence): - return [self._replace_form_coefficients_recompute(f, replace_map) for f in form] - else: - raise TypeError(f"Cannot replace coefficients in form of type {type(form)}") - def prepare_recompute_component( self, inputs: typing.Sequence[typing.Any], relevant_outputs: typing.Sequence[typing.Any] ) -> dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function]: """Prepare for recomputing the block with different control inputs.""" - replace_map = self._create_recompute_replace_map(inputs) - - # 1. Substitute forms using active candidate inputs instead of static tape checkpoints - lhs = self._replace_form_coefficients_recompute(self._lhs, replace_map) - rhs = self._replace_form_coefficients_recompute(self._rhs, replace_map) - preconditioner = ( - self._replace_form_coefficients_recompute(self._preconditioner, replace_map) - if self._preconditioner is not None - else None - ) - - # 2. Recompile UFL forms with candidate inputs - compiled_lhs = dolfinx.fem.form( - lhs, # type: ignore[arg-type] - jit_options=self._jit_options, - form_compiler_options=self._form_compiler_options, - entity_maps=self._entity_maps, - ) - compiled_rhs = dolfinx.fem.form( - rhs, # type: ignore[arg-type] - jit_options=self._jit_options, - form_compiler_options=self._form_compiler_options, - entity_maps=self._entity_maps, - ) - compiled_preconditioner = ( - dolfinx.fem.form( - preconditioner, # type: ignore[arg-type] - jit_options=self._jit_options, - form_compiler_options=self._form_compiler_options, - entity_maps=self._entity_maps, - ) - if preconditioner is not None - else None - ) + # The forward solver (self._problem()) is bound, forever, to compiled + # forms referencing dedicated placeholder coefficients rather than the + # user's own dependency objects (see LinearProblem._value_placeholders, + # which mirrors NonlinearProblem's): writing this call's + # candidate/checkpointed values into the placeholders -- never into + # block_variable.output itself -- is what the next solve sees, + # without ever mutating an object the user (or a Taylor test + # perturbing a control directly) holds a live reference to. + problem = self._problem() + for block_variable in self.get_dependencies(): + placeholder = problem._value_placeholders.get(block_variable.output) + if placeholder is not None: + placeholder.x.array[:] = block_variable.saved_output.x.array[:] + placeholder.x.scatter_forward() - # 3. Hot-swap solver forms - self._forward_solver._a = compiled_lhs # type: ignore[assignment] - self._forward_solver._L = compiled_rhs # type: ignore[assignment] - self._forward_solver._preconditioner = compiled_preconditioner - self._forward_solver.bcs = self._bcs - self._forward_solver._u = self._u + # Re-establish this block's own bcs on the shared forward solver (the + # Problem itself -- see _problem()), since another block may have + # used it with different bcs in between. + problem.bcs = self._bcs + problem._u = self._u # Clear solution vector if isinstance(self._u, dolfinx.fem.Function): @@ -475,9 +403,11 @@ def prepare_recompute_component( else: for ui in self._u: ui.x.array[:] = 0.0 - # 4. Solve forward state while halting annotation + # Solve forward state while halting annotation. Call the base-class + # solve() directly (not problem.solve()), which would record another + # block onto the tape. with pyadjoint.stop_annotating(): - self._forward_solver.solve() + dolfinx.fem.petsc.LinearProblem.solve(problem) return self._u def recompute_component( @@ -558,82 +488,69 @@ def prepare_evaluate_tlm( self, inputs, tlm_inputs, relevant_outputs ) -> tuple[typing.Union[list[ufl.Form], ufl.Form], dolfinx.fem.Form]: - F_form, replacement_map = self._compute_residual() - if self._tlm_solver is None: - self._tlm_solver = self.construct_tlm_solver() - # Even if the solver is cached, we need to replace the form, as the output from pyadjoint - # is stored in a new function. - assert isinstance(self._tlm_solver, LinearAdjointProblem) - self._tlm_solver._a = dolfinx.fem.form( - self._compute_residual_derivative(), - jit_options=self._jit_options, - form_compiler_options=self._form_compiler_options, - entity_maps=self._entity_maps, - ) - # Build RHS (dFdm) for the monolithic system - if isinstance(self._u, list): - test_funcs = get_sorted_arguments(F_form.arguments(), 0) - dFdm = sum([ufl.ZeroBaseForm((test,)) for test in test_funcs]) - else: - test_funcs = [F_form.arguments()[0]] - dFdm = ufl.ZeroBaseForm((test_funcs[0],)) + # The TLM solver -- and the compiled LHS it solves with, shared + # verbatim with dF/du (see LinearProblem._get_or_build_dFdu_template) + # -- are shared across every block this Problem records; likewise the + # per-dependency TLM right-hand-side templates (see + # LinearProblem._get_or_build_tlm_rhs_templates) are each compiled + # once. Refresh this block's own checkpointed values into the + # placeholders and re-establish this block's own bcs on every call, + # since another block may have used the same solver in between -- but + # never rebuild or recompile any of these forms. + problem = self._problem() + tlm_solver = problem._get_or_build_tlm_solver() + tlm_solver.bcs = self._bcs + templates, seed_placeholders, state_placeholder = problem._get_or_build_tlm_rhs_templates() + for block_variable in self.get_dependencies(): + placeholder = problem._value_placeholders.get(block_variable.output) + if placeholder is not None: + placeholder.x.array[:] = block_variable.saved_output.x.array[:] + placeholder.x.scatter_forward() + state_list = state_placeholder if isinstance(state_placeholder, list) else [state_placeholder] + for placeholder, out_bv in zip(state_list, self.get_outputs(), strict=True): + placeholder.x.array[:] = out_bv.saved_output.x.array[:] + placeholder.x.scatter_forward() + + # 3. Assemble RHS Vector utilizing the shared solver's cached vector, + # accumulating only the dependencies that actually have a + # tangent-linear value this call -- see + # LinearProblem._get_or_build_tlm_rhs_templates for why an inactive + # dependency's term must be skipped entirely rather than evaluated + # with a zeroed direction. + b_petsc = tlm_solver._b + with b_petsc.localForm() as b_loc: + b_loc.set(0.0) for block_variable in self.get_dependencies(): tlm_value = block_variable.tlm_value - c_rep = block_variable.saved_output if tlm_value is None: continue - assert c_rep in replacement_map.values() - # Accumulate sensitivities across all block components - dFdm += ufl.derivative(-F_form, c_rep, tlm_value) - - # Safely wrap zero forms to prevent compilation crashes - dFdm = ufl.algorithms.expand_derivatives(dFdm) - if isinstance(self._u, list): - blocks = ufl.extract_blocks(dFdm) - if len(blocks) != len(self._u): - # Some zero blocks, manually pad with zero forms - _dFdm = [ufl.ZeroBaseForm((test,)) for test in test_funcs] - for block in blocks: - args = block.arguments() - assert len(args) == 1, "Expected a single test function in the block." - _dFdm[args[0].part()] = block - dFdm = _dFdm - else: - if dFdm == 0 or dFdm.empty(): - dFdm = ufl.ZeroBaseForm((test_funcs[0],)) - - dFdm_compiled = dolfinx.fem.form( - dFdm, - jit_options=self._jit_options, - form_compiler_options=self._form_compiler_options, - entity_maps=self._entity_maps, - ) - - # 3. Assemble RHS Vector utilizing the internal block-allocated vector - b_petsc = self._tlm_solver._b - with b_petsc.localForm() as b_loc: - b_loc.set(0.0) - - dolfinx.fem.petsc.assemble_vector(b_petsc, dFdm_compiled) + template = templates.get(block_variable.output) + if template is None: + continue + seed = seed_placeholders[block_variable.output] + seed.x.array[:] = tlm_value.x.array[:] + seed.x.scatter_forward() + dolfinx.fem.petsc.assemble_vector(b_petsc, template) dolfinx.la.petsc._ghost_update(b_petsc, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) - # 4. Apply Homogeneous Boundary Conditions safely directly to the block vector - if self._bcs: - try: - for bc in self._bcs: - bc.set(b_petsc.array_w, alpha=0.0) - except RuntimeError: - # FEniCSx throws RuntimeError for flat .set() on a blocked array. - # We use the bcs_by_block utility to handle the nested extraction. - from dolfinx.fem.bcs import bcs_by_block - - V_ext = dolfinx.fem.extract_function_spaces(dFdm_compiled) - bcs_lift = bcs_by_block(V_ext, self._bcs) - dolfinx.fem.petsc.set_bc(b_petsc, bcs_lift, alpha=0.0) - - # 5. Solve the full monolithic TLM system - self._tlm_solver.solve() + # Homogeneous boundary conditions are applied to b_petsc by + # tlm_solver.solve() itself (it homogenizes using tlm_solver.bcs, + # already re-established above), so there is no need to duplicate + # that here. + + # 4. Solve the full monolithic TLM system. The solver's own solution + # storage is shared across every block, so copy the result out into + # this block's own buffer immediately (mirroring how the adjoint path + # copies out of the shared adjoint solver in prepare_evaluate_adj) + # rather than relying on solver-owned storage identity. + tlm_solver.solve() + if isinstance(self._tlm_solutions, list): + for tlm_sol, sol in zip(self._tlm_solutions, tlm_solver.u): + tlm_sol.x.array[:] = sol.x.array[:] + else: + assert isinstance(self._tlm_solutions, dolfinx.fem.Function) + self._tlm_solutions.x.array[:] = tlm_solver.u.x.array[:] return self._tlm_solutions @@ -654,17 +571,38 @@ def prepare_evaluate_adj( ) -> tuple[ufl.Form, dict[Function, Function]]: """Prepare the block for evaluating the adjoint.""" + # The adjoint solver -- and the compiled LHS it solves with -- are + # shared across every block this Problem records: adjoint(dF/du) does + # not actually depend on the state u for a linear problem (F is + # linear in u, so its derivative doesn't reference u's value at all), + # so once every *other* dependency is routed through a placeholder + # (see LinearProblem._value_placeholders), that operator is fixed for + # the life of the Problem and was compiled once, in + # LinearProblem._get_or_build_adjoint_solver. Refresh this block's own + # checkpointed values into the placeholders and re-establish this + # block's own bcs on every call, since another block may have used + # the same solver in between -- but never rebuild or recompile the + # form itself. + problem = self._problem() + adjoint_solver = problem._get_or_build_adjoint_solver() + adjoint_solver.bcs = self._bcs + for block_variable in self.get_dependencies(): + placeholder = problem._value_placeholders.get(block_variable.output) + if placeholder is not None: + placeholder.x.array[:] = block_variable.saved_output.x.array[:] + placeholder.x.scatter_forward() + # Extract dJ/du[v] from the adjoint inputs. if len(adj_inputs) == 1: adj_rhs = adj_inputs[0] - dJdu = self._adjoint_solver._b + dJdu = adjoint_solver._b with dJdu.localForm() as dJdu_loc, adj_rhs.petsc_vec.localForm() as adj_rhs_loc: dJdu_loc.array[:] = adj_rhs_loc.array[:] else: assert len(adj_inputs) == len(self.get_outputs()), ( f"Expected {len(self.get_outputs())} adjoint inputs, got {len(adj_inputs)})" ) - dJdu = self._adjoint_solver._b + dJdu = adjoint_solver._b with dJdu.localForm() as dJdu_loc: dJdu_loc.set(0.0) @@ -678,28 +616,18 @@ def prepare_evaluate_adj( dolfinx.la.petsc.assign(arrs, dJdu) dJdu.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD) # type: ignore[arg-type] - # Compute (dF/du[v])* for the linear problem. + # F_form/replacement_map are still needed by evaluate_adj_component + # (to build each dependency's own sensitivity form), but the adjoint + # LHS itself is already correct on adjoint_solver -- no rebuild, no + # recompile. F_form, replacement_map = self._compute_residual() - dFdu = self._compute_residual_derivative() - summed = sum_form(dFdu) - dFdu_adj = compute_form_adjoint(summed) - dFdu_adj = ufl.algorithms.apply_derivatives.apply_derivatives(ufl.algorithms.expand_derivatives(dFdu_adj)) - assert dFdu_adj.empty() is False, "Adjoint of dF/du[v] is empty. Check if the problem is linear." - # Solve adjoint problem - compiled_dFdu = dolfinx.fem.form( - ufl.extract_blocks(dFdu_adj), # type: ignore[arg-type] - jit_options=self._jit_options, - form_compiler_options=self._form_compiler_options, - entity_maps=self._entity_maps, - ) - self._adjoint_solver._a = compiled_dFdu - self._adjoint_solver.solve() + adjoint_solver.solve() if isinstance(self._adjoint_solutions, list): - for adj_sol, sol in zip(self._adjoint_solutions, self._adjoint_solver.u): + for adj_sol, sol in zip(self._adjoint_solutions, adjoint_solver.u): adj_sol.x.array[:] = sol.x.array[:] else: assert isinstance(self._adjoint_solutions, dolfinx.fem.Function) - self._adjoint_solutions.x.array[:] = self._adjoint_solver.u.x.array[:] + self._adjoint_solutions.x.array[:] = adjoint_solver.u.x.array[:] return F_form, replacement_map def evaluate_adj_component( @@ -752,59 +680,93 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ if (hessian_inputs is None) or (len(tlm_output) == 0): return - # Using the equation Form we derive dF/du, d^2F/du^2 * du/dm * direction. - dFdu_form = self._compute_residual_derivative() - - # For linear forms d2Fdu2 is zero, but we include it for completeness. - - unknowns = [output.saved_output for output in self.get_outputs()] - summed_form = sum_form(dFdu_form) - d2Fdu2 = ufl.algorithms.expand_derivatives(ufl.derivative(summed_form, unknowns, tlm_output)) - - # bdy = self._should_compute_boundary_adjoint(relevant_dependencies) - - # Assemble right hand side of second order adjoint equation - # Note this term should always be zero for linear problems, but we include it for completeness. - if not d2Fdu2.empty(): - raise RuntimeError(f"This term {d2Fdu2:s} should be zero for linear problems.") - b_form = d2Fdu2 if d2Fdu2.empty() else ufl.action(ufl.adjoint(d2Fdu2), self._adjoint_solutions) - dFdu_adj = self._compute_adjoint(sum_form(dFdu_form)) - for bo in self.get_dependencies(): - c = bo.output - c_rep = bo.saved_output - tlm_input = bo.tlm_value - if tlm_input is None: - continue - if isinstance(c, (dolfinx.mesh.Mesh, dolfinx.fem.DirichletBC)): - raise NotImplementedError(f"Hessian computation for {type(c)} control not implemented yet.") - else: - summed_form = sum_form(dFdu_adj) - dFdu_adj_applied = ufl.action(summed_form, self._adjoint_solutions) - b_form += ufl.derivative(dFdu_adj_applied, c_rep, tlm_input) + # The adjoint solver -- and the compiled LHS it solves with, shared + # verbatim with the first-order adjoint equation (see + # LinearProblem._get_or_build_adjoint_solver) -- are shared across + # every block this Problem records. Refresh this block's own + # checkpointed values into the placeholders and re-establish this + # block's own bcs on every call, since another block may have used + # the same solver in between -- but never rebuild or recompile the + # LHS itself. + problem = self._problem() + adjoint_solver = problem._get_or_build_adjoint_solver() + adjoint_solver.bcs = self._bcs + for block_variable in self.get_dependencies(): + placeholder = problem._value_placeholders.get(block_variable.output) + if placeholder is not None: + placeholder.x.array[:] = block_variable.saved_output.x.array[:] + placeholder.x.scatter_forward() if len(outputs) == 1: - b = self._adjoint_solver._b + # Scalar problem: use the cached per-dependency Hessian templates + # (see LinearProblem._get_or_build_hessian_templates) instead of + # rebuilding and recompiling the SOA right-hand side symbolically + # on every call. + _, seed_placeholders, state_placeholder = problem._get_or_build_tlm_rhs_templates() + soa_templates, _, _ = problem._get_or_build_hessian_templates() + + state_placeholder.x.array[:] = outputs[0].saved_output.x.array[:] # type: ignore[union-attr] + state_placeholder.x.scatter_forward() # type: ignore[union-attr] + problem._adjoint_solution_placeholder.x.array[:] = self._adjoint_solutions.x.array[:] # type: ignore[union-attr] + problem._adjoint_solution_placeholder.x.scatter_forward() # type: ignore[union-attr] + problem._hessian_u_seed.x.array[:] = tlm_output[0].x.array[:] # type: ignore[union-attr] + problem._hessian_u_seed.x.scatter_forward() # type: ignore[union-attr] + + b = adjoint_solver._b with b.localForm() as b_loc: b_loc.set(0.0) - form_i = ufl.algorithms.apply_derivatives.apply_derivatives(b_form) - if not form_i.empty(): - compiled_soa_rhs = dolfinx.fem.form( - form_i, - jit_options=self._jit_options, - form_compiler_options=self._form_compiler_options, - entity_maps=self._entity_maps, - ) - dolfinx.fem.petsc.assemble_vector(b, compiled_soa_rhs) - b.ghostUpdate(PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) - - b.scale(-1) + for block_variable in self.get_dependencies(): + tlm_input = block_variable.tlm_value + if tlm_input is None: + continue + c = block_variable.output + if isinstance(c, (dolfinx.mesh.Mesh, dolfinx.fem.DirichletBC)): + raise NotImplementedError(f"Hessian computation for {type(c)} control not implemented yet.") + template = soa_templates.get(c) + if template is None: + continue + seed = seed_placeholders[c] + seed.x.array[:] = tlm_input.x.array[:] + seed.x.scatter_forward() + dolfinx.fem.petsc.assemble_vector(b, template) + dolfinx.la.petsc._ghost_update(b, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) + b.scale(-1) with b.localForm() as b_loc: b_loc.array[:] += hessian_inputs[0].array[:] b.ghostUpdate(PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) - self._adjoint_solver._b = b + adjoint_solver._b = b else: + # Blocked problem: Hessian cross-term templating is not + # implemented for this case (see the module-level plan notes); + # this keeps the pre-templating, per-call ufl.replace + recompile + # path, with the shared adjoint solver applying only the + # ownership-move. + dFdu_form = self._compute_residual_derivative() + unknowns = [output.saved_output for output in self.get_outputs()] + summed_form = sum_form(dFdu_form) + d2Fdu2 = ufl.algorithms.expand_derivatives(ufl.derivative(summed_form, unknowns, tlm_output)) + + # Assemble right hand side of second order adjoint equation + # Note this term should always be zero for linear problems, but we include it for completeness. + if not d2Fdu2.empty(): + raise RuntimeError(f"This term {d2Fdu2:s} should be zero for linear problems.") + b_form = d2Fdu2 + dFdu_adj = self._compute_adjoint(sum_form(dFdu_form)) + for bo in self.get_dependencies(): + c = bo.output + c_rep = bo.saved_output + tlm_input = bo.tlm_value + if tlm_input is None: + continue + if isinstance(c, (dolfinx.mesh.Mesh, dolfinx.fem.DirichletBC)): + raise NotImplementedError(f"Hessian computation for {type(c)} control not implemented yet.") + else: + summed_form = sum_form(dFdu_adj) + dFdu_adj_applied = ufl.action(summed_form, self._adjoint_solutions) + b_form += ufl.derivative(dFdu_adj_applied, c_rep, tlm_input) + bs = [] b_form = ufl.extract_blocks(b_form) for i, hess_input in enumerate(hessian_inputs): @@ -834,27 +796,26 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ bi.array[:] += hess_input.array bi.scatter_forward() - b = self._adjoint_solver._b + b = adjoint_solver._b local_arrays = [bi.array[: bi.index_map.size_local * bi.block_size] for bi in bs] dolfinx.la.petsc.assign(local_arrays, b) b.ghostUpdate(PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) - # Compile SOA LHS - dFdu_adj = dolfinx.fem.form( - self._compute_adjoint(sum_form(dFdu_form)), - jit_options=self._jit_options, - form_compiler_options=self._form_compiler_options, - entity_maps=self._entity_maps, - ) - - # Solve adjoint problem - self._adjoint_solver._a = dFdu_adj - self._adjoint_solver.solve() + # The SOA (second-order-adjoint) equation shares its LHS verbatim with + # the first-order adjoint equation (both are adjoint(dF/du), computed + # from dFdu_form above) -- already correct and permanent on + # adjoint_solver, so no rebuild or recompile needed here either. + adjoint_solver.solve() if isinstance(self._second_adjoint_solutions, list): - for adj_sol, sol in zip(self._second_adjoint_solutions, self._adjoint_solver.u): + for adj_sol, sol in zip(self._second_adjoint_solutions, adjoint_solver.u): adj_sol.x.array[:] = sol.x.array[:] else: - self._second_adjoint_solutions.x.array[:] = self._adjoint_solver.u.x.array[:] + self._second_adjoint_solutions.x.array[:] = adjoint_solver.u.x.array[:] + if len(outputs) == 1: + problem._second_adjoint_solution_placeholder.x.array[:] = ( # type: ignore[union-attr] + self._second_adjoint_solutions.x.array[:] + ) + problem._second_adjoint_solution_placeholder.x.scatter_forward() # type: ignore[union-attr] return self._compute_residual(), self._adjoint_solutions, self._second_adjoint_solutions @@ -895,6 +856,50 @@ def evaluate_hessian_component( assert isinstance(c, dolfinx.fem.Function) W = c.function_space + if len(outputs) == 1: + # Scalar problem: use the cached per-dependency Hessian templates + # (see LinearProblem._get_or_build_hessian_templates) instead of + # rebuilding and recompiling the Hessian-action output + # symbolically on every call. All the placeholders these + # templates reference (the dependency values, the state, and both + # adjoint solutions) were already refreshed by + # prepare_evaluate_hessian above; only the per-dependency + # "direction" seeds need setting here, and only for dependencies + # that actually have a tangent-linear value this call -- see + # LinearProblem._get_or_build_hessian_templates for why an + # inactive dependency's cross term must be skipped entirely + # rather than evaluated with a zeroed direction. + problem = self._problem() + _, fixed_templates, cross_templates = problem._get_or_build_hessian_templates() + _, seed_placeholders, _ = problem._get_or_build_tlm_rhs_templates() + + fixed_template = fixed_templates[c] + hessian_output = _create_vector(fixed_template, W) + hessian_output.array[:] = 0.0 + assemble_compiled_form(fixed_template, hessian_output) + + for _, bv in relevant_dependencies: + c2 = bv.output + if isinstance(c2, dolfinx.fem.DirichletBC): + continue + tlm_input = bv.tlm_value + if tlm_input is None: + continue + template = cross_templates.get((c, c2)) + if template is None: + continue + seed2 = seed_placeholders[c2] + seed2.x.array[:] = tlm_input.x.array[:] + seed2.x.scatter_forward() + assemble_compiled_form(template, hessian_output) + + hessian_output.array[:] *= -1.0 + return hessian_output + + # Blocked problem: Hessian cross-term templating is not implemented + # for this case (see the module-level plan notes); this keeps the + # pre-templating, per-call ufl.replace + recompile path. + # # We are trying to compute (dF/dm)^T lambda_1 # and (dF_dm)^T lambda_ 2. However, standard approach of UFL # does not work for MixedFunctionSpaces, as the control space is not @@ -970,15 +975,11 @@ def __init__( u: dolfinx.fem.Function | None = None, J: ufl.Form | None = None, P: ufl.Form | None = None, - kind: str | None = None, - petsc_options: dict | None = None, form_compiler_options: dict | None = None, jit_options: dict | None = None, entity_maps: typing.Sequence[dolfinx.mesh.EntityMap] | None = None, ad_block_tag: str | None = None, - adjoint_petsc_options: dict | None = None, - tlm_petsc_options: dict | None = None, - petsc_options_prefix: str = "dxa_nonlinear_block_", + problem: "NonlinearProblem" = ..., ) -> None: ... @typing.overload @@ -989,15 +990,11 @@ def __init__( u: typing.Sequence[dolfinx.fem.Function] | None = None, J: typing.Sequence[typing.Sequence[ufl.Form]] | None = None, P: typing.Sequence[typing.Sequence[ufl.Form]] | None = None, - kind: str | typing.Sequence[typing.Sequence[str]] | None = None, - petsc_options: dict | None = None, form_compiler_options: dict | None = None, jit_options: dict | None = None, entity_maps: typing.Sequence[dolfinx.mesh.EntityMap] | None = None, ad_block_tag: str | None = None, - adjoint_petsc_options: dict | None = None, - tlm_petsc_options: dict | None = None, - petsc_options_prefix: str = "dxa_nonlinear_block_", + problem: "NonlinearProblem" = ..., ) -> None: ... def __init__( @@ -1007,19 +1004,17 @@ def __init__( u: dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function] | None = None, J: ufl.Form | typing.Sequence[typing.Sequence[ufl.Form]] | None = None, P: ufl.Form | typing.Sequence[typing.Sequence[ufl.Form]] | None = None, - kind: str | typing.Sequence[typing.Sequence[str]] | None = None, - petsc_options: dict | None = None, form_compiler_options: dict | None = None, jit_options: dict | None = None, entity_maps: typing.Sequence[dolfinx.mesh.EntityMap] | None = None, ad_block_tag: str | None = None, - adjoint_petsc_options: dict | None = None, - tlm_petsc_options: dict | None = None, - petsc_options_prefix: str = "dxa_nonlinear_block_", + problem: "NonlinearProblem" = None, # type: ignore[assignment] ) -> None: - self._adjoint_petsc_options = adjoint_petsc_options - self._tlm_petsc_options = tlm_petsc_options + assert problem is not None, "problem must be provided." + # See LinearProblemBlock.__init__ for the rationale for holding a + # plain (strong) reference here. + self._problem_obj = problem super().__init__(ad_block_tag=ad_block_tag) self._preconditioner = P @@ -1056,25 +1051,12 @@ def __init__( self._jit_options = jit_options self._form_compiler_options = form_compiler_options self._entity_maps = entity_maps - self._petsc_options = petsc_options if petsc_options is not None else {} - self._petsc_options_prefix = petsc_options_prefix self._bcs = bcs if bcs is not None else [] - # Solver for recomputing the linear problem - self._forward_solver = dolfinx.fem.petsc.NonlinearProblem( - J=J, # type: ignore[arg-type] - F=self._rhs, # type: ignore[arg-type] - bcs=self._bcs, - u=self._u, # type: ignore[arg-type] - P=self._preconditioner, # type: ignore[arg-type] - petsc_options=self._petsc_options, - petsc_options_prefix=petsc_options_prefix, - form_compiler_options=self._form_compiler_options, - jit_options=self._jit_options, - kind=kind, # type: ignore[arg-type] - entity_maps=self._entity_maps, - ) # type: ignore[misc] - self._kind = "nest" if self._forward_solver.A.getType() == "nest" else kind + # No forward/adjoint solver is built here: this block shares the ones + # owned by self._problem() (see NonlinearProblem in ../solvers.py), + # built once and reused across every block that Problem records + # instead of once per solve() call. if isinstance(self._u, dolfinx.fem.Function): self._adjoint_solutions = self._u.copy() # type: ignore[assignment] @@ -1086,22 +1068,9 @@ def __init__( self._second_adjoint_solutions = [u.copy() for u in self._u] self._tlm_solutions = [u.copy() for u in self._u] - if isinstance(F, ufl.Form): - dFdu_adj = ufl.adjoint(ufl.derivative(F, u)) - else: - raise NotImplementedError("Blocked systems not implemented yet.") - self._adjoint_solver = LinearAdjointProblem( - dFdu_adj, # type: ignore[arg-type] - self._rhs, # type: ignore[arg-type] - bcs=self._bcs, - P=self._preconditioner, # type: ignore[arg-type] - form_compiler_options=self._form_compiler_options, - jit_options=self._jit_options, - petsc_options=self._adjoint_petsc_options, - petsc_options_prefix=self._petsc_options_prefix, - kind=kind, # type: ignore[arg-type] - entity_maps=self._entity_maps, - ) # type: ignore[misc] + def _problem(self) -> "NonlinearProblem": + """Return this block's owning Problem, which owns the shared solvers.""" + return self._problem_obj def _recover_bcs(self): bcs = [] @@ -1134,16 +1103,28 @@ def _replace_coefficients_in_form(self, form: ufl.Form) -> ufl.Form: def prepare_recompute_component(self, inputs, relevant_outputs): """Prepare for recomputing the block with different control inputs.""" - # As opposed to the linear problem, we need to update the coefficients in place, - # as the nonlinear problem snes.setContext doesn't reflect in place updates on the solver. + # The SNES the shared forward solver (self._problem()) owns is bound, + # forever, to a fixed set of placeholder coefficients rather than the + # user's own dependency objects (see + # NonlinearProblem._value_placeholders): writing this call's + # candidate/checkpointed values into the placeholders -- never into + # block_variable.output itself -- is what makes the SNES see them, + # without ever mutating an object the user (or a Taylor test + # perturbing a control directly) holds a live reference to. + problem = self._problem() for block_variable in self.get_dependencies(): - coeff = block_variable.output - if isinstance(coeff, dolfinx.fem.Function): - coeff.x.array[:] = block_variable.saved_output.x.array[:] - coeff.x.scatter_forward() + placeholder = problem._value_placeholders.get(block_variable.output) + if placeholder is not None: + placeholder.x.array[:] = block_variable.saved_output.x.array[:] + placeholder.x.scatter_forward() + + # Re-establish this block's own bcs on the shared forward solver (the + # Problem itself -- see _problem()), since another block may have used + # it with different bcs in between. + problem.bcs = self._bcs # Warm-start original unknown objects in place - u_list = self._forward_solver._u if isinstance(self._forward_solver._u, list) else [self._forward_solver._u] + u_list = problem._u if isinstance(problem._u, list) else [problem._u] for idx, out_bv in relevant_outputs: u_list[idx].x.array[:] = out_bv.saved_output.x.array[:] u_list[idx].x.scatter_forward() @@ -1154,12 +1135,15 @@ def recompute_component( self, inputs: typing.Iterable[Function], block_variable, idx: int, prepared: None ) -> Function: """Recompute the block with the prepared linear problem.""" + problem = self._problem() + # Call the base-class solve() directly (not problem.solve()), which + # would record another block onto the tape. with pyadjoint.tape.stop_annotating(): - self._forward_solver.solve() - if isinstance(self._forward_solver._u, list): - output = self._forward_solver._u[idx] + dolfinx.fem.petsc.NonlinearProblem.solve(problem) + if isinstance(problem._u, list): + output = problem._u[idx] else: - output = self._forward_solver._u + output = problem._u assert isinstance(output, Function) return output @@ -1247,70 +1231,74 @@ def _compute_residual_derivative(self) -> typing.Union[ufl.Form, list[list[ufl.F def prepare_evaluate_tlm( self, inputs, tlm_inputs, relevant_outputs - ) -> tuple[typing.Union[list[ufl.Form], ufl.Form], dolfinx.fem.Form]: - F_form = self._compute_residual() - - dFdu_compiled = dolfinx.fem.form( - self._compute_residual_derivative(), # type: ignore[arg-type] - jit_options=self._jit_options, - form_compiler_options=self._form_compiler_options, - entity_maps=self._entity_maps, - ) - return F_form, dFdu_compiled # type: ignore[return-value] - - def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepared=None) -> dolfinx.fem.Function: - """Solve the TLM equation for the block variable. - - .. math:: + ) -> typing.Union[dolfinx.fem.Function, typing.Sequence[dolfinx.fem.Function]]: + # The TLM solver -- and the compiled LHS it solves with, shared + # verbatim with dF/du (see NonlinearProblem._get_or_build_dFdu_template) + # -- are shared across every block this Problem records; likewise the + # per-dependency TLM right-hand-side templates (see + # NonlinearProblem._get_or_build_tlm_rhs_templates) are each compiled + # once. Refresh this block's own checkpointed values into the + # placeholders and re-establish this block's own bcs on every call, + # since another block may have used the same solver in between -- but + # never rebuild or recompile any of these forms. + problem = self._problem() + tlm_solver = problem._get_or_build_tlm_solver() + tlm_solver.bcs = self._bcs + templates, seed_placeholders, state_placeholder = problem._get_or_build_tlm_rhs_templates() - \frac{\\partial F}{\\partial u} \frac{\\partial u}{\\partial m} = \frac{\\partial F}{\\partial m} - - """ - F, dFdu = prepared - - V = self.get_outputs()[idx].output.function_space - - # FIXME: DirichletBC not block variable yet. Required later on. Currently all bcs should be homogenized - bcs = [] - for bc in self._bcs: - bcs.append(bc) - - dFdm = ufl.ZeroBaseForm((ufl.TestFunction(V),)) + for block_variable in self.get_dependencies(): + placeholder = problem._value_placeholders.get(block_variable.output) + if placeholder is not None: + placeholder.x.array[:] = block_variable.saved_output.x.array[:] + placeholder.x.scatter_forward() + out_bv = self.get_outputs()[0] + state_placeholder.x.array[:] = out_bv.saved_output.x.array[:] + state_placeholder.x.scatter_forward() + + # Assemble RHS vector using the shared solver's cached vector, + # accumulating only the dependencies that actually have a + # tangent-linear value this call -- see + # NonlinearProblem._get_or_build_tlm_rhs_templates for why an + # inactive dependency's term must be skipped entirely rather than + # evaluated with a zeroed direction. + b_petsc = tlm_solver._b + with b_petsc.localForm() as b_loc: + b_loc.set(0.0) for block_variable in self.get_dependencies(): tlm_value = block_variable.tlm_value - c_rep = block_variable.saved_output if tlm_value is None: continue - dFdm += ufl.derivative(-F, c_rep, tlm_value) + template = templates.get(block_variable.output) + if template is None: + continue + seed = seed_placeholders[block_variable.output] + seed.x.array[:] = tlm_value.x.array[:] + seed.x.scatter_forward() + dolfinx.fem.petsc.assemble_vector(b_petsc, template) + dolfinx.la.petsc._ghost_update(b_petsc, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) + + # Homogeneous boundary conditions are applied to b_petsc by + # tlm_solver.solve() itself (it homogenizes using tlm_solver.bcs, + # already re-established above), so there is no need to duplicate + # that here. + tlm_solver.solve() + if isinstance(self._tlm_solutions, list): + for tlm_sol, sol in zip(self._tlm_solutions, tlm_solver.u): + tlm_sol.x.array[:] = sol.x.array[:] + else: + assert isinstance(self._tlm_solutions, dolfinx.fem.Function) + self._tlm_solutions.x.array[:] = tlm_solver.u.x.array[:] - if isinstance(dFdm, float): - v = dFdu.arguments()[0] - dFdm = ufl.ZeroBaseForm((v,)) + return self._tlm_solutions - dFdm = ufl.algorithms.expand_derivatives(dFdm) - dFdm_compiled = dolfinx.fem.form( - dFdm, - jit_options=self._jit_options, - form_compiler_options=self._form_compiler_options, - entity_maps=self._entity_maps, - ) - dudm = dolfinx.fem.Function(V, name="du_dm_tlm_linearblock") - A_tlm = dolfinx.fem.petsc.assemble_matrix(dFdu, bcs=bcs) - A_tlm.assemble() - b_tlm = dolfinx.fem.create_vector(dolfinx.fem.extract_function_spaces(dFdm_compiled)) # type: ignore[arg-type] - b_tlm.array[:] = 0.0 - dolfinx.fem.petsc.assemble_vector(b_tlm.petsc_vec, dFdm_compiled) - - if bcs is not None: - # This system should never be "blocked" - dolfinx.fem.petsc.apply_lifting(b_tlm.petsc_vec, [dFdu], bcs=[bcs], alpha=0) - dolfinx.la.petsc._ghost_update(b_tlm.petsc_vec, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) # type: ignore [arg-type] - for bc in bcs: - bc.set(b_tlm.array, alpha=0) + def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepared=None) -> dolfinx.fem.Function: + # The system was solved natively in prepare_evaluate_tlm. + # Return the corresponding requested sub-function. + if isinstance(self._tlm_solutions, list): + return self._tlm_solutions[idx] else: - dolfinx.la.petsc._ghost_update(b_tlm, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) # type: ignore [arg-type] - solve_linear_problem(A_tlm, dudm.x, b_tlm, petsc_options=self._tlm_petsc_options) - return dudm + assert isinstance(self._tlm_solutions, dolfinx.fem.Function) + return self._tlm_solutions def prepare_evaluate_adj( self, @@ -1320,32 +1308,46 @@ def prepare_evaluate_adj( ) -> typing.Union[ufl.Form, typing.Iterable[ufl.Form]]: """Prepare the block for evaluating the adjoint.""" - # Compute (dF/du[v])* for the linear problem. + # The adjoint solver -- and the compiled LHS it solves with -- are + # shared across every block this Problem records: once every non-u + # dependency is routed through its placeholder and "u at this + # evaluation point" through its own dedicated placeholder (see + # NonlinearProblem._get_or_build_adjoint_solver), that operator is + # fixed for the life of the Problem and was compiled once. Refresh + # this block's own checkpointed values into the placeholders and + # re-establish this block's own bcs on every call, since another + # block may have used the same solver in between -- but never + # rebuild or recompile the LHS itself. + problem = self._problem() + adjoint_solver = problem._get_or_build_adjoint_solver() + adjoint_solver.bcs = self._bcs + for block_variable in self.get_dependencies(): + placeholder = problem._value_placeholders.get(block_variable.output) + if placeholder is not None: + placeholder.x.array[:] = block_variable.saved_output.x.array[:] + placeholder.x.scatter_forward() + out_bv = self.get_outputs()[0] + problem._state_placeholder.x.array[:] = out_bv.saved_output.x.array[:] # type: ignore[union-attr] + problem._state_placeholder.x.scatter_forward() # type: ignore[union-attr] + + # F_form is still needed by evaluate_adj_component (to build each + # dependency's own sensitivity form), but the adjoint LHS itself is + # already correct on adjoint_solver -- no rebuild, no recompile. F_form = self._compute_residual() - dFdu = self._compute_residual_derivative() - dFdu_adj = self._compute_adjoint(dFdu) # Extract dJ/du[v] from the adjoint inputs. assert len(adj_inputs) == 1 adj_rhs = adj_inputs[0] - dJdu = self._adjoint_solver._b + dJdu = adjoint_solver._b with dJdu.localForm() as dJdu_loc, adj_rhs.petsc_vec.localForm() as adj_rhs_loc: dJdu_loc.array[:] = adj_rhs_loc.array[:] - # Solve adjoint problem - compiled_dFdu = dolfinx.fem.form( - dFdu_adj, # type: ignore[arg-type] - jit_options=self._jit_options, - form_compiler_options=self._form_compiler_options, - entity_maps=self._entity_maps, - ) - self._adjoint_solver._a = compiled_dFdu - self._adjoint_solver.solve() + adjoint_solver.solve() if isinstance(self._adjoint_solutions, list): - for adj_sol, sol in zip(self._adjoint_solutions, self._adjoint_solver.u): + for adj_sol, sol in zip(self._adjoint_solutions, adjoint_solver.u): adj_sol.x.array[:] = sol.x.array[:] else: assert isinstance(self._adjoint_solutions, dolfinx.fem.Function) - self._adjoint_solutions.x.array[:] = self._adjoint_solver.u.x.array[:] + self._adjoint_solutions.x.array[:] = adjoint_solver.u.x.array[:] return F_form def evaluate_adj_component( @@ -1393,60 +1395,85 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ if (hessian_inputs is None) or (len(tlm_output) == 0): return - # Using the equation Form we derive dF/du, d^2F/du^2 * du/dm * direction. - dFdu_form = self._compute_residual_derivative() + # The adjoint solver -- and the compiled LHS it solves with, shared + # verbatim with the first-order adjoint equation (both are + # adjoint(dF/du), see NonlinearProblem._get_or_build_adjoint_solver) + # -- are shared across every block this Problem records. Refresh this + # block's own checkpointed values into the placeholders and + # re-establish this block's own bcs on every call, since another + # block may have used the same solver in between -- but never + # rebuild or recompile the LHS itself. + problem = self._problem() + adjoint_solver = problem._get_or_build_adjoint_solver() + adjoint_solver.bcs = self._bcs + for block_variable in self.get_dependencies(): + placeholder = problem._value_placeholders.get(block_variable.output) + if placeholder is not None: + placeholder.x.array[:] = block_variable.saved_output.x.array[:] + placeholder.x.scatter_forward() + problem._state_placeholder.x.array[:] = outputs[0].saved_output.x.array[:] # type: ignore[union-attr] + problem._state_placeholder.x.scatter_forward() # type: ignore[union-attr] + assert len(outputs) == 1, "Hessian computation only implemented for single output blocks." assert len(tlm_output) == 1, "Hessian computation only implemented for single TLM output blocks." - d2Fdu2 = ufl.algorithms.expand_derivatives(ufl.derivative(dFdu_form, outputs[0].saved_output, tlm_output[0])) - - # bdy = self._should_compute_boundary_adjoint(relevant_dependencies) assert len(hessian_inputs) == 1, "Hessian computation only implemented for single hessian input blocks." - # Assemble right hand side of second order adjoint equation - b_form = d2Fdu2 if d2Fdu2.empty() else ufl.action(ufl.adjoint(d2Fdu2), self._adjoint_solutions) - dFdu_adj = ufl.action(ufl.adjoint(dFdu_form), self._adjoint_solutions) - for bo in self.get_dependencies(): - c = bo.output - c_rep = bo.saved_output - tlm_input = bo.tlm_value + # The SOA solver -- and the compiled LHS it solves with, shared + # verbatim with the first-order adjoint equation -- is shared across + # every block this Problem records; likewise the SOA right-hand-side + # templates (see NonlinearProblem._get_or_build_hessian_templates) + # are each compiled once. Refresh the placeholders they reference and + # never rebuild or recompile any of these forms. + soa_self_template, soa_cross_templates, _, _ = problem._get_or_build_hessian_templates() + _, seed_placeholders, _ = problem._get_or_build_tlm_rhs_templates() + + problem._adjoint_solution_placeholder.x.array[:] = self._adjoint_solutions.x.array[:] # type: ignore[union-attr] + problem._adjoint_solution_placeholder.x.scatter_forward() # type: ignore[union-attr] + problem._hessian_u_seed.x.array[:] = tlm_output[0].x.array[:] # type: ignore[union-attr] + problem._hessian_u_seed.x.scatter_forward() # type: ignore[union-attr] + + b = adjoint_solver._b + with b.localForm() as b_loc: + b_loc.set(0.0) + if soa_self_template is not None: + dolfinx.fem.petsc.assemble_vector(b, soa_self_template) + for block_variable in self.get_dependencies(): + tlm_input = block_variable.tlm_value if tlm_input is None: continue + c = block_variable.output if isinstance(c, (dolfinx.mesh.Mesh, dolfinx.fem.DirichletBC)): raise NotImplementedError(f"Hessian computation for {type(c)} control not implemented yet.") - else: - b_form += ufl.derivative(dFdu_adj, c_rep, tlm_input) - b = self._adjoint_solver._b - with b.localForm() as b_loc: - b_loc.set(0.0) - if not ufl.algorithms.apply_derivatives.apply_derivatives(b_form).empty(): - compiled_soa_rhs = dolfinx.fem.form( - b_form, - jit_options=self._jit_options, - form_compiler_options=self._form_compiler_options, - entity_maps=self._entity_maps, - ) - dolfinx.fem.petsc.assemble_vector(b, compiled_soa_rhs) - b.ghostUpdate(PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) # type: ignore [arg-type] - b.scale(-1) + template = soa_cross_templates.get(c) + if template is None: + continue + seed = seed_placeholders[c] + seed.x.array[:] = tlm_input.x.array[:] + seed.x.scatter_forward() + dolfinx.fem.petsc.assemble_vector(b, template) + dolfinx.la.petsc._ghost_update(b, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) + b.scale(-1) with b.localForm() as b_loc, hessian_inputs[0].petsc_vec.localForm() as hess_loc: b_loc.array[:] += hess_loc.array[:] - # Compile SOA LHS - dFdu_adj = dolfinx.fem.form( - ufl.adjoint(dFdu_form), - jit_options=self._jit_options, - form_compiler_options=self._form_compiler_options, - entity_maps=self._entity_maps, - ) - - self._adjoint_solver._a = dFdu_adj - self._adjoint_solver._u = self._second_adjoint_solutions - self._adjoint_solver.solve() + # The SOA (second-order-adjoint) equation shares its LHS verbatim + # with the first-order adjoint equation (both are + # adjoint(dF/du) -- already correct and permanent on adjoint_solver, + # so no rebuild or recompile needed here either. + # + # Redirect the shared solver's scratch solution storage at this + # block's own second-adjoint buffer for the duration of this solve. + adjoint_solver._u = self._second_adjoint_solutions + adjoint_solver.solve() if isinstance(self._second_adjoint_solutions, list): - for adj_sol, sol in zip(self._second_adjoint_solutions, self._adjoint_solver.u): + for adj_sol, sol in zip(self._second_adjoint_solutions, adjoint_solver.u): adj_sol.x.array[:] = sol.x.array[:] else: - self._second_adjoint_solutions.x.array[:] = self._adjoint_solver.u.x.array[:] + self._second_adjoint_solutions.x.array[:] = adjoint_solver.u.x.array[:] + problem._second_adjoint_solution_placeholder.x.array[:] = ( # type: ignore[union-attr] + self._second_adjoint_solutions.x.array[:] + ) + problem._second_adjoint_solution_placeholder.x.scatter_forward() # type: ignore[union-attr] return self._compute_residual(), self._adjoint_solutions, self._second_adjoint_solutions def evaluate_hessian_component( @@ -1461,11 +1488,12 @@ def evaluate_hessian_component( ): c = block_variable.output - F_form, adj_sol, adj_sol2 = prepared - + # prepared (F_form, adj_sol, adj_sol2) is unused here: every quantity + # this method needs is already baked into the cached Hessian + # templates below, refreshed from those same values by + # prepare_evaluate_hessian. outputs = self.get_outputs() assert len(outputs) == 1, "Hessian computation only implemented for single output blocks." - tlm_output = outputs[0].tlm_value c_rep = block_variable.saved_output @@ -1487,54 +1515,40 @@ def evaluate_hessian_component( assert isinstance(c, dolfinx.fem.Function) W = c.function_space - dc = ufl.TestFunction(W) - form_adj = ufl.action(F_form, adj_sol) - form_adj2 = ufl.action(F_form, adj_sol2) - if isinstance(c, dolfinx.mesh.Mesh): - raise NotImplementedError("Hessian computation for Mesh control not implemented yet.") - # dFdm_adj = ufl.derivative(form_adj, X, dc) - # dFdm_adj2 = ufl.derivative(form_adj2, X, dc) - else: - # Assume Function - dFdm_adj = ufl.derivative(form_adj, c_rep, dc) - dFdm_adj2 = ufl.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(ufl.derivative(dFdm_adj, outputs[0].saved_output, tlm_output)) + # Use the cached per-dependency Hessian templates (see + # NonlinearProblem._get_or_build_hessian_templates) instead of + # rebuilding and recompiling the Hessian-action output symbolically + # on every call. All the placeholders these templates reference (the + # dependency values, the state, and both adjoint solutions) were + # already refreshed by prepare_evaluate_hessian above; only the + # per-dependency "direction" seeds need setting here, and only for + # dependencies that actually have a tangent-linear value this call -- + # see NonlinearProblem._get_or_build_hessian_templates for why an + # inactive dependency's cross term must be skipped entirely rather + # than evaluated with a zeroed direction. + problem = self._problem() + _, _, fixed_templates, cross_templates = problem._get_or_build_hessian_templates() + _, seed_placeholders, _ = problem._get_or_build_tlm_rhs_templates() + + fixed_template = fixed_templates[c] + hessian_output = _create_vector(fixed_template, W) + hessian_output.array[:] = 0.0 + assemble_compiled_form(fixed_template, hessian_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, dolfinx.fem.DirichletBC): continue tlm_input = bv.tlm_value if tlm_input is None: continue - - if c2 == self._u and not self.linear: + template = cross_templates.get((c, c2)) + if template is None: continue + seed2 = seed_placeholders[c2] + seed2.x.array[:] = tlm_input.x.array[:] + seed2.x.scatter_forward() + assemble_compiled_form(template, hessian_output) - # TODO: If tlm_input is a Sum, this crashes in some instances? - if isinstance(c2_rep, dolfinx.mesh.Mesh): - X = ufl.SpatialCoordinate(c2_rep) - d2Fdm2 += ufl.algorithms.expand_derivatives(ufl.derivative(dFdm_adj, X, tlm_input)) - else: - d2Fdm2 += ufl.algorithms.expand_derivatives(ufl.derivative(dFdm_adj, c2_rep, tlm_input)) - - hessian_form = ufl.algorithms.expand_derivatives(d2Fdm2 + dFdm_adj2 + d2Fdudm) - - compiled_hessian = dolfinx.fem.form( - hessian_form, - jit_options=self._jit_options, - form_compiler_options=self._form_compiler_options, - entity_maps=self._entity_maps, - ) - hessian_output = _create_vector(compiled_hessian, hessian_form.arguments()[0].ufl_function_space()) - hessian_output.array[:] = 0.0 - assemble_compiled_form(compiled_hessian, hessian_output) hessian_output.array[:] *= -1.0 return hessian_output diff --git a/src/dolfinx_adjoint/solvers.py b/src/dolfinx_adjoint/solvers.py index 2f9386b..27a9d13 100644 --- a/src/dolfinx_adjoint/solvers.py +++ b/src/dolfinx_adjoint/solvers.py @@ -7,10 +7,63 @@ import ufl from dolfinx.fem.function import Function as _Function -from .blocks.solvers import LinearProblemBlock, NonlinearProblemBlock +from .blocks.solvers import ( + LinearProblemBlock, + NonlinearProblemBlock, + assign_mixed_parts, + get_sorted_arguments, + sum_form, +) +from .petsc_utils import LinearAdjointProblem from .types import Function +def _collect_coefficients(form: ufl.Form | typing.Sequence | None) -> set: + """Return the set of UFL coefficients appearing anywhere in ``form``. + + ``form`` may be a single form or an arbitrarily nested sequence of forms + (entries may be ``None``, e.g. a zero block in a blocked system). Plain set + union rather than ``sum_form``: unlike summing, this never requires the + sub-forms' arguments to be mutually compatible (e.g. carry matching + ``part()`` tags), which blocked NonlinearProblem forms are not. + """ + if form is None: + return set() + if isinstance(form, ufl.Form): + return set(form.coefficients()) + coefficients: set = set() + for f in form: + coefficients |= _collect_coefficients(f) + return coefficients + + +def _replace_with_placeholders( + form: ufl.Form | typing.Sequence | None, placeholders: dict +) -> ufl.Form | typing.Sequence | None: + """Recursively apply ``ufl.replace(form, placeholders)`` to a (possibly nested) form + structure. + + A module-level function, not a nested closure: a nested function that + recurses by calling itself by name captures *itself* as a free variable, + which makes the function object (and, via ``self`` if the closure also + needs it) part of a reference cycle -- collected only by the cyclic + garbage collector, at a moment that differs between MPI ranks, not by + ordinary refcounting. That is exactly the hazard ``Problem`` owning its + solvers (rather than each ``Block``) exists to avoid: a self-referential + ``_replace`` closure inside ``LinearProblem``/``NonlinearProblem.__init__`` + would keep the ``Problem`` itself -- and its PETSc solvers -- alive as + cyclic garbage. Taking ``placeholders`` as a plain argument instead of + capturing ``self`` sidesteps this entirely: a module-level function + referring to itself by name is looked up through the module's namespace, + not a closure cell, so no cycle is created. + """ + if form is None: + return None + if isinstance(form, ufl.Form): + return ufl.replace(form, placeholders) + return [_replace_with_placeholders(f, placeholders) for f in form] + + @typing.overload def resolve_u(u: _Function | None, L: ufl.Form) -> _Function: ... @typing.overload @@ -117,6 +170,17 @@ def __init__( self.ad_block_tag = ad_block_tag self._adj_options = adjoint_petsc_options self._tlm_options = tlm_petsc_options + + # Assign mixed-space `part` indices to Test/Trial arguments once, + # here, for blocked systems (mirroring what LinearProblemBlock used to + # redo per block): needed so a blocked bilinear/linear form can be + # safely combined into one whole-system form (via sum_form) when + # building the adjoint solver below. + if not isinstance(a, ufl.Form): + a, L = assign_mixed_parts(a, L) # type: ignore[arg-type] + if P is not None: + P, _ = assign_mixed_parts(P, L) # type: ignore[arg-type] + self._u = resolve_u(u, L) # type: ignore[arg-type] # Cache some objects @@ -126,15 +190,37 @@ def __init__( self._form_compiler_options = form_compiler_options self._entity_maps = entity_maps self._petsc_options = petsc_options + self._petsc_options_prefix = petsc_options_prefix self._kind = kind + # The forward solver's compiled forms reference dedicated placeholder + # coefficients rather than the user's own dependency objects -- + # exactly like NonlinearProblem, so both classes share the same + # data-handling story: a solve always means "refresh the + # placeholders' values, then call the solver", never "recompile a + # form" or "mutate the user's own coefficient in place". solve() + # (below) refreshes them from the user's own current values; + # LinearProblemBlock.prepare_recompute_component refreshes them from + # a block's checkpointed/candidate values instead. Neither ever + # writes into the user's own coefficient objects, so a Taylor test + # that perturbs the original control directly + # (`pyadjoint.taylor_test(Jh, m, dm)`) always sees a pristine `m`. + u_list = self._u if isinstance(self._u, list) else [self._u] + coefficients = _collect_coefficients(a) | _collect_coefficients(L) + if P is not None: + coefficients |= _collect_coefficients(P) + coefficients -= set(u_list) + self._value_placeholders: dict[dolfinx.fem.Function, dolfinx.fem.Function] = { + c: dolfinx.fem.Function(c.function_space) for c in coefficients + } + # Initialize linear solver super().__init__( - a=a, # type: ignore[arg-type] - L=L, # type: ignore[arg-type] + a=_replace_with_placeholders(a, self._value_placeholders), # type: ignore[arg-type] + L=_replace_with_placeholders(L, self._value_placeholders), # type: ignore[arg-type] bcs=bcs, u=self._u, # type: ignore[arg-type] - P=P, # type: ignore[arg-type] + P=_replace_with_placeholders(P, self._value_placeholders), # type: ignore[arg-type] kind=kind, # type: ignore[arg-type] petsc_options_prefix=petsc_options_prefix, petsc_options=petsc_options, @@ -143,6 +229,317 @@ def __init__( entity_maps=entity_maps, ) # type: ignore[misc] + # Match the adjoint/TLM solvers' matrix layout to whatever `kind` the + # forward solver actually resolved to (kind=None can auto-resolve to + # "nest" for blocked problems). + self._kind = "nest" if self.A.getType() == "nest" else kind + + # Adjoint and tangent-linear solvers: built lazily (on first use, see + # _get_or_build_adjoint_solver/_get_or_build_tlm_solver below) and shared + # by every LinearProblemBlock this Problem records, rather than one per + # block/solve() call. Blocks only ever hold a weak reference back to this + # Problem (see LinearProblemBlock._problem), so dropping this Problem + # releases the forward, adjoint and TLM solvers' PETSc objects + # deterministically instead of leaving that to pyadjoint's tape/cyclic-GC + # schedule -- see the "mpi-collective-destruction-hazard" note in the + # dolfinx-adjoint-knowledge repository for why that matters. Laziness + # keeps pure forward (non-annotated) use from paying for a symbolic + # adjoint form it never needs. + self._adjoint_solver: typing.Optional[LinearAdjointProblem] = None + self._tlm_solver: typing.Optional[LinearAdjointProblem] = None + self._dFdu_template: typing.Optional[ufl.Form | typing.Sequence] = None + self._dFdu_adj_template: typing.Optional[ufl.Form | typing.Sequence] = None + self._tlm_rhs_templates: typing.Optional[dict] = None + self._tlm_seed_placeholders: dict[dolfinx.fem.Function, dolfinx.fem.Function] = {} + self._residual_state_placeholder: typing.Union[ + dolfinx.fem.Function, typing.Sequence[dolfinx.fem.Function], None + ] = None + self._hessian_templates: typing.Optional[tuple[dict, dict, dict]] = None + self._adjoint_solution_placeholder: typing.Optional[dolfinx.fem.Function] = None + self._second_adjoint_solution_placeholder: typing.Optional[dolfinx.fem.Function] = None + self._hessian_u_seed: typing.Optional[dolfinx.fem.Function] = None + + def _get_or_build_dFdu_template(self) -> ufl.Form | typing.Sequence: + """Build (once) and return dF/du with every non-u coefficient replaced by its + placeholder. + + dF/du does not actually depend on the state u for a linear problem: + F(u, v) = a(u, v) - L(v) is linear in u, so its derivative doesn't + reference u's value at all, only whatever *other* coefficients a + itself depends on. This is exactly a (placeholder-substituted), and + is the shared basis for both the adjoint operator + (``_get_or_build_adjoint_solver``, which just adjoints it) and the + TLM operator (``_get_or_build_tlm_solver``, used as-is): built once, + for the life of this Problem, so neither ever needs to rebuild or + recompile it -- only refresh the placeholders' values (see + ``LinearProblemBlock.prepare_evaluate_adj``/``prepare_evaluate_hessian``/``prepare_evaluate_tlm``). + """ + if self._dFdu_template is None: + self._dFdu_template = ufl.replace(sum_form(self._lhs), self._value_placeholders) # type: ignore[arg-type] + return self._dFdu_template + + def _get_or_build_dFdu_adj_template(self) -> ufl.Form | typing.Sequence: + """Build (once) and return adjoint(dF/du), shared by the adjoint solver + (``_get_or_build_adjoint_solver``) and, for scalar problems, the Hessian + SOA right-hand-side's cross-dependency templates + (``_get_or_build_hessian_templates``). + + Kept exactly as ``_compute_adjoint`` returns it -- a nested list of + forms for a blocked problem -- since that structure is what + ``LinearAdjointProblem``/``dolfinx.fem.petsc.LinearProblem`` needs for + block matrix assembly; callers that need a single summed form (Hessian + templating, scalar-only) apply ``sum_form`` themselves. + """ + if self._dFdu_adj_template is None: + self._dFdu_adj_template = LinearProblemBlock._compute_adjoint( + self._get_or_build_dFdu_template() # type: ignore[arg-type] + ) + return self._dFdu_adj_template + + def _get_or_build_adjoint_solver(self) -> LinearAdjointProblem: + """Build (once) and return the adjoint solver shared by every block this Problem records.""" + if self._adjoint_solver is None: + self._adjoint_solver = LinearAdjointProblem( + self._get_or_build_dFdu_adj_template(), # type: ignore[arg-type] + self._rhs, # type: ignore[arg-type] + bcs=self.bcs, + P=self._preconditioner, # type: ignore[arg-type] + form_compiler_options=self._form_compiler_options, + jit_options=self._jit_options, + petsc_options=self._adj_options, + petsc_options_prefix=f"{self._petsc_options_prefix}adjoint_", + kind=self._kind, # type: ignore[arg-type] + entity_maps=self._entity_maps, + ) # type: ignore[misc] + return self._adjoint_solver + + def _get_or_build_tlm_solver(self) -> LinearAdjointProblem: + """Build (once) and return the TLM solver shared by every block this Problem records. + + No explicit ``u=`` is passed: like the adjoint solver, this gets its own + scratch solution Function from the base class, and callers copy the + result out (see ``LinearProblemBlock.prepare_evaluate_tlm``) rather than + relying on solver-owned storage identity, since that storage is now + shared across every block instead of private to one. + + Unlike the adjoint operator (which decomposes dF/du back into blocks + itself, inside ``compute_form_adjoint``/``_compute_adjoint``), dF/du + is used here as-is, so for a blocked problem it must be decomposed + with ``ufl.extract_blocks`` before compiling: a summed multi-part + form is a perfectly good UFL object to keep substituting into and + differentiating, but it is not, on its own, a compilable one -- the + parts must be split apart first (mirroring + ``_compute_residual_derivative``'s ``ufl.extract_blocks(dFdu)`` in the + pre-templating code this replaces). + """ + if self._tlm_solver is None: + dFdu_template = self._get_or_build_dFdu_template() + if isinstance(self._u, list): + dFdu_template = ufl.extract_blocks(dFdu_template) # type: ignore[arg-type] + self._tlm_solver = LinearAdjointProblem( + dFdu_template, # type: ignore[arg-type] + self._rhs, # type: ignore[arg-type] + bcs=self.bcs, + P=self._preconditioner, # type: ignore[arg-type] + form_compiler_options=self._form_compiler_options, + jit_options=self._jit_options, + petsc_options=self._tlm_options, + petsc_options_prefix=f"{self._petsc_options_prefix}tlm_", + kind=self._kind, # type: ignore[arg-type] + entity_maps=self._entity_maps, + ) # type: ignore[misc] + return self._tlm_solver + + def _get_or_build_tlm_rhs_templates( + self, + ) -> tuple[ + dict[dolfinx.fem.Function, typing.Any], + dict[dolfinx.fem.Function, dolfinx.fem.Function], + typing.Union[dolfinx.fem.Function, typing.Sequence[dolfinx.fem.Function]], + ]: + """Build (once) and return the per-dependency TLM right-hand-side templates. + + Unlike dF/du, dF/dm genuinely depends on the state u even for a + linear problem (a is bilinear, so differentiating w.r.t. a + coefficient embedded in a while holding u fixed leaves u in the + result), so this needs its own "current state" placeholder, + ``_residual_state_placeholder``, distinct from the live self._u the + forward solve owns. + + One compiled one-form is built per dependency, using a dedicated + "direction" placeholder for that dependency (``_tlm_seed_placeholders``) + rather than a single combined form summed over every dependency: + summing symbolically would require deciding, once and for all, which + dependencies contribute, but which ones actually have a tangent-linear + value varies from call to call. Refreshing an unused dependency's + seed to zero and evaluating its term anyway is not a safe substitute + for skipping it: if that dependency appears in a way that is singular + at its current value (e.g. a `1/c` term, with `c` legitimately zero + somewhere in the domain), the assembled contribution would be `0 * + inf = NaN` there even though the *seed* is zero, silently corrupting + the sum. Keeping every dependency's contribution as its own compiled + form, only ever assembled when that dependency actually has a + tangent-linear value (see ``LinearProblemBlock.prepare_evaluate_tlm``), + avoids that entirely by never evaluating an inactive dependency's term + at all -- exactly matching what skipping it symbolically did before. + """ + if self._tlm_rhs_templates is None: + u_list = self._u if isinstance(self._u, list) else [self._u] + if isinstance(self._u, list): + self._residual_state_placeholder = [ + dolfinx.fem.Function(ui.function_space) # type: ignore[union-attr] + for ui in u_list + ] + state_arg: typing.Any = self._residual_state_placeholder + else: + self._residual_state_placeholder = dolfinx.fem.Function(self._u.function_space) # type: ignore[union-attr] + state_arg = self._residual_state_placeholder + + a_template = self._get_or_build_dFdu_template() + L_template = ufl.replace(sum_form(self._rhs), self._value_placeholders) # type: ignore[arg-type] + F_template = ufl.action(a_template, state_arg) - L_template # type: ignore[arg-type] + + if isinstance(self._u, list): + test_funcs = list(get_sorted_arguments(F_template.arguments(), 0)) + else: + test_funcs = [F_template.arguments()[0]] + + templates: dict[dolfinx.fem.Function, typing.Any] = {} + for c, c_placeholder in self._value_placeholders.items(): + seed = dolfinx.fem.Function(c.function_space) + dFdm_c = ufl.algorithms.expand_derivatives(-ufl.derivative(F_template, c_placeholder, seed)) + if isinstance(self._u, list): + blocks = ufl.extract_blocks(dFdm_c) + padded = [ufl.ZeroBaseForm((test,)) for test in test_funcs] + for block in blocks: + args = block.arguments() + assert len(args) == 1, "Expected a single test function in the block." + padded[args[0].part()] = block + dFdm_c = padded + else: + if dFdm_c == 0 or dFdm_c.empty(): + dFdm_c = ufl.ZeroBaseForm((test_funcs[0],)) + templates[c] = dolfinx.fem.form( + dFdm_c, + jit_options=self._jit_options, + form_compiler_options=self._form_compiler_options, + entity_maps=self._entity_maps, + ) + self._tlm_seed_placeholders[c] = seed + self._tlm_rhs_templates = templates + return self._tlm_rhs_templates, self._tlm_seed_placeholders, self._residual_state_placeholder # type: ignore[return-value] + + def _get_or_build_hessian_templates(self) -> tuple[dict, dict, dict]: + """Build (once) and return the per-dependency Hessian templates used by + ``LinearProblemBlock.prepare_evaluate_hessian``'s SOA right-hand side and + ``evaluate_hessian_component``'s own Hessian-action output. + + Scalar (non-blocked) problems only -- the blocked Hessian cross-term + stays on the pre-templating, per-call ``ufl.replace`` + recompile path + (see the module-level plan notes: this is deferred, ownership-move + only, for the blocked case). + + Three dicts are returned: + + - ``soa_templates[c]``: the SOA right-hand-side's contribution from + dependency ``c``'s tangent-linear direction, a 1-form in the state's + own test function. + - ``fixed_templates[c]``: the part of dependency ``c``'s own + Hessian-action output that does not depend on any *other* + dependency's tangent-linear value (``dL2dm + d2Fdudm``) -- always + assembled. + - ``cross_templates[(c, c2)]``: dependency ``c``'s Hessian-action + contribution from *another* dependency ``c2``'s tangent-linear + direction (``d2Fdm2``). + + Each is kept as its own compiled one-form, using a dedicated + "direction" placeholder (the same ``_tlm_seed_placeholders`` the TLM + right-hand side already uses -- safe to share, since the TLM forward + sweep has always finished computing every tangent-linear value before + the reverse (adjoint/Hessian) sweep that needs these runs), for the + same reason as ``_get_or_build_tlm_rhs_templates``: summing every + dependency's cross-term contribution into one combined form and + zeroing an inactive dependency's seed has the same ``0 * inf = NaN`` + hazard there does. + """ + if self._hessian_templates is None: + assert not isinstance(self._u, list), "Hessian templating is only implemented for scalar problems." + _, seed_placeholders, state_placeholder = self._get_or_build_tlm_rhs_templates() + dFdu_template = self._get_or_build_dFdu_template() + dFdu_adj_template = sum_form(self._get_or_build_dFdu_adj_template()) # type: ignore[arg-type] + assert isinstance(dFdu_template, ufl.Form) + assert isinstance(dFdu_adj_template, ufl.Form) + assert isinstance(state_placeholder, dolfinx.fem.Function) + + self._adjoint_solution_placeholder = dolfinx.fem.Function(self._u.function_space) # type: ignore[union-attr] + self._second_adjoint_solution_placeholder = dolfinx.fem.Function( + self._u.function_space # type: ignore[union-attr] + ) + self._hessian_u_seed = dolfinx.fem.Function(self._u.function_space) # type: ignore[union-attr] + + L_template = ufl.replace(sum_form(self._rhs), self._value_placeholders) # type: ignore[arg-type] + F_template = ufl.action(dFdu_template, state_placeholder) - L_template # type: ignore[arg-type] + + # dF/du does not depend on u for a linear problem, so its second + # derivative w.r.t. u is always exactly zero: there is no SOA + # "self" term to template here, only the cross-dependency terms + # below (contrast NonlinearProblem, where F is nonlinear in u). + # Verify this invariant once, here, rather than on every call. + d2Fdu2_check = ufl.algorithms.expand_derivatives( + ufl.derivative(dFdu_template, state_placeholder, self._hessian_u_seed) + ) + if not d2Fdu2_check.empty(): + raise RuntimeError(f"This term {d2Fdu2_check} should be zero for linear problems.") + + dFdu_adj_applied = ufl.action(dFdu_adj_template, self._adjoint_solution_placeholder) + L1 = ufl.action(F_template, self._adjoint_solution_placeholder) + L2 = ufl.action(F_template, self._second_adjoint_solution_placeholder) + + soa_templates: dict = {} + fixed_templates: dict = {} + cross_templates: dict = {} + for c, c_placeholder in self._value_placeholders.items(): + seed = seed_placeholders[c] + + soa_form = ufl.algorithms.expand_derivatives(ufl.derivative(dFdu_adj_applied, c_placeholder, seed)) + if not (soa_form == 0 or soa_form.empty()): + soa_templates[c] = dolfinx.fem.form( + soa_form, + jit_options=self._jit_options, + form_compiler_options=self._form_compiler_options, + entity_maps=self._entity_maps, + ) + + dc = ufl.TestFunction(c.function_space) + dL1dm = ufl.derivative(L1, c_placeholder, dc) + dL2dm = ufl.derivative(L2, c_placeholder, dc) + d2Fdudm = ufl.algorithms.expand_derivatives( + ufl.derivative(dL1dm, state_placeholder, self._hessian_u_seed) + ) + fixed_form = ufl.algorithms.expand_derivatives(dL2dm + d2Fdudm) + if fixed_form == 0 or fixed_form.empty(): + fixed_form = ufl.ZeroBaseForm((dc,)) + fixed_templates[c] = dolfinx.fem.form( + fixed_form, + jit_options=self._jit_options, + form_compiler_options=self._form_compiler_options, + entity_maps=self._entity_maps, + ) + + for c2, c2_placeholder in self._value_placeholders.items(): + seed2 = seed_placeholders[c2] + cross_form = ufl.algorithms.expand_derivatives(ufl.derivative(dL1dm, c2_placeholder, seed2)) + if cross_form == 0 or cross_form.empty(): + continue + cross_templates[(c, c2)] = dolfinx.fem.form( + cross_form, + jit_options=self._jit_options, + form_compiler_options=self._form_compiler_options, + entity_maps=self._entity_maps, + ) + self._hessian_templates = (soa_templates, fixed_templates, cross_templates) + return self._hessian_templates + def solve(self, annotate: bool = True) -> typing.Union[dolfinx.fem.Function, typing.Sequence[dolfinx.fem.Function]]: """ Solve the linear problem and return the solution. @@ -155,18 +552,23 @@ def solve(self, annotate: bool = True) -> typing.Union[dolfinx.fem.Function, typ bcs=self.bcs, u=self.u, # type: ignore[arg-type] P=self._preconditioner, # type: ignore[arg-type] - kind=self._kind, # type: ignore[arg-type] - petsc_options=self._petsc_options, form_compiler_options=self._form_compiler_options, jit_options=self._jit_options, entity_maps=self._entity_maps, ad_block_tag=self.ad_block_tag, - adjoint_petsc_options=self._adj_options, - tlm_petsc_options=self._tlm_options, + problem=self, ) # type: ignore[misc] tape = pyadjoint.get_working_tape() tape.add_block(block) + # Refresh the forward solver's placeholder coefficients from the + # user's own, current values before an ordinary solve: a prior + # recompute (see LinearProblemBlock.prepare_recompute_component) may + # have left them holding a checkpointed/candidate value instead. + for original, placeholder in self._value_placeholders.items(): + placeholder.x.array[:] = original.x.array[:] + placeholder.x.scatter_forward() + out = dolfinx.fem.petsc.LinearProblem.solve(self) if annotate: if isinstance(out, Function): @@ -268,13 +670,43 @@ def __init__( self._form_compiler_options = form_compiler_options self._entity_maps = entity_maps self._petsc_options = petsc_options + self._petsc_options_prefix = petsc_options_prefix self._kind = kind - # Initialize linear solver + # The SNES built by super().__init__() below binds to the exact + # compiled F/J/P Form objects passed to it, forever: its residual and + # Jacobian callbacks close over those objects in a context dict set up + # once (see dolfinx.fem.petsc.NonlinearProblem.__init__'s + # jacobian_ctx/function_ctx), so reassigning self._F/self._J later -- + # the trick LinearProblem.solve() uses to switch between "live" and + # "recompute" forms -- would have no effect on what the SNES actually + # assembles. The only way to make the SNES see a different value for a + # coefficient is to mutate the exact Function object its compiled + # forms reference. + # + # To keep that mutation from ever touching an object the user (or a + # Taylor test perturbing a control directly) holds a live reference + # to, every non-u coefficient is routed through a dedicated + # placeholder Function from the very start: the SNES is built against + # F/J/P with every such coefficient replaced by its placeholder, and + # the placeholders are (re)populated -- from the user's own current + # values for an ordinary solve() (see solve() below), or from a + # block's checkpointed/candidate values for a recompute (see + # NonlinearProblemBlock.prepare_recompute_component) -- before every + # solve, never the other way around. + u_list = self._u if isinstance(self._u, list) else [self._u] + coefficients = _collect_coefficients(F) - set(u_list) + if J is not None: + coefficients |= _collect_coefficients(J) - set(u_list) + self._value_placeholders: dict[dolfinx.fem.Function, dolfinx.fem.Function] = { + c: dolfinx.fem.Function(c.function_space) for c in coefficients + } + + # Initialize nonlinear solver super().__init__( - F=F, # type: ignore[arg-type] - J=J, # type: ignore[arg-type] - P=P, # type: ignore[arg-type] + F=_replace_with_placeholders(F, self._value_placeholders), # type: ignore[arg-type] + J=_replace_with_placeholders(J, self._value_placeholders), # type: ignore[arg-type] + P=_replace_with_placeholders(P, self._value_placeholders), # type: ignore[arg-type] bcs=self._bcs, u=self._u, # type: ignore[arg-type] kind=kind, # type: ignore[arg-type] @@ -285,6 +717,241 @@ def __init__( entity_maps=entity_maps, ) # type: ignore[misc] + # Adjoint and tangent-linear solvers: built lazily (see + # _get_or_build_adjoint_solver/_get_or_build_tlm_solver) and shared by + # every NonlinearProblemBlock this Problem records, rather than one + # per block/solve() call -- same rationale as LinearProblem. + self._adjoint_solver: typing.Optional[LinearAdjointProblem] = None + self._tlm_solver: typing.Optional[LinearAdjointProblem] = None + self._dFdu_template: typing.Optional[ufl.Form] = None + self._state_placeholder: typing.Optional[dolfinx.fem.Function] = None + self._tlm_rhs_templates: typing.Optional[dict] = None + self._tlm_seed_placeholders: dict[dolfinx.fem.Function, dolfinx.fem.Function] = {} + self._hessian_templates: typing.Optional[tuple] = None + self._adjoint_solution_placeholder: typing.Optional[dolfinx.fem.Function] = None + self._second_adjoint_solution_placeholder: typing.Optional[dolfinx.fem.Function] = None + self._hessian_u_seed: typing.Optional[dolfinx.fem.Function] = None + + def _get_or_build_dFdu_template(self) -> ufl.Form: + """Build (once) and return dF/du with every non-u coefficient replaced by its + placeholder, and u itself replaced by a dedicated "state" placeholder + standing in for "u at this evaluation point". + + Unlike the linear case, dF/du genuinely depends on u's current value + here (F is nonlinear in u), so it needs a coefficient slot for that -- + but that slot need not be ``self._u`` itself (which the live + SNES/forward path owns): a dedicated placeholder, refreshed from a + block's own checkpointed output before each adjoint/TLM/Hessian solve + (see ``NonlinearProblemBlock.prepare_evaluate_adj``/ + ``prepare_evaluate_tlm``/``prepare_evaluate_hessian``), keeps this + operator's compiled form fixed for the life of the Problem, exactly + like the non-u dependencies already routed through + ``self._value_placeholders``. Shared, verbatim, by the adjoint + operator (which just adjoints it) and the TLM operator (used as-is), + mirroring ``LinearProblem._get_or_build_dFdu_template``. + """ + if self._dFdu_template is None: + if not isinstance(self._rhs, ufl.Form): + raise NotImplementedError("Blocked systems not implemented yet.") + self._state_placeholder = dolfinx.fem.Function(self._u.function_space) # type: ignore[union-attr] + replace_map: dict = {**self._value_placeholders, self._u: self._state_placeholder} + self._dFdu_template = ufl.replace(self._lhs, replace_map) # type: ignore[arg-type] + return self._dFdu_template + + def _get_or_build_adjoint_solver(self) -> LinearAdjointProblem: + """Build (once) and return the adjoint solver shared by every block this Problem records.""" + if self._adjoint_solver is None: + dFdu_adj = ufl.adjoint(self._get_or_build_dFdu_template()) + self._adjoint_solver = LinearAdjointProblem( + dFdu_adj, # type: ignore[arg-type] + self._rhs, # type: ignore[arg-type] + bcs=self._bcs, + P=self._preconditioner, # type: ignore[arg-type] + form_compiler_options=self._form_compiler_options, + jit_options=self._jit_options, + petsc_options=self._adj_options, + petsc_options_prefix=f"{self._petsc_options_prefix}adjoint_", + kind=self._kind, # type: ignore[arg-type] + entity_maps=self._entity_maps, + ) # type: ignore[misc] + return self._adjoint_solver + + def _get_or_build_tlm_solver(self) -> LinearAdjointProblem: + """Build (once) and return the TLM solver shared by every block this Problem records. + + dF/du is used here as-is (unlike the adjoint operator, which adjoints + it), mirroring ``LinearProblem._get_or_build_tlm_solver``. + """ + if self._tlm_solver is None: + self._tlm_solver = LinearAdjointProblem( + self._get_or_build_dFdu_template(), # type: ignore[arg-type] + self._rhs, # type: ignore[arg-type] + bcs=self._bcs, + P=self._preconditioner, # type: ignore[arg-type] + form_compiler_options=self._form_compiler_options, + jit_options=self._jit_options, + petsc_options=self._tlm_options, + petsc_options_prefix=f"{self._petsc_options_prefix}tlm_", + kind=self._kind, # type: ignore[arg-type] + entity_maps=self._entity_maps, + ) # type: ignore[misc] + return self._tlm_solver + + def _get_or_build_tlm_rhs_templates( + self, + ) -> tuple[ + dict[dolfinx.fem.Function, typing.Any], + dict[dolfinx.fem.Function, dolfinx.fem.Function], + dolfinx.fem.Function, + ]: + """Build (once) and return the per-dependency TLM right-hand-side templates. + + One compiled one-form is built per dependency, using a dedicated + "direction" placeholder for that dependency + (``_tlm_seed_placeholders``) rather than a single combined form summed + over every dependency, for the same reason as + ``LinearProblem._get_or_build_tlm_rhs_templates``: which dependencies + actually have a tangent-linear value varies from call to call, and + evaluating an inactive dependency's term with a zeroed seed instead of + skipping it outright risks ``0 * inf = NaN`` if that dependency's + derivative is singular where it is currently valued (e.g. a `1/c` + term with `c` legitimately zero somewhere in the domain). + """ + if self._tlm_rhs_templates is None: + self._get_or_build_dFdu_template() # ensures self._state_placeholder exists + assert self._state_placeholder is not None + assert isinstance(self._rhs, ufl.Form) + replace_map: dict = {**self._value_placeholders, self._u: self._state_placeholder} + F_template = ufl.replace(self._rhs, replace_map) + test_func = F_template.arguments()[0] + + templates: dict[dolfinx.fem.Function, typing.Any] = {} + for c, c_placeholder in self._value_placeholders.items(): + seed = dolfinx.fem.Function(c.function_space) + dFdm_c = ufl.algorithms.expand_derivatives(-ufl.derivative(F_template, c_placeholder, seed)) + if dFdm_c == 0 or dFdm_c.empty(): + dFdm_c = ufl.ZeroBaseForm((test_func,)) + templates[c] = dolfinx.fem.form( + dFdm_c, + jit_options=self._jit_options, + form_compiler_options=self._form_compiler_options, + entity_maps=self._entity_maps, + ) + self._tlm_seed_placeholders[c] = seed + self._tlm_rhs_templates = templates + return self._tlm_rhs_templates, self._tlm_seed_placeholders, self._state_placeholder # type: ignore[return-value] + + def _get_or_build_hessian_templates( + self, + ) -> tuple[typing.Optional[dolfinx.fem.Form], dict, dict, dict]: + """Build (once) and return the per-dependency Hessian templates used by + ``NonlinearProblemBlock.prepare_evaluate_hessian``'s SOA right-hand side + and ``evaluate_hessian_component``'s own Hessian-action output. + + Four values are returned: + + - ``soa_self_template``: the SOA right-hand-side's contribution from + dF/du's own second derivative w.r.t. u (``d2Fdu2``) -- genuinely + nonzero here since F is nonlinear in u (contrast + ``LinearProblem._get_or_build_hessian_templates``, where this term + is always zero and skipped entirely) -- or ``None`` if it happens + to vanish structurally. Always assembled when not ``None``. + - ``soa_cross_templates[c]``: the SOA right-hand-side's contribution + from dependency ``c``'s tangent-linear direction, a 1-form in the + state's own test function. + - ``fixed_templates[c]``: the part of dependency ``c``'s own + Hessian-action output that does not depend on any *other* + dependency's tangent-linear value (``dL2dm + d2Fdudm``) -- always + assembled. + - ``cross_templates[(c, c2)]``: dependency ``c``'s Hessian-action + contribution from *another* dependency ``c2``'s tangent-linear + direction (``d2Fdm2``). + + Mirrors ``LinearProblem._get_or_build_hessian_templates`` exactly for + the cross-dependency terms (same per-dependency-template rationale, + including the ``0 * inf = NaN`` hazard of a combined, zeroed-seed + form), with one addition: the SOA "self" term, which for a linear + problem is always zero and so needs no template at all. + """ + if self._hessian_templates is None: + dFdu_template = self._get_or_build_dFdu_template() + dFdu_adj_template = ufl.adjoint(dFdu_template) + assert self._state_placeholder is not None + state_placeholder = self._state_placeholder + assert isinstance(self._rhs, ufl.Form) + + self._adjoint_solution_placeholder = dolfinx.fem.Function(self._u.function_space) # type: ignore[union-attr] + self._second_adjoint_solution_placeholder = dolfinx.fem.Function( + self._u.function_space # type: ignore[union-attr] + ) + self._hessian_u_seed = dolfinx.fem.Function(self._u.function_space) # type: ignore[union-attr] + + d2Fdu2_template = ufl.algorithms.expand_derivatives( + ufl.derivative(dFdu_template, state_placeholder, self._hessian_u_seed) + ) + soa_self_template = None + if not d2Fdu2_template.empty(): + soa_self_form = ufl.action(ufl.adjoint(d2Fdu2_template), self._adjoint_solution_placeholder) + soa_self_template = dolfinx.fem.form( + soa_self_form, + jit_options=self._jit_options, + form_compiler_options=self._form_compiler_options, + entity_maps=self._entity_maps, + ) + + dFdu_adj_applied = ufl.action(dFdu_adj_template, self._adjoint_solution_placeholder) + + _, seed_placeholders, _ = self._get_or_build_tlm_rhs_templates() + replace_map: dict = {**self._value_placeholders, self._u: state_placeholder} + F_template = ufl.replace(self._rhs, replace_map) + L1 = ufl.action(F_template, self._adjoint_solution_placeholder) + L2 = ufl.action(F_template, self._second_adjoint_solution_placeholder) + + soa_cross_templates: dict = {} + fixed_templates: dict = {} + cross_templates: dict = {} + for c, c_placeholder in self._value_placeholders.items(): + seed = seed_placeholders[c] + + soa_form = ufl.algorithms.expand_derivatives(ufl.derivative(dFdu_adj_applied, c_placeholder, seed)) + if not (soa_form == 0 or soa_form.empty()): + soa_cross_templates[c] = dolfinx.fem.form( + soa_form, + jit_options=self._jit_options, + form_compiler_options=self._form_compiler_options, + entity_maps=self._entity_maps, + ) + + dc = ufl.TestFunction(c.function_space) + dL1dm = ufl.derivative(L1, c_placeholder, dc) + dL2dm = ufl.derivative(L2, c_placeholder, dc) + d2Fdudm = ufl.algorithms.expand_derivatives( + ufl.derivative(dL1dm, state_placeholder, self._hessian_u_seed) + ) + fixed_form = ufl.algorithms.expand_derivatives(dL2dm + d2Fdudm) + if fixed_form == 0 or fixed_form.empty(): + fixed_form = ufl.ZeroBaseForm((dc,)) + fixed_templates[c] = dolfinx.fem.form( + fixed_form, + jit_options=self._jit_options, + form_compiler_options=self._form_compiler_options, + entity_maps=self._entity_maps, + ) + + for c2, c2_placeholder in self._value_placeholders.items(): + seed2 = seed_placeholders[c2] + cross_form = ufl.algorithms.expand_derivatives(ufl.derivative(dL1dm, c2_placeholder, seed2)) + if cross_form == 0 or cross_form.empty(): + continue + cross_templates[(c, c2)] = dolfinx.fem.form( + cross_form, + jit_options=self._jit_options, + form_compiler_options=self._form_compiler_options, + entity_maps=self._entity_maps, + ) + self._hessian_templates = (soa_self_template, soa_cross_templates, fixed_templates, cross_templates) + return self._hessian_templates # type: ignore[return-value] + def solve(self, annotate: bool = True) -> typing.Union[dolfinx.fem.Function, typing.Sequence[dolfinx.fem.Function]]: """ Solve the linear problem and return the solution. @@ -297,18 +964,23 @@ def solve(self, annotate: bool = True) -> typing.Union[dolfinx.fem.Function, typ bcs=self._bcs, u=self.u, # type: ignore[arg-type] P=self._preconditioner, # type: ignore[arg-type] - kind=self._kind, # type: ignore[arg-type] - petsc_options=self._petsc_options, form_compiler_options=self._form_compiler_options, jit_options=self._jit_options, entity_maps=self._entity_maps, ad_block_tag=self.ad_block_tag, - adjoint_petsc_options=self._adj_options, - tlm_petsc_options=self._tlm_options, + problem=self, ) # type: ignore[misc] tape = pyadjoint.get_working_tape() tape.add_block(block) + # Refresh the SNES-facing placeholder coefficients from the user's + # own, current values before an ordinary solve: a prior recompute + # (see NonlinearProblemBlock.prepare_recompute_component) may have + # left them holding a checkpointed/candidate value instead. + for original, placeholder in self._value_placeholders.items(): + placeholder.x.array[:] = original.x.array[:] + placeholder.x.scatter_forward() + out = dolfinx.fem.petsc.NonlinearProblem.solve(self) if annotate: if isinstance(out, Function): diff --git a/tests/test_dirichlet_bc.py b/tests/test_dirichlet_bc.py index 0687a11..336cf19 100644 --- a/tests/test_dirichlet_bc.py +++ b/tests/test_dirichlet_bc.py @@ -142,3 +142,24 @@ def test_time_dependent_bc_replay(): J_replay = Jhat(m) assert np.isclose(J_replay, J_forward, atol=1e-10, rtol=1e-10) + + # Gradient/Taylor equivalence through the multi-solve() loop: replaying, + # differentiating, and perturbing a control across several timesteps on + # one shared LinearProblem must give the same Taylor convergence as + # before the solver/form-reuse refactor (see the module-level plan in + # dolfinx-adjoint-knowledge's solver-reuse spec, "Behaviour is + # unchanged"). + # J is dominated by the (non-controlled) time-varying Dirichlet value, so + # J's sensitivity to m is comparatively small: scale the perturbation up + # so that pyadjoint.taylor_test's fixed step sizes resolve the quadratic + # remainder well above solver/roundoff noise (matching the disproportionate + # perturbation-to-control scaling already used in + # test_nonlinear_problem.py's own taylor tests). + pert = Function(V) + pert.interpolate(lambda x: 20.0 * np.cos(x[1] * np.pi)) + + Jhat(m) + pyadjoint.taylor_test(Jhat, m, pert, dJdm=0) + Jhat(m) + min_rate_grad = pyadjoint.taylor_test(Jhat, m, pert) + assert np.isclose(min_rate_grad, 2.0, rtol=1e-2, atol=5e-2), f"Expected 2.0, got {min_rate_grad}" diff --git a/tests/test_solver_reuse.py b/tests/test_solver_reuse.py new file mode 100644 index 0000000..ce91e34 --- /dev/null +++ b/tests/test_solver_reuse.py @@ -0,0 +1,649 @@ +"""Regression tests for solver/form reuse across repeated LinearProblem.solve() calls. + +These guard the "problem owns its solvers, blocks share them" refactor: recomputing +a LinearProblem's forward solve (e.g. during a tape replay, Taylor test, or +optimisation iteration) must not recompile the underlying UFL forms via FFCx on +every call, and must never mutate the user's own dependency/control objects in +place (doing so would corrupt them for any later use of that same object, such as +a Taylor test that perturbs the original control directly). +""" + +import gc +import weakref + +from mpi4py import MPI + +import dolfinx +import numpy as np +import pyadjoint +import ufl + +from dolfinx_adjoint import Function, LinearProblem, NonlinearProblem, assemble_scalar, assign + + +def _run_heat_steps(num_steps: int, monkeypatch) -> int: + """Solve a small time-stepping heat problem for ``num_steps`` steps on a single, + reused LinearProblem, then replay it once via a ReducedFunctional at a different + control value, counting calls to ``dolfinx.fem.form`` during that replay only. + """ + pyadjoint.get_working_tape().clear_tape() + + mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 6, 6) + V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) + dt = 0.1 + + m = Function(V, name="control") + m.interpolate(lambda x: np.sin(x[0] * np.pi)) + + u, v = ufl.TrialFunction(V), ufl.TestFunction(V) + uh = Function(V, name="state") + u_prev = Function(V, name="state_prev") + + F = (u - u_prev) / dt * v * ufl.dx + ufl.inner(ufl.grad(u), ufl.grad(v)) * ufl.dx - m * v * ufl.dx + a, L = ufl.system(F) + + mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim) + boundary_facets = dolfinx.mesh.exterior_facet_indices(mesh.topology) + boundary_dofs = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim - 1, boundary_facets) + bc = dolfinx.fem.dirichletbc(dolfinx.default_scalar_type(0.0), boundary_dofs, V) + + petsc_options = { + "ksp_type": "preonly", + "pc_type": "lu", + "ksp_error_if_not_converged": True, + "pc_factor_mat_solver_type": "mumps", + } + problem = LinearProblem(a, L, bcs=[bc], u=uh, petsc_options=petsc_options) + + for _ in range(num_steps): + problem.solve() + assign(uh, u_prev) + + # A single objective computed once, after the loop: its own AssembleBlock + # recompiles independently of this test and independently of num_steps, so + # it contributes a fixed, step-count-independent cost that does not + # confound the assertion this test is making about LinearProblemBlock. + J = assemble_scalar(0.5 * ufl.inner(uh, uh) * ufl.dx) + control = pyadjoint.Control(m) + Jhat = pyadjoint.ReducedFunctional(J, control) + + m2 = Function(V) + m2.interpolate(lambda x: 1.0 + np.cos(x[0] * np.pi)) + + real_form = dolfinx.fem.form + calls = [] + + def counting_form(*args, **kwargs): + calls.append(1) + return real_form(*args, **kwargs) + + with monkeypatch.context() as ctx: + ctx.setattr(dolfinx.fem, "form", counting_form) + Jhat(m2) + + return len(calls) + + +def test_recompute_form_count_independent_of_step_count(monkeypatch): + """Replaying more timesteps on the same LinearProblem must not compile more forms. + + The recompute-time placeholder forms (see + ``LinearProblem._value_placeholders``) are compiled exactly once, when the + Problem is constructed -- not once per block/timestep, and not again on + replay. So the number of ``dolfinx.fem.form`` calls triggered by replaying + the whole tape must be the same whether the tape has 3 timesteps or 8. + """ + count_short = _run_heat_steps(3, monkeypatch) + count_long = _run_heat_steps(8, monkeypatch) + assert count_short > 0 + assert count_short == count_long + + +def test_recompute_does_not_corrupt_original_control(): + """A Taylor test that perturbs the original control object directly must see + that object's pristine value at every perturbation, not whatever a previous + recompute happened to leave written into it. + + This is a direct regression test for a bug caught while implementing + recompute-time placeholders: mutating a block's dependency Function in place + during recompute (rather than into a dedicated placeholder) corrupts the + control for any later computation relative to it -- exactly what + ``pyadjoint.taylor_test(Jh, m, dm)`` does by calling ``m._ad_add(...)`` on the + same, live ``m`` object between evaluations. + """ + pyadjoint.get_working_tape().clear_tape() + + mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 6, 6) + V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) + + uh = Function(V, name="state") + v = ufl.TestFunction(V) + u_trial = ufl.TrialFunction(V) + + m = Function(V, name="control") + m.interpolate(lambda x: 1.0 + x[0] ** 2 + x[1] ** 2) + m_original = m.x.array.copy() + + f = dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(1.0)) + a = m * ufl.inner(ufl.grad(u_trial), ufl.grad(v)) * ufl.dx + L = f * v * ufl.dx + + mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim) + boundary_facets = dolfinx.mesh.exterior_facet_indices(mesh.topology) + boundary_dofs = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim - 1, boundary_facets) + bc = dolfinx.fem.dirichletbc(dolfinx.default_scalar_type(0.0), boundary_dofs, V) + + petsc_options = { + "ksp_type": "preonly", + "pc_type": "lu", + "ksp_error_if_not_converged": True, + "pc_factor_mat_solver_type": "mumps", + } + problem = LinearProblem(a, L, bcs=[bc], u=uh, petsc_options=petsc_options) + problem.solve() + + d = Function(V) + d.interpolate(lambda x: np.sin(np.pi * x[0])) + J = assemble_scalar(0.5 * ufl.inner(uh - d, uh - d) * ufl.dx) + + control = pyadjoint.Control(m) + Jh = pyadjoint.ReducedFunctional(J, control) + + dm = Function(V) + dm.interpolate(lambda x: 2 * np.sin(x[0] * np.pi) * np.cos(x[1] * np.pi)) + + # Passing `m` itself (not a copy) as the expansion point, matching a + # perfectly ordinary way to write this call. + min_rate = pyadjoint.taylor_test(Jh, m, dm, dJdm=0) + + assert np.allclose(m.x.array, m_original), ( + "the original control object was mutated by recompute; it must stay pristine" + ) + assert np.isclose(min_rate, 1.0, rtol=1e-1, atol=1e-1), f"Expected convergence rate close to 1.0, got {min_rate}" + + +def test_nonlinear_recompute_does_not_corrupt_original_control(): + """As above, but for NonlinearProblem. + + NonlinearProblem's SNES binds its residual/Jacobian callbacks to fixed + compiled Form objects at construction time, so there is no way to later + swap in a differently-compiled form the way a KSP-based solve can; the + only way to make the SNES see a different coefficient value is to mutate + the exact object its compiled forms reference. Every non-``u`` coefficient + is therefore routed through a dedicated placeholder from the moment the + SNES is built (``NonlinearProblem._value_placeholders``), and + ``LinearProblem`` now uses the identical mechanism -- both classes always + solve by refreshing placeholder values and calling an unchanging, + already-compiled solver, never by mutating the user's own coefficient or + recompiling a form. + """ + pyadjoint.get_working_tape().clear_tape() + + mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 8, 7) + V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) + + f = Function(V, name="control") + f.interpolate(lambda x: 2.0 + np.sin(x[0])) + f_original = f.x.array.copy() + + u1 = Function(V, name="state") + u1.interpolate(lambda x: np.ones_like(x[0])) + v1 = ufl.TestFunction(V) + F1 = (1 + u1**2) * ufl.inner(ufl.grad(u1), ufl.grad(v1)) * ufl.dx - f * v1 * ufl.dx + + mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim) + boundary_facets = dolfinx.mesh.exterior_facet_indices(mesh.topology) + boundary_dofs = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim - 1, boundary_facets) + bc_val = dolfinx.fem.Constant(mesh, np.dtype(dolfinx.default_scalar_type).type(1.0)) + bc = dolfinx.fem.dirichletbc(bc_val, boundary_dofs, V) + + options = { + "snes_error_if_not_converged": True, + "ksp_type": "preonly", + "pc_type": "lu", + "pc_factor_mat_solver_type": "mumps", + } + problem = NonlinearProblem(F1, u=u1, bcs=[bc], petsc_options=options, adjoint_petsc_options=options) + problem.solve() + + d = pyadjoint.AdjFloat(0.2) + J = assemble_scalar((u1 - d) ** 3 * ufl.dx) + control = pyadjoint.Control(f) + Jh = pyadjoint.ReducedFunctional(J, control) + + dm = Function(V) + dm.interpolate(lambda x: 2 * np.sin(x[0] * np.pi)) + + # Passing `f` itself (not a copy) as the expansion point. + min_rate = pyadjoint.taylor_test(Jh, f, dm, dJdm=0) + + assert np.allclose(f.x.array, f_original), ( + "the original control object was mutated by recompute; it must stay pristine" + ) + assert np.isclose(min_rate, 1.0, rtol=1e-1, atol=1e-1), f"Expected convergence rate close to 1.0, got {min_rate}" + + +def test_linear_adjoint_lhs_compiled_once(): + """adjoint(dF/du) does not depend on the state for a linear problem, so it + should be compiled exactly once per Problem (in + ``LinearProblem._get_or_build_adjoint_solver``) and never rebuilt -- + including across repeated ``derivative()``/``hessian()`` calls at + different control values, which is exactly the scenario + (``test_tlm_update.py``'s "warm up at another point" tests) that would + catch a stale operator being silently reused with the wrong coefficient + values. + + The control ``m`` sits inside the bilinear form itself (not just the + right-hand side), so the adjoint operator actually depends on it -- a + control that only appeared in ``L`` wouldn't exercise this at all. + """ + pyadjoint.get_working_tape().clear_tape() + + mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 6, 6) + V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) + + uh = Function(V, name="state") + v = ufl.TestFunction(V) + u_trial = ufl.TrialFunction(V) + + m = Function(V, name="control") + m.interpolate(lambda x: 1.0 + x[0] ** 2 + x[1] ** 2) + + f = dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(1.0)) + a = m * ufl.inner(ufl.grad(u_trial), ufl.grad(v)) * ufl.dx + L = f * v * ufl.dx + + mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim) + boundary_facets = dolfinx.mesh.exterior_facet_indices(mesh.topology) + boundary_dofs = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim - 1, boundary_facets) + bc = dolfinx.fem.dirichletbc(dolfinx.default_scalar_type(0.0), boundary_dofs, V) + + petsc_options = { + "ksp_type": "preonly", + "pc_type": "lu", + "ksp_error_if_not_converged": True, + "pc_factor_mat_solver_type": "mumps", + } + problem = LinearProblem(a, L, bcs=[bc], u=uh, petsc_options=petsc_options) + problem.solve() + + d = Function(V) + d.interpolate(lambda x: np.sin(np.pi * x[0])) + J = assemble_scalar(0.5 * ufl.inner(uh - d, uh - d) * ufl.dx) + control = pyadjoint.Control(m) + Jh = pyadjoint.ReducedFunctional(J, control) + + adjoint_solver = problem._get_or_build_adjoint_solver() + compiled_lhs = adjoint_solver._a + assert compiled_lhs is not None + + Jh.derivative() + assert adjoint_solver._a is compiled_lhs, "adjoint LHS was rebuilt by derivative()" + + m2 = Function(V) + m2.interpolate(lambda x: 2.0 + np.sin(x[0])) + Jh(m2) + Jh.derivative() + assert adjoint_solver._a is compiled_lhs, "adjoint LHS was rebuilt after evaluating at a new point" + + dm = Function(V) + dm.interpolate(lambda x: np.cos(x[0] * np.pi)) + Jh.hessian(dm) + assert adjoint_solver._a is compiled_lhs, "the SOA (Hessian) solve rebuilt the shared adjoint LHS" + + +def test_nonlinear_adjoint_lhs_compiled_once(): + """As above, but for NonlinearProblem: adjoint(dF/du) does depend on u's + current value here, but that dependency is routed through a dedicated + placeholder (``NonlinearProblem._state_placeholder``), refreshed per + call, so the compiled LHS itself is still built exactly once. + """ + pyadjoint.get_working_tape().clear_tape() + + mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 6, 6) + V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) + + f = Function(V, name="control") + f.interpolate(lambda x: 2.0 + np.sin(x[0])) + + u1 = Function(V, name="state") + u1.interpolate(lambda x: np.ones_like(x[0])) + v1 = ufl.TestFunction(V) + F1 = (1 + u1**2) * ufl.inner(ufl.grad(u1), ufl.grad(v1)) * ufl.dx - f * v1 * ufl.dx + + mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim) + boundary_facets = dolfinx.mesh.exterior_facet_indices(mesh.topology) + boundary_dofs = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim - 1, boundary_facets) + bc_val = dolfinx.fem.Constant(mesh, np.dtype(dolfinx.default_scalar_type).type(1.0)) + bc = dolfinx.fem.dirichletbc(bc_val, boundary_dofs, V) + + options = { + "snes_error_if_not_converged": True, + "ksp_type": "preonly", + "pc_type": "lu", + "pc_factor_mat_solver_type": "mumps", + } + problem = NonlinearProblem(F1, u=u1, bcs=[bc], petsc_options=options, adjoint_petsc_options=options) + problem.solve() + + d = pyadjoint.AdjFloat(0.2) + J = assemble_scalar((u1 - d) ** 3 * ufl.dx) + control = pyadjoint.Control(f) + Jh = pyadjoint.ReducedFunctional(J, control) + + adjoint_solver = problem._get_or_build_adjoint_solver() + compiled_lhs = adjoint_solver._a + assert compiled_lhs is not None + + Jh.derivative() + assert adjoint_solver._a is compiled_lhs, "adjoint LHS was rebuilt by derivative()" + + f2 = Function(V) + f2.interpolate(lambda x: 3.0 + np.cos(x[0])) + Jh(f2) + Jh.derivative() + assert adjoint_solver._a is compiled_lhs, "adjoint LHS was rebuilt after evaluating at a new point" + + dm = Function(V) + dm.interpolate(lambda x: np.sin(x[0] * np.pi)) + Jh.hessian(dm) + assert adjoint_solver._a is compiled_lhs, "the SOA (Hessian) solve rebuilt the shared adjoint LHS" + + +def test_tlm_rhs_templates_compiled_once(): + """Each dependency's TLM right-hand-side template + (``LinearProblem._get_or_build_tlm_rhs_templates``) must be compiled exactly + once and never rebuilt, including across repeated ``hessian()`` calls at + different control values (a Hessian evaluation always drives a TLM sweep + first). + """ + pyadjoint.get_working_tape().clear_tape() + + mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 6, 6) + V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) + + uh = Function(V, name="state") + v = ufl.TestFunction(V) + u_trial = ufl.TrialFunction(V) + + m = Function(V, name="control") + m.interpolate(lambda x: 1.0 + x[0] ** 2 + x[1] ** 2) + + f = dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(1.0)) + a = m * ufl.inner(ufl.grad(u_trial), ufl.grad(v)) * ufl.dx + L = f * v * ufl.dx + + mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim) + boundary_facets = dolfinx.mesh.exterior_facet_indices(mesh.topology) + boundary_dofs = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim - 1, boundary_facets) + bc = dolfinx.fem.dirichletbc(dolfinx.default_scalar_type(0.0), boundary_dofs, V) + + petsc_options = { + "ksp_type": "preonly", + "pc_type": "lu", + "ksp_error_if_not_converged": True, + "pc_factor_mat_solver_type": "mumps", + } + problem = LinearProblem(a, L, bcs=[bc], u=uh, petsc_options=petsc_options) + problem.solve() + + d = Function(V) + d.interpolate(lambda x: np.sin(np.pi * x[0])) + J = assemble_scalar(0.5 * ufl.inner(uh - d, uh - d) * ufl.dx) + control = pyadjoint.Control(m) + Jh = pyadjoint.ReducedFunctional(J, control) + + dm = Function(V) + dm.interpolate(lambda x: np.cos(x[0] * np.pi)) + Jh.hessian(dm) + + templates, _, _ = problem._get_or_build_tlm_rhs_templates() + compiled_ids = {c: id(form) for c, form in templates.items()} + assert compiled_ids, "no TLM RHS templates were built" + + m2 = Function(V) + m2.interpolate(lambda x: 2.0 + np.sin(x[0])) + Jh(m2) + Jh.hessian(dm) + + templates_after, _, _ = problem._get_or_build_tlm_rhs_templates() + for c, form in templates_after.items(): + assert id(form) == compiled_ids[c], f"TLM RHS template for {c.name} was rebuilt" + + +def test_tlm_skips_inactive_dependency_with_singular_derivative(): + """An inactive dependency (no tangent-linear value) whose derivative would be + singular at its current value must never be evaluated at all -- not even + with a zeroed seed. + + ``c`` enters the bilinear form as ``sqrt(c)``, which is finite at ``c == 0`` + but whose derivative, ``1 / (2 * sqrt(c))``, is not. ``c`` is set to exactly + zero over part of the domain and is never declared a control, so its + tangent-linear value is always ``None``. If that dependency's contribution + were assembled with a zeroed seed instead of skipped outright, the ``0 * + inf`` there would poison the whole tangent-linear (and, downstream, + Hessian) result with ``NaN``. Only ``m`` is seeded. + """ + pyadjoint.get_working_tape().clear_tape() + + mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 6, 6) + V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) + Z = dolfinx.fem.functionspace(mesh, ("DG", 0)) + + uh = Function(V, name="state") + v = ufl.TestFunction(V) + u_trial = ufl.TrialFunction(V) + + c = Function(Z, name="not_a_control") + c.interpolate(lambda x: np.maximum(x[0] - 0.5, 0.0)) + + m = Function(V, name="control") + m.interpolate(lambda x: 1.0 + x[0] ** 2 + x[1] ** 2) + + a = (1.0 + ufl.sqrt(c)) * ufl.inner(ufl.grad(u_trial), ufl.grad(v)) * ufl.dx + L = m * v * ufl.dx + + mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim) + boundary_facets = dolfinx.mesh.exterior_facet_indices(mesh.topology) + boundary_dofs = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim - 1, boundary_facets) + bc = dolfinx.fem.dirichletbc(dolfinx.default_scalar_type(0.0), boundary_dofs, V) + + petsc_options = { + "ksp_type": "preonly", + "pc_type": "lu", + "ksp_error_if_not_converged": True, + "pc_factor_mat_solver_type": "mumps", + } + problem = LinearProblem(a, L, bcs=[bc], u=uh, petsc_options=petsc_options) + problem.solve() + assert np.isfinite(uh.x.array).all() + + d = Function(V) + d.interpolate(lambda x: np.sin(np.pi * x[0])) + J = assemble_scalar(0.5 * ufl.inner(uh - d, uh - d) * ufl.dx) + control = pyadjoint.Control(m) + Jh = pyadjoint.ReducedFunctional(J, control) + + dm = Function(V) + dm.interpolate(lambda x: np.cos(x[0] * np.pi)) + + grad = Jh.derivative() + assert np.isfinite(grad.x.array).all(), "gradient contains NaN/Inf from the inactive singular dependency" + + hessian_action = Jh.hessian(dm) + assert np.isfinite(hessian_action.x.array).all(), ( + "Hessian action contains NaN/Inf from the inactive singular dependency" + ) + + +def test_nonlinear_tlm_rhs_templates_compiled_once(): + """As ``test_tlm_rhs_templates_compiled_once``, but for NonlinearProblem: + the TLM solver (``NonlinearProblem._get_or_build_tlm_solver``) and each + dependency's TLM right-hand-side template + (``NonlinearProblem._get_or_build_tlm_rhs_templates``) must be compiled + exactly once and never rebuilt, including across repeated ``hessian()`` + calls at different control values (a Hessian evaluation always drives a + TLM sweep first). + """ + pyadjoint.get_working_tape().clear_tape() + + mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 6, 6) + V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) + + f = Function(V, name="control") + f.interpolate(lambda x: 2.0 + np.sin(x[0])) + + u1 = Function(V, name="state") + u1.interpolate(lambda x: np.ones_like(x[0])) + v1 = ufl.TestFunction(V) + F1 = (1 + u1**2) * ufl.inner(ufl.grad(u1), ufl.grad(v1)) * ufl.dx - f * v1 * ufl.dx + + mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim) + boundary_facets = dolfinx.mesh.exterior_facet_indices(mesh.topology) + boundary_dofs = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim - 1, boundary_facets) + bc_val = dolfinx.fem.Constant(mesh, np.dtype(dolfinx.default_scalar_type).type(1.0)) + bc = dolfinx.fem.dirichletbc(bc_val, boundary_dofs, V) + + options = { + "snes_error_if_not_converged": True, + "ksp_type": "preonly", + "pc_type": "lu", + "pc_factor_mat_solver_type": "mumps", + } + problem = NonlinearProblem(F1, u=u1, bcs=[bc], petsc_options=options, adjoint_petsc_options=options) + problem.solve() + + d = pyadjoint.AdjFloat(0.2) + J = assemble_scalar((u1 - d) ** 3 * ufl.dx) + control = pyadjoint.Control(f) + Jh = pyadjoint.ReducedFunctional(J, control) + + dm = Function(V) + dm.interpolate(lambda x: np.sin(x[0] * np.pi)) + Jh.hessian(dm) + + tlm_solver = problem._get_or_build_tlm_solver() + compiled_tlm_lhs = tlm_solver._a + assert compiled_tlm_lhs is not None + + templates, _, _ = problem._get_or_build_tlm_rhs_templates() + compiled_ids = {c: id(form) for c, form in templates.items()} + assert compiled_ids, "no TLM RHS templates were built" + + f2 = Function(V) + f2.interpolate(lambda x: 3.0 + np.cos(x[0])) + Jh(f2) + Jh.hessian(dm) + + assert tlm_solver._a is compiled_tlm_lhs, "TLM LHS was rebuilt after evaluating at a new point" + templates_after, _, _ = problem._get_or_build_tlm_rhs_templates() + for c, form in templates_after.items(): + assert id(form) == compiled_ids[c], f"TLM RHS template for {c.name} was rebuilt" + + +def test_linear_problem_released_by_refcounting_not_gc(): + """Dropping a LinearProblem (and clearing the tape) must release it -- and its + PETSc solvers -- via ordinary reference counting, never leaving it as cyclic + garbage collected only by a later ``gc.collect()``. + + Regression test for a real bug found while building the placeholder-substitution + machinery: a recursive nested ``_replace`` helper inside ``LinearProblem.__init__`` + captured *itself* (for the recursive call) and ``self`` in its closure, forming a + reference cycle. That cycle kept the Problem -- and hence its PETSc Mat/KSP + objects -- alive until whenever the cyclic garbage collector next ran, exactly + the "freed at a rank-nondeterministic moment" hazard the whole + Problem-owns-its-solvers refactor exists to eliminate: confirmed live under + ``mpirun -n 2`` by observing the two ranks diverge into two different collective + calls -- one inside a PETSc Mat's collective MUMPS-termination destructor, the + other already building an unrelated dofmap for the next test -- while this bug + was present. Fixed by making the recursive helper + (``dolfinx_adjoint.solvers._replace_with_placeholders``) a plain module-level + function taking the placeholder dict as an explicit argument, so it does not need + to capture itself or ``self``. + """ + pyadjoint.get_working_tape().clear_tape() + + mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 6, 6) + V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) + + uh = Function(V, name="state") + v = ufl.TestFunction(V) + u_trial = ufl.TrialFunction(V) + m = Function(V, name="control") + m.interpolate(lambda x: 1.0 + x[0] ** 2) + + a = m * ufl.inner(ufl.grad(u_trial), ufl.grad(v)) * ufl.dx + L = ufl.inner(dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(1.0)), v) * ufl.dx + + mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim) + boundary_facets = dolfinx.mesh.exterior_facet_indices(mesh.topology) + boundary_dofs = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim - 1, boundary_facets) + bc = dolfinx.fem.dirichletbc(dolfinx.default_scalar_type(0.0), boundary_dofs, V) + + petsc_options = { + "ksp_type": "preonly", + "pc_type": "lu", + "ksp_error_if_not_converged": True, + "pc_factor_mat_solver_type": "mumps", + } + gc.disable() + try: + problem = LinearProblem(a, L, bcs=[bc], u=uh, petsc_options=petsc_options) + problem.solve() + + problem_ref = weakref.ref(problem) + del problem + pyadjoint.get_working_tape().clear_tape() + + assert problem_ref() is None, ( + "LinearProblem was not released by ordinary refcounting -- it is cyclic " + "garbage, collected only by gc.collect() at a moment that can differ " + "between MPI ranks" + ) + finally: + gc.enable() + + +def test_nonlinear_problem_released_by_refcounting_not_gc(): + """As above, but for NonlinearProblem.""" + pyadjoint.get_working_tape().clear_tape() + + mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 6, 6) + V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) + + f = Function(V, name="control") + f.interpolate(lambda x: 2.0 + np.sin(x[0])) + + u1 = Function(V, name="state") + u1.interpolate(lambda x: np.ones_like(x[0])) + v1 = ufl.TestFunction(V) + F1 = (1 + u1**2) * ufl.inner(ufl.grad(u1), ufl.grad(v1)) * ufl.dx - f * v1 * ufl.dx + + mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim) + boundary_facets = dolfinx.mesh.exterior_facet_indices(mesh.topology) + boundary_dofs = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim - 1, boundary_facets) + bc_val = dolfinx.fem.Constant(mesh, np.dtype(dolfinx.default_scalar_type).type(1.0)) + bc = dolfinx.fem.dirichletbc(bc_val, boundary_dofs, V) + + options = { + "snes_error_if_not_converged": True, + "ksp_type": "preonly", + "pc_type": "lu", + "pc_factor_mat_solver_type": "mumps", + } + gc.disable() + try: + problem = NonlinearProblem(F1, u=u1, bcs=[bc], petsc_options=options, adjoint_petsc_options=options) + problem.solve() + + problem_ref = weakref.ref(problem) + del problem + pyadjoint.get_working_tape().clear_tape() + + assert problem_ref() is None, ( + "NonlinearProblem was not released by ordinary refcounting -- it is cyclic " + "garbage, collected only by gc.collect() at a moment that can differ " + "between MPI ranks" + ) + finally: + gc.enable() diff --git a/tests/test_tlm_update.py b/tests/test_tlm_update.py index 9577309..45ee2cb 100644 --- a/tests/test_tlm_update.py +++ b/tests/test_tlm_update.py @@ -121,6 +121,82 @@ def test_hessian_is_independent_of_previous_evaluation_points(warm_up_at_another assert np.isclose(min_rate, 3.0, rtol=0.1, atol=0.1), f"Expected convergence rate close to 3.0, got {min_rate}" +def _diffusive_poisson(mesh): + """Scalar Poisson problem whose control ``m`` sits inside ``a`` (the diffusivity). + + A scalar-space sibling of ``_viscous_stokes``: the control has to enter the + bilinear form for this test to have any teeth, for the same reason noted there. + """ + V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) + Z = dolfinx.fem.functionspace(mesh, ("DG", 0)) + dx = ufl.Measure("dx", domain=mesh) + + m = Function(Z, name="diffusivity") + m.interpolate(lambda x: 1.0 + 0.5 * np.sin(np.pi * x[0])) + + u, v = ufl.TrialFunction(V), ufl.TestFunction(V) + x = ufl.SpatialCoordinate(mesh) + f = 1e2 * ufl.sin(ufl.pi * x[0]) * ufl.cos(ufl.pi * x[1]) + + a = ufl.inner(m * ufl.grad(u), ufl.grad(v)) * dx + L = ufl.inner(f, v) * dx + + mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim) + facets = dolfinx.mesh.exterior_facet_indices(mesh.topology) + dofs = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim - 1, facets) + bc = dolfinx.fem.dirichletbc(dolfinx.default_scalar_type(0.0), dofs, V) + + uh = Function(V, name="state") + problem = LinearProblem( + a, + L, + u=uh, + bcs=[bc], + petsc_options=direct_solve, + adjoint_petsc_options=direct_solve, + tlm_petsc_options=direct_solve, + ) + problem.solve() + + # Quartic in the state, for the same round-off-avoidance reason as + # ``_viscous_stokes``'s objective. + J = assemble_scalar(uh**4 * dx) + return pyadjoint.ReducedFunctional(J, pyadjoint.Control(m)), Z + + +@pytest.mark.parametrize("warm_up_at_another_point", [False, True]) +def test_hessian_is_independent_of_previous_evaluation_points_scalar(warm_up_at_another_point, mesh_2D): + """Scalar-space sibling of ``test_hessian_is_independent_of_previous_evaluation_points``. + + That test only covers the blocked/Stokes path; a plain (non-blocked) ``LinearProblem`` + with the control inside the bilinear form is the common case the cached, compiled-once + adjoint/TLM/Hessian templates (``LinearProblem._get_or_build_adjoint_solver``/ + ``_get_or_build_tlm_solver``/``_get_or_build_hessian_templates``) exist to serve, and + deserves the identical regression coverage: the Hessian at ``m2`` must not depend on + whether ``J`` was evaluated at ``m1`` first. + """ + pyadjoint.get_working_tape().clear_tape() + Jh, Z = _diffusive_poisson(mesh_2D) + + m1, m2 = Function(Z), Function(Z) + m1.interpolate(lambda x: 1.0 + 0.5 * np.sin(np.pi * x[0])) + m2.interpolate(lambda x: 2.0 + 0.5 * np.cos(np.pi * x[1])) + h = Function(Z) + h.interpolate(lambda x: 1.0 + 0.3 * np.sin(3 * x[0])) + + if warm_up_at_another_point: + Jh(m1) + Jh.derivative() + Jh.hessian(h) + + Jh(m2) + dJdm = Jh.derivative()._ad_dot(h) + Hm = Jh.hessian(h)._ad_dot(h) + + min_rate = pyadjoint.taylor_test(Jh, m2, h, dJdm=dJdm, Hm=Hm) + assert np.isclose(min_rate, 3.0, rtol=0.1, atol=0.1), f"Expected convergence rate close to 3.0, got {min_rate}" + + def test_hessian_mpi_breakdown(mesh_2D): pyadjoint.get_working_tape().clear_tape() Jh, Z = _viscous_stokes(mesh_2D)