From c8f7432d612706480bd59ccf2fe58cfe9059ff03 Mon Sep 17 00:00:00 2001 From: jorgensd Date: Mon, 31 Aug 2026 20:51:09 +0000 Subject: [PATCH 1/8] Run new claude session on the code. Plan to verify tomorrow. --- src/dolfinx_adjoint/assembly.py | 8 +- src/dolfinx_adjoint/blocks/solvers.py | 1391 ++++++++++--------------- src/dolfinx_adjoint/petsc_utils.py | 26 +- src/dolfinx_adjoint/solvers.py | 1239 +++++++++++----------- tests/test_tlm_update.py | 106 +- 5 files changed, 1256 insertions(+), 1514 deletions(-) diff --git a/src/dolfinx_adjoint/assembly.py b/src/dolfinx_adjoint/assembly.py index a489d94..c55b341 100644 --- a/src/dolfinx_adjoint/assembly.py +++ b/src/dolfinx_adjoint/assembly.py @@ -57,10 +57,10 @@ def error_norm( u_ex: ufl.core.expr.Expr, u: ufl.core.expr.Expr, norm_type=typing.Literal["L2", "H1"], - jit_options: typing.Optional[dict] = None, - form_compiler_options: typing.Optional[dict] = None, - entity_map: typing.Optional[dict[dolfinx.mesh.Mesh, npt.NDArray[numpy.int32]]] = None, - ad_block_tag: typing.Optional[str] = None, + jit_options: dict | None = None, + form_compiler_options: dict | None = None, + entity_map: dict[dolfinx.mesh.Mesh, npt.NDArray[numpy.int32]] | None = None, + ad_block_tag: str | None = None, annotate: bool = True, ) -> float: """Compute the error norm between the exact solution and the computed solution. diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index d074f37..4fba1ba 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -132,6 +132,52 @@ def get_sorted_arguments(arguments: typing.Iterable[ufl.Argument], number: int) return sorted(filter(lambda x: x.number() == number, arguments), key=lambda a: a.part()) +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 a blocked ``NonlinearProblem``'s forms are not + required to be before ``assign_mixed_parts`` runs. + """ + 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 _map_block_variables_to_form( + form: ufl.Form | NestedMutableSequence[ufl.Form] | None, + block_variables: typing.Iterable[pyadjoint.block_variable.BlockVariable], +) -> dict[Function, Function]: + """Map each ``block_variable``'s output coefficient, where it appears in ``form``, to its + checkpointed value. + + ``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); recurses once per nesting level, + independently of how many ``block_variables`` are given. + """ + if form is None: + return {} + if isinstance(form, ufl.Form): + coefficients = form.coefficients() + return { + block_variable.output: block_variable.saved_output + for block_variable in block_variables + if block_variable.output in coefficients + } + replace_map: dict = {} + for f in form: + replace_map.update(_map_block_variables_to_form(f, block_variables)) + return replace_map + + def sum_form(form: NestedSequence[ufl.Form | None]) -> ufl.Form | None: """Sum a blocked form into a single form.""" # Handle top-level None @@ -160,165 +206,61 @@ def sum_form(form: NestedSequence[ufl.Form | None]) -> ufl.Form | None: raise TypeError(f"Cannot sum form of type {type(form)}") -class LinearProblemBlock(pyadjoint.Block): - """A linear problem that can be used with adjoint methods. - - This class extends the `dolfinx.fem.petsc.LinearProblem` to support adjoint methods. +class _ProblemBlockBase(pyadjoint.Block): + """Shared tape-block machinery for ``LinearProblemBlock``/``NonlinearProblemBlock``. + + Holds what is unconditionally identical or shares one implementation + between the two Problem kinds: fetching the owning Problem, recovering + Dirichlet BC dependencies, detecting a boundary-condition-only adjoint, + transposing a (possibly blocked) bilinear form, the shared warm-started + recompute flow, and -- since ``LinearProblem._compute_residual`` and + ``NonlinearProblem._compute_residual`` both settle on the same output + shape (a single summed ``ufl.Form`` plus its dependency replacement map, + see each subclass's docstring) -- every first-order adjoint, TLM and + Hessian method built *on top of* that residual, both scalar and blocked. + Each concrete subclass still implements its own ``__init__`` (constructor + kwargs differ: ``a``/``L`` vs ``J``/``F``) and ``_compute_residual`` + itself, since building the residual is the one place a shared algorithm + isn't possible -- the two classes start from different user-supplied + data. """ + _problem_obj: typing.Any + _bcs: typing.Sequence[dolfinx.fem.DirichletBC] + _u: _Function | typing.Sequence[_Function] _adjoint_solutions: dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function] - _tlm_solutions: dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function] _second_adjoint_solutions: dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function] + _tlm_solutions: dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function] + _jit_options: dict | None + _form_compiler_options: dict | None + _entity_maps: typing.Sequence[dolfinx.mesh.EntityMap] | None - # 2. Overload for the SCALAR case - @typing.overload - def __init__( - self, - a: ufl.Form, - L: ufl.Form, - *, - bcs: typing.Sequence[dolfinx.fem.DirichletBC] | None = None, - u: _Function | None = None, - P: ufl.Form | 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, - problem: "LinearProblem" = ..., - ) -> None: ... - - @typing.overload - def __init__( - self, - a: typing.Sequence[typing.Sequence[ufl.Form]], - L: typing.Sequence[ufl.Form], - *, - bcs: typing.Sequence[dolfinx.fem.DirichletBC] | None = None, - u: typing.Sequence[_Function] | None = None, - P: typing.Sequence[typing.Sequence[ufl.Form]] | 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, - problem: "LinearProblem" = ..., - ) -> None: ... - - def __init__( - self, - a: ufl.Form | typing.Sequence[typing.Sequence[ufl.Form]], - L: ufl.Form | typing.Sequence[ufl.Form], - *, - 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, - 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, - problem: "LinearProblem" = None, # type: ignore[assignment] - ) -> None: - - assert problem is not None, "problem must be provided." - # Strong reference, deliberately: a throwaway LinearProblem (solve it, - # then only touch the ReducedFunctional) is a common pattern, so the - # block must keep the Problem -- and its shared solvers -- alive for - # as long as the block itself is reachable. Not cyclic garbage on its - # own (verified), so this doesn't reintroduce the MPI - # collective-destruction hazard from dolfinx-adjoint-knowledge's - # mpi-collective-destruction-hazard note -- that needs an unmerged - # checkpoint schedule making the *tape* cyclic. - self._problem_obj = problem - super().__init__(ad_block_tag=ad_block_tag) - - # Collect all arguments in variational forms and replace them with similar - # once that is based on a mixed functionspace. - if not isinstance(a, ufl.Form): - a, L = assign_mixed_parts(a, L) - if P is not None: - P, _ = assign_mixed_parts(P, L) - - self._lhs = a - self._rhs = L - self._preconditioner = P - - # Create overloaded functions - self._u: _Function | typing.Sequence[_Function] - if isinstance(u, dolfinx.fem.Function): - self._u = pyadjoint.create_overloaded_object(u) - elif u is None: - try: - # Extract function space for unknown from the right hand - # side of the equation. - self._u = Function(L.arguments()[0].ufl_function_space()) # type: ignore - except AttributeError: - self._u = [Function(Li.arguments()[0].ufl_function_space()) for Li in L] # type: ignore[union-attr] - else: - self._u = [pyadjoint.create_overloaded_object(ui) for ui in u] - - # NOTE: Add mesh and constants as dependencies later on - - # To ensure that the solver can be recycled in time dependent loops, the unknown is also added as a dependency - # if present in the form. - if isinstance(self._u, dolfinx.fem.Function): - assert isinstance(self._lhs, ufl.Form) - assert isinstance(self._rhs, ufl.Form) - if self._u in self._lhs.coefficients() or self._u in self._rhs.coefficients(): - raise RuntimeError("The unknown function u should not be present in the variational forms a or L.") - for c in self._lhs.coefficients(): - self.add_dependency(c, no_duplicates=True) - for c in self._rhs.coefficients(): - self.add_dependency(c, no_duplicates=True) - elif isinstance(self._u, typing.Iterable): - for Ai in self._lhs: # type: ignore - for Aij in Ai: - if Aij is not None: - assert isinstance(Aij, ufl.Form) - for c in Aij.coefficients(): - if c in self._u: - raise RuntimeError( - "The unknown function u should not be present in the variational forms a or L." - ) - self.add_dependency(c, no_duplicates=True) - for part in self._rhs: # type: ignore - for c in part.coefficients(): - if c in self._u: - raise RuntimeError( - "The unknown function u should not be present in the variational forms a or L." - ) - self.add_dependency(c, no_duplicates=True) - else: - raise RuntimeError(f"Unknown type for unknown function u={type(self._u)}.") - # 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._bcs = bcs if bcs is not None else [] + @property + def _problem(self) -> typing.Any: + """Return this block's owning Problem, which owns the shared solvers.""" + return self._problem_obj - # Add dependencies from the boundary conditions - if self._bcs is not None: - for bc in self._bcs: - if hasattr(bc, "block_variable"): - self.add_dependency(bc, no_duplicates=True) + def _compute_residual(self) -> tuple[ufl.Form, dict[Function, Function]]: + """Build this block's residual ``F(u, v) = 0`` at its checkpointed dependency values. - # 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. + The one genuinely irreducible difference between the two Problem kinds: ``LinearProblem`` + derives it from ``a``/``L`` via ``ufl.action``, ``NonlinearProblem`` already has ``F`` + directly. Both settle on the same output shape -- a single summed ``ufl.Form`` plus its + dependency replacement map -- which is what lets every method built on top of this one + (below) be shared. Not implemented on the base; each subclass overrides it. + """ + raise NotImplementedError - if isinstance(self._u, dolfinx.fem.Function): - self._adjoint_solutions = self._u.copy() - self._second_adjoint_solutions = self._u.copy() - self._tlm_solutions = self._u.copy() - else: - assert isinstance(self._u, typing.Iterable) - self._adjoint_solutions = [u.copy() for u in self._u] - self._second_adjoint_solutions = [u.copy() for u in self._u] - self._tlm_solutions = [u.copy() for u in self._u] + def _refresh_dFdu_state(self, problem: typing.Any) -> None: + """Refresh whichever coefficient stands in for "the state" in ``dF/du``, if any. - def _problem(self) -> "LinearProblem": - """Return this block's owning Problem, which owns the shared solvers.""" - return self._problem_obj + A no-op by default: ``LinearProblem``'s ``dF/du`` (``a``) is bilinear and never + references the state at all, so there is nothing to refresh before using the shared + adjoint solver. ``NonlinearProblemBlock`` overrides this, since ``dF/du`` genuinely + depends on ``u``'s current value there -- and, unlike the TLM path (which already loops + over ``problem._get_or_build_tlm_rhs_templates()``'s state placeholder(s) generically), + neither ``prepare_evaluate_adj`` nor ``prepare_evaluate_hessian`` otherwise touches it. + """ def _recover_bcs(self): bcs = [] @@ -330,86 +272,6 @@ def _recover_bcs(self): bcs.append(c_rep) return bcs - def _create_replace_map(self, form: ufl.Form | NestedMutableSequence[ufl.Form] | None) -> dict[Function, Function]: - """Replace dependencies with latest checkpoint.""" - replace_map = {} - for block_variable in self.get_dependencies(): - coeff = block_variable.output - if isinstance(form, ufl.Form): - if coeff in form.coefficients(): - replace_map[coeff] = block_variable.saved_output - elif form is None: - return {} - else: - for f in form: - replace_map.update(self._create_replace_map(f)) - for block_variable in self.get_outputs(): - coeff = block_variable.output - if isinstance(form, ufl.Form): - if coeff in form.coefficients(): - replace_map[coeff] = block_variable.saved_output - elif form is None: - return {} - else: - for f in form: - replace_map.update(self._create_replace_map(f)) - return replace_map - - 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.""" - - # 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() - - # 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): - self._u.x.array[:] = 0.0 - else: - for ui in self._u: - ui.x.array[:] = 0.0 - # 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(): - dolfinx.fem.petsc.LinearProblem.solve(problem) - return self._u - - def recompute_component( - self, - inputs: typing.Iterable[Function], - block_variable: pyadjoint.block_variable.BlockVariable, - idx: int, - prepared: dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function], - ) -> dolfinx.fem.Function: - """Recompute and return an isolated copy of the solution state.""" - if isinstance(prepared, dolfinx.fem.Function): - assert idx == 0 - # Return an explicit copy so each tape block gets an isolated state snapshot - return prepared.copy() - else: - assert isinstance(prepared, typing.Iterable) - return prepared[idx].copy() - def _should_compute_boundary_adjoint( self, relevant_dependencies: typing.List[tuple[int, pyadjoint.block_variable.BlockVariable]] ) -> bool: @@ -428,29 +290,21 @@ def _compute_adjoint(cls, form: ufl.Form) -> typing.Sequence[typing.Sequence[ufl """ return ufl.extract_blocks(compute_form_adjoint(form)) - def _compute_residual(self) -> tuple[ufl.Form, dict[Function, Function]]: - """Convert the formulation :math:`a(u, v)=L(v)` into a residual :math:`F(u_b, v) = 0` where - :math:`u_b` is the solution of the forward problem at the current time and all coefficients are updated. - """ - # NOTE: Should probably be possible to compile this form once. - replacement_functions = self.get_outputs() - r_funcs = ( - [r.saved_output for r in replacement_functions] - if len(replacement_functions) > 1 - else replacement_functions[0].saved_output - ) - summed_form = sum_form(self._lhs) - F_form = ufl.action(summed_form, r_funcs) - sum_form(self._rhs) - replacement_map = self._create_replace_map(F_form) - F_form = ufl.replace(F_form, replacement_map) - return F_form, replacement_map + def _create_replace_map(self, form: ufl.Form | NestedMutableSequence[ufl.Form] | None) -> dict[Function, Function]: + """Replace dependencies with latest checkpoint.""" + replace_map: dict = {} + replace_map.update(_map_block_variables_to_form(form, self.get_dependencies())) + replace_map.update(_map_block_variables_to_form(form, self.get_outputs())) + return replace_map - def _compute_residual_derivative(self) -> typing.Union[ufl.Form, list[list[ufl.Form]]]: - """Compute the derivative of the residual with respect to the outputs.""" + def _compute_residual_derivative(self) -> ufl.Form | list[list[ufl.Form]]: + """Compute the derivative of the residual with respect to the outputs. - res = self._compute_residual() - F_form = res[0] if isinstance(res, tuple) else res - assert isinstance(F_form, ufl.Form), "Residual form must be a single UFL form." + Shared by both Problem kinds: built purely from ``self._compute_residual()``'s output + (a single summed form, the same shape for both classes), so no per-class override is + needed even though the two classes construct that residual very differently. + """ + F_form, _ = self._compute_residual() outputs = self.get_outputs() # Use r.saved_output directly; no lookup in replacement_map needed! @@ -473,15 +327,15 @@ def prepare_evaluate_tlm( ) -> typing.Sequence[Function] | dolfinx.fem.Function: # 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 + # verbatim with dF/du (see *Problem._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() + # *Problem._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() @@ -499,10 +353,10 @@ def prepare_evaluate_tlm( # 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 + # *Problem._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 + b_petsc = tlm_solver.b with b_petsc.localForm() as b_loc: b_loc.set(0.0) for block_variable in self.get_dependencies(): @@ -551,18 +405,14 @@ def prepare_evaluate_adj( """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() + # shared across every block this Problem records. Refresh this + # block's own checkpointed values into the placeholders, refresh + # whatever "state" dF/du is evaluated at if it depends on one (a + # no-op for LinearProblem, see _refresh_dFdu_state), 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(): @@ -570,18 +420,19 @@ def prepare_evaluate_adj( if placeholder is not None: placeholder.x.array[:] = block_variable.saved_output.x.array[:] placeholder.x.scatter_forward() + self._refresh_dFdu_state(problem) # Extract dJ/du[v] from the adjoint inputs. if len(adj_inputs) == 1: adj_rhs = adj_inputs[0] - dJdu = 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 = adjoint_solver._b + dJdu = adjoint_solver.b with dJdu.localForm() as dJdu_loc: dJdu_loc.set(0.0) @@ -652,22 +503,102 @@ def evaluate_adj_component( assemble_compiled_form(compiled_sensitivity, tensor=vec) return vec + def prepare_recompute_component( + self, inputs: typing.Sequence[typing.Any], relevant_outputs: typing.Sequence[typing.Any] + ) -> _Function | typing.Sequence[_Function]: + """Prepare for recomputing the block with different control inputs, and solve. + + 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``/ + ``NonlinearProblem._value_placeholders``): 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. + + The Problem's own unknown(s) are warm-started from this block's own + saved outputs rather than zeroed: required for SNES/Newton + convergence, and applied to a KSP-based ``LinearProblem`` the same + way, so an iterative solver configured with a nonzero initial guess + benefits from it too -- both classes now solve by refreshing + placeholder/unknown values and calling an unchanging, already-built + solver, never by mutating the user's own coefficient or recompiling a + form. + + Solving happens once, here -- not once per output in + ``recompute_component`` -- which matters for a multi-output (blocked) + problem. The base ``dolfinx`` ``solve()`` is called directly (via + ``problem._dolfinx_solve()``), not ``problem.solve()``, which would + record another block onto the tape. + """ + 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() + + # 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 the Problem's own unknown(s) from this block's own saved + # outputs. + 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() + + with pyadjoint.stop_annotating(): + problem._dolfinx_solve() + return problem._u + + def recompute_component( + self, + inputs: typing.Iterable[Function], + block_variable: pyadjoint.block_variable.BlockVariable, + idx: int, + prepared: _Function | typing.Sequence[_Function], + ) -> Function: + """Return an isolated copy of this block's own share of the already-recomputed state.""" + if isinstance(prepared, dolfinx.fem.Function): + assert idx == 0 + # Return an explicit copy so each tape block gets an isolated state snapshot + return prepared.copy() + else: + assert isinstance(prepared, typing.Sequence) + return prepared[idx].copy() + def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_dependencies): - # First fetch all relevant values + """Assemble and solve the second-order-adjoint (SOA) equation. + + Scalar (single-output) problems share one implementation for both + Problem kinds, built entirely from the Problem's cached Hessian/TLM + templates (see ``HessianTemplates``) -- no block-specific residual + construction needed, since the SOA self-term is already correctly + zero/nonzero per class (see ``_build_soa_self_template``). A blocked + (multi-output) problem has no templated fast path; its right-hand + side is assembled by ``_evaluate_hessian_blocked_rhs``, also shared + by both Problem kinds, built from ``_compute_residual_derivative``/ + ``_compute_adjoint`` -- themselves built only from the per-class + ``_compute_residual`` hook (the two Problem kinds start from + different user-supplied data there). + """ outputs = self.get_outputs() tlm_output = [output.tlm_value for output in outputs if output is not None] if (hessian_inputs is None) or (len(tlm_output) == 0): return # 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 + # verbatim with the first-order adjoint equation -- is 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() + problem = self._problem adjoint_solver = problem._get_or_build_adjoint_solver() adjoint_solver.bcs = self._bcs for block_variable in self.get_dependencies(): @@ -675,14 +606,15 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ if placeholder is not None: placeholder.x.array[:] = block_variable.saved_output.x.array[:] placeholder.x.scatter_forward() + self._refresh_dFdu_state(problem) 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 SOA right-hand side symbolically - # on every call. + # Use the cached per-dependency Hessian templates (see + # *Problem._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() + hessian_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] @@ -691,9 +623,10 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ 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 + b = adjoint_solver.b with b.localForm() as b_loc: b_loc.set(0.0) + dolfinx.fem.petsc.assemble_vector(b, hessian_templates.soa_self) for block_variable in self.get_dependencies(): tlm_input = block_variable.tlm_value if tlm_input is None: @@ -701,7 +634,7 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ 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) + template = hessian_templates.soa_cross.get(c) if template is None: continue seed = seed_placeholders[c] @@ -714,76 +647,13 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ with b.localForm() as b_loc: b_loc.array[:] += hessian_inputs[0].array[:] b.ghostUpdate(PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) - 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): - if hess_input is not None: - bi = dolfinx.la.vector(hess_input.index_map, hess_input.block_size) - else: - out_i = self.get_outputs()[i].saved_output - bi = dolfinx.la.vector( - out_i.function_space.dofmap.index_map, out_i.function_space.dofmap.index_map_bs - ) - bs.append(bi) - bi.array[:] = 0.0 - form_i = ufl.algorithms.apply_derivatives.apply_derivatives(b_form[i]) - 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.assemble_vector(bi.array, compiled_soa_rhs) - bi.scatter_reverse(dolfinx.la.InsertMode.add) - bi.scatter_forward() - bi.array[:] *= -1 - - if hess_input is not None: - bi.array[:] += hess_input.array - - bi.scatter_forward() - 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) + self._evaluate_hessian_blocked_rhs(adjoint_solver, hessian_inputs, tlm_output, relevant_dependencies) # 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. + # 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. adjoint_solver.solve() if isinstance(self._second_adjoint_solutions, list): for adj_sol, sol in zip(self._second_adjoint_solutions, adjoint_solver.u): @@ -798,6 +668,76 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ return self._compute_residual(), self._adjoint_solutions, self._second_adjoint_solutions + def _evaluate_hessian_blocked_rhs(self, adjoint_solver, hessian_inputs, tlm_output, relevant_dependencies): + """Assemble the SOA right-hand side for a blocked (multi-output) problem into ``adjoint_solver.b``. + + No templated fast path exists for a blocked problem (see + ``*Problem._get_or_build_hessian_templates``): the right-hand side is built symbolically, + per call, from ``_compute_residual_derivative``/``_compute_adjoint`` -- the same two + hooks the scalar path's templates are themselves built from -- so this one implementation + is shared by both Problem kinds. ``d2Fdu2`` (the SOA self-term) is always included + unconditionally rather than asserted zero: it comes out zero for a linear residual as a + result of the maths, exactly like ``_build_soa_self_template``'s scalar-path self-term -- + and, mirroring that helper, is reduced from bilinear (test and trial) to linear (test + only) by adjointing and acting on the first-order adjoint solution before use, since + ``ufl.derivative`` w.r.t. a *coefficient* (``unknowns``) leaves ``dF/du``'s own trial + argument untouched. + """ + 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)) + + if d2Fdu2.empty(): + b_form = d2Fdu2 + else: + b_form = 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) + + bs = [] + b_form = ufl.extract_blocks(b_form) + for i, hess_input in enumerate(hessian_inputs): + if hess_input is not None: + bi = dolfinx.la.vector(hess_input.index_map, hess_input.block_size) + else: + out_i = self.get_outputs()[i].saved_output + bi = dolfinx.la.vector(out_i.function_space.dofmap.index_map, out_i.function_space.dofmap.index_map_bs) + bs.append(bi) + bi.array[:] = 0.0 + form_i = ufl.algorithms.apply_derivatives.apply_derivatives(b_form[i]) + 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.assemble_vector(bi.array, compiled_soa_rhs) + bi.scatter_reverse(dolfinx.la.InsertMode.add) + bi.scatter_forward() + bi.array[:] *= -1 + + if hess_input is not None: + bi.array[:] += hess_input.array + + bi.scatter_forward() + 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) + def evaluate_hessian_component( self, inputs, @@ -808,25 +748,27 @@ def evaluate_hessian_component( relevant_dependencies, prepared=None, ): - c = block_variable.output - - (F_form, replacement_map), adj_sol, adj_sol2 = prepared + """Return this dependency's contribution to the Hessian action. + Scalar (single-output) problems share one implementation for both + Problem kinds (see ``prepare_evaluate_hessian``); a blocked + (multi-output) problem falls back to + ``_evaluate_hessian_component_blocked``. + """ + c = block_variable.output + residual_prepared, adj_sol, adj_sol2 = prepared outputs = self.get_outputs() tlm_output = [output.tlm_value for output in outputs] - c_rep = block_variable.saved_output # If m = DirichletBC then d^2F(u,m)/dm^2 = 0 and d^2F(u,m)/dudm = 0, # so we only have the term dF(u,m)/dm * adj_sol2 if isinstance(c, dolfinx.fem.DirichletBC): raise NotImplementedError("Hessian computation for DirichletBC control not implemented yet.") - if isinstance(c_rep, dolfinx.fem.Constant): raise NotImplementedError("Hessian computation for Constant control not implemented yet.") # mesh = extract_mesh_from_form(F_form) # W = c._ad_function_space(mesh) - elif isinstance(c, dolfinx.mesh.Mesh): raise NotImplementedError("Hessian computation for Mesh control not implemented yet.") # X = dolfin.SpatialCoordinate(c) @@ -836,8 +778,7 @@ def evaluate_hessian_component( 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 + # Use the cached per-dependency 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 @@ -845,14 +786,14 @@ def evaluate_hessian_component( # 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() + # *Problem._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 + hessian_templates = problem._get_or_build_hessian_templates() _, seed_placeholders, _ = problem._get_or_build_tlm_rhs_templates() - fixed_template = fixed_templates[c] + fixed_template = hessian_templates.fixed[c] hessian_output = _create_vector(fixed_template, W) hessian_output.array[:] = 0.0 assemble_compiled_form(fixed_template, hessian_output) @@ -864,7 +805,7 @@ def evaluate_hessian_component( tlm_input = bv.tlm_value if tlm_input is None: continue - template = cross_templates.get((c, c2)) + template = hessian_templates.cross.get((c, c2)) if template is None: continue seed2 = seed_placeholders[c2] @@ -875,15 +816,29 @@ def evaluate_hessian_component( 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 - # mixed. Therefore, we instead we compute it as dL/dm = d(lambda_i^T F(m))/dm, - # which is equivalent. + return self._evaluate_hessian_component_blocked( + residual_prepared, adj_sol, adj_sol2, c, c_rep, W, tlm_output, relevant_dependencies + ) + + def _evaluate_hessian_component_blocked( + self, residual_prepared, adj_sol, adj_sol2, c, c_rep, W, tlm_output, relevant_dependencies + ): + """Return this dependency's Hessian-action contribution for a blocked (multi-output) problem. + + No templated fast path exists for a blocked problem; this builds it symbolically, per + call, from the residual ``prepare_evaluate_hessian`` already recorded + (``residual_prepared``) -- shared by both Problem kinds, since both settle on the same + ``_compute_residual`` output shape. + + 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 + mixed. Therefore, we instead we compute it as dL/dm = d(lambda_i^T F(m))/dm, + which is equivalent. + """ + F_form, replacement_map = residual_prepared + outputs = self.get_outputs() + F_summed = sum_form(F_form) L1 = ufl.action(F_summed, adj_sol) L2 = ufl.action(F_summed, adj_sol2) @@ -935,7 +890,181 @@ def evaluate_hessian_component( return hessian_output -class NonlinearProblemBlock(pyadjoint.Block): +class LinearProblemBlock(_ProblemBlockBase): + """A linear problem that can be used with adjoint methods. + + This class extends the `dolfinx.fem.petsc.LinearProblem` to support adjoint methods. + """ + + _adjoint_solutions: dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function] + _tlm_solutions: dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function] + _second_adjoint_solutions: dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function] + + # 2. Overload for the SCALAR case + @typing.overload + def __init__( + self, + a: ufl.Form, + L: ufl.Form, + *, + bcs: typing.Sequence[dolfinx.fem.DirichletBC] | None = None, + u: _Function | None = None, + P: ufl.Form | 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, + problem: "LinearProblem" = ..., + ) -> None: ... + + @typing.overload + def __init__( + self, + a: typing.Sequence[typing.Sequence[ufl.Form]], + L: typing.Sequence[ufl.Form], + *, + bcs: typing.Sequence[dolfinx.fem.DirichletBC] | None = None, + u: typing.Sequence[_Function] | None = None, + P: typing.Sequence[typing.Sequence[ufl.Form]] | 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, + problem: "LinearProblem" = ..., + ) -> None: ... + + def __init__( + self, + a: ufl.Form | typing.Sequence[typing.Sequence[ufl.Form]], + L: ufl.Form | typing.Sequence[ufl.Form], + *, + 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, + 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, + problem: "LinearProblem" = None, # type: ignore[assignment] + ) -> None: + + assert problem is not None, "problem must be provided." + # Strong reference, deliberately: a throwaway LinearProblem (solve it, + # then only touch the ReducedFunctional) is a common pattern, so the + # block must keep the Problem -- and its shared solvers -- alive for + # as long as the block itself is reachable. Not cyclic garbage on its + # own (verified), so this doesn't reintroduce the MPI + # collective-destruction hazard from dolfinx-adjoint-knowledge's + # mpi-collective-destruction-hazard note -- that needs an unmerged + # checkpoint schedule making the *tape* cyclic. + self._problem_obj = problem + super().__init__(ad_block_tag=ad_block_tag) + + # Collect all arguments in variational forms and replace them with similar + # once that is based on a mixed functionspace. + if not isinstance(a, ufl.Form): + a, L = assign_mixed_parts(a, L) + if P is not None: + P, _ = assign_mixed_parts(P, L) + + self._lhs = a + self._rhs = L + self._preconditioner = P + + # Create overloaded functions + self._u: _Function | typing.Sequence[_Function] + if isinstance(u, dolfinx.fem.Function): + self._u = pyadjoint.create_overloaded_object(u) + elif u is None: + try: + # Extract function space for unknown from the right hand + # side of the equation. + self._u = Function(L.arguments()[0].ufl_function_space()) # type: ignore + except AttributeError: + self._u = [Function(Li.arguments()[0].ufl_function_space()) for Li in L] # type: ignore[union-attr] + else: + self._u = [pyadjoint.create_overloaded_object(ui) for ui in u] + + # NOTE: Add mesh and constants as dependencies later on + + # To ensure that the solver can be recycled in time dependent loops, the unknown is also added as a dependency + # if present in the form. + if isinstance(self._u, dolfinx.fem.Function): + assert isinstance(self._lhs, ufl.Form) + assert isinstance(self._rhs, ufl.Form) + if self._u in self._lhs.coefficients() or self._u in self._rhs.coefficients(): + raise RuntimeError("The unknown function u should not be present in the variational forms a or L.") + for c in self._lhs.coefficients(): + self.add_dependency(c, no_duplicates=True) + for c in self._rhs.coefficients(): + self.add_dependency(c, no_duplicates=True) + elif isinstance(self._u, typing.Iterable): + for Ai in self._lhs: # type: ignore + for Aij in Ai: + if Aij is not None: + assert isinstance(Aij, ufl.Form) + for c in Aij.coefficients(): + if c in self._u: + raise RuntimeError( + "The unknown function u should not be present in the variational forms a or L." + ) + self.add_dependency(c, no_duplicates=True) + for part in self._rhs: # type: ignore + for c in part.coefficients(): + if c in self._u: + raise RuntimeError( + "The unknown function u should not be present in the variational forms a or L." + ) + self.add_dependency(c, no_duplicates=True) + else: + raise RuntimeError(f"Unknown type for unknown function u={type(self._u)}.") + # 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._bcs = bcs if bcs is not None else [] + + # Add dependencies from the boundary conditions + if self._bcs is not None: + for bc in self._bcs: + if hasattr(bc, "block_variable"): + self.add_dependency(bc, no_duplicates=True) + + # 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() + self._second_adjoint_solutions = self._u.copy() + self._tlm_solutions = self._u.copy() + else: + assert isinstance(self._u, typing.Iterable) + self._adjoint_solutions = [u.copy() for u in self._u] + self._second_adjoint_solutions = [u.copy() for u in self._u] + self._tlm_solutions = [u.copy() for u in self._u] + + def _compute_residual(self) -> tuple[ufl.Form, dict[Function, Function]]: + """Convert the formulation :math:`a(u, v)=L(v)` into a residual :math:`F(u_b, v) = 0` where + :math:`u_b` is the solution of the forward problem at the current time and all coefficients are updated. + """ + # NOTE: Should probably be possible to compile this form once. + replacement_functions = self.get_outputs() + r_funcs = ( + [r.saved_output for r in replacement_functions] + if len(replacement_functions) > 1 + else replacement_functions[0].saved_output + ) + summed_form = sum_form(self._lhs) + F_form = ufl.action(summed_form, r_funcs) - sum_form(self._rhs) + replacement_map = self._create_replace_map(F_form) + F_form = ufl.replace(F_form, replacement_map) + return F_form, replacement_map + + +class NonlinearProblemBlock(_ProblemBlockBase): """A linear problem that can be used with adjoint methods. This class extends the `dolfinx.fem.petsc.LinearProblem` to support adjoint methods. @@ -1014,16 +1143,10 @@ def __init__( # NOTE: Add mesh and constants as dependencies later on u_list = self._u if isinstance(self._u, list) else [self._u] - if J is not None: - assert isinstance(J, ufl.Form) - for c in J.coefficients(): - if c not in u_list: # Exclude unknown - self.add_dependency(c, no_duplicates=True) - if self._rhs is not None: - assert isinstance(self._rhs, ufl.Form) - for c in self._rhs.coefficients(): - if c not in u_list: # Exclude unknown - self.add_dependency(c, no_duplicates=True) + for c in _collect_coefficients(J) - set(u_list): + self.add_dependency(c, no_duplicates=True) + for c in _collect_coefficients(self._rhs) - set(u_list): + self.add_dependency(c, no_duplicates=True) # Cache form parameters for later # NOTE: Should probably be in a struct @@ -1033,7 +1156,7 @@ def __init__( self._bcs = bcs if bcs is not None else [] # No forward/adjoint solver is built here: this block shares the ones - # owned by self._problem() (see NonlinearProblem in ../solvers.py), + # 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. @@ -1047,483 +1170,35 @@ def __init__( self._second_adjoint_solutions = [u.copy() for u in self._u] self._tlm_solutions = [u.copy() for u in self._u] - def _problem(self) -> "NonlinearProblem": - """Return this block's owning Problem, which owns the shared solvers.""" - return self._problem_obj - - def _recover_bcs(self): - bcs = [] - for block_variable in self.get_dependencies(): - c = block_variable.output - c_rep = block_variable.saved_output - - if isinstance(c, dolfinx.fem.DirichletBC): - bcs.append(c_rep) - return bcs - - def _create_replace_map(self, form: ufl.Form) -> dict[Function, Function]: - """Replace dependencies with latest checkpoint.""" - replace_map = {} - for block_variable in self.get_dependencies(): - coeff = block_variable.output - if coeff in form.coefficients(): - replace_map[coeff] = block_variable.saved_output - return replace_map - - def _replace_coefficients_in_form(self, form: ufl.Form) -> ufl.Form: - """Replace coefficients in the form with saved outputs. - - Args: - form: The UFL form to replace coefficients in. - """ - replace_map = self._create_replace_map(form) - return ufl.replace(form, replace_map) - - def prepare_recompute_component(self, inputs, relevant_outputs): - """Prepare for recomputing the block with different control inputs.""" - - # 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(): - 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 = 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() - - return None - - 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(): - dolfinx.fem.petsc.NonlinearProblem.solve(problem) - if isinstance(problem._u, list): - output = problem._u[idx] - else: - output = problem._u - assert isinstance(output, Function) - return output - - def _should_compute_boundary_adjoint( - self, relevant_dependencies: typing.List[tuple[int, pyadjoint.block_variable.BlockVariable]] - ) -> bool: - bdy = False - for _, dep in relevant_dependencies: - if isinstance(dep.output, dolfinx.fem.DirichletBC): - bdy = True - break - return bdy - - @classmethod - @typing.overload - def _compute_adjoint( - cls, form: typing.Sequence[typing.Sequence[ufl.Form]] - ) -> typing.Sequence[typing.Sequence[ufl.Form]]: ... - - @classmethod - @typing.overload - def _compute_adjoint(cls, form: ufl.Form) -> ufl.Form: ... + def _compute_residual(self) -> tuple[ufl.Form, dict[Function, Function]]: + """Build the residual :math:`F(u_b, v) = 0` at the current checkpointed dependency values. - @classmethod - def _compute_adjoint( - cls, form: ufl.Form | typing.Sequence[typing.Sequence[ufl.Form]] - ) -> ufl.Form | typing.Sequence[typing.Sequence[ufl.Form]]: - """ - Compute adjoint of a bilinear form :math:`a(u, v)`, which could be written as a blocked system. - """ - if isinstance(form, ufl.Form): - return ufl.adjoint(form) - else: - assert isinstance(form, typing.Iterable) - adj_form: list[list[ufl.Form]] = [] - tmp_form: list[list[ufl.Form]] = [] - for i, f_i in enumerate(form): - tmp_form.append([]) - adj_form.append([]) - for j, form_ij in enumerate(f_i): - tmp_form[i].append(ufl.adjoint(form_ij)) - adj_form[i].append(ufl.adjoint(form_ij)) - for i, f_i in enumerate(tmp_form): - for j, form_ij in enumerate(f_i): - adj_form[j][i] = form_ij - return adj_form - - def _compute_residual(self) -> typing.Union[ufl.Form, list[ufl.Form]]: - """Convert the formulation :math:`a(u, v)=L(v)` into a residual :math:`F(u_b, v) = 0` where - :math:`u_b` is the solution of the forward problem at the current time and all coefficients are updated. + Unlike ``LinearProblemBlock`` (which derives ``F`` from ``a``/``L`` via ``ufl.action``), + ``F`` is already the residual here -- this only needs to substitute in the checkpointed + dependency and output values. Settles on the same output shape as + ``LinearProblemBlock._compute_residual`` (a single summed form plus its replacement map) + so every method built on top of it (adjoint, TLM, Hessian) is shared on the base class. """ - # NOTE: Should probably be possible to compile this form once. replacement_functions = self.get_outputs() - assert isinstance(self._rhs, (ufl.Form, typing.Sequence)) - assert isinstance(self._rhs, ufl.Form) replacement_map = self._create_replace_map(self._rhs) u_list = self._u if isinstance(self._u, list) else [self._u] for u, block in zip(u_list, replacement_functions): replacement_map[u] = block.saved_output - if isinstance(self._u, dolfinx.fem.Function): - F_form = ufl.replace(self._rhs, replacement_map) - else: - assert isinstance(self._rhs, typing.Iterable) - F_form = [ufl.replace(rhs_j, replacement_map) for rhs_j in self._rhs] - return F_form - - def _compute_residual_derivative(self) -> typing.Union[ufl.Form, list[list[ufl.Form]]]: - """Compute the derivative of the residual with respect to the outputs.""" - - F_form = self._compute_residual() - outputs = [output.saved_output for output in self.get_outputs()] - if len(outputs) == 1: - assert isinstance(F_form, ufl.Form) - dFdu = ufl.derivative(F_form, outputs[0], ufl.TrialFunction(outputs[0].function_space)) - else: - assert isinstance(F_form, list) - dFdu = [] - for i in range(len(outputs)): - dFdu.append([]) - for j in range(len(outputs)): - dFdu[-1].append(ufl.derivative(F_form[i], outputs[j], ufl.TrialFunction(outputs[j].function_space))) - return dFdu - - def prepare_evaluate_tlm( - self, inputs, tlm_inputs, relevant_outputs - ) -> 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() - - 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 - if tlm_value is None: - continue - 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) # type: ignore[arg-type] - - 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 - - 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: - assert isinstance(self._tlm_solutions, dolfinx.fem.Function) - return self._tlm_solutions - - def prepare_evaluate_adj( - self, - inputs: typing.Sequence[Function], - adj_inputs: typing.Sequence[dolfinx.la.Vector], - relevant_dependencies: typing.List[tuple[int, pyadjoint.block_variable.BlockVariable]], - ) -> typing.Union[ufl.Form, typing.Iterable[ufl.Form]]: - """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: 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() - # Extract dJ/du[v] from the adjoint inputs. - assert len(adj_inputs) == 1 - adj_rhs = adj_inputs[0] - 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[:] - - adjoint_solver.solve() - if isinstance(self._adjoint_solutions, list): - 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[:] = adjoint_solver.u.x.array[:] - return F_form - - def evaluate_adj_component( - self, - inputs: typing.Iterable[Function], - adj_inputs: typing.Iterable[dolfinx.la.Vector], - block_variable: pyadjoint.block_variable.BlockVariable, - idx: int, - prepared: typing.Union[ufl.Form, typing.Iterable[ufl.Form]], - ) -> typing.Union[_SpecialVector, typing.Iterable[_SpecialVector]]: - """Evaluate the adjoint component, i.e. :math:`\frac{\\partial Au - b}{\\partial c}`.""" - - residual = prepared - - c = block_variable.output - - c_rep = block_variable.saved_output - if isinstance(c, dolfinx.fem.Function): - dc = ufl.TrialFunction(c.function_space) - else: - raise NotImplementedError(f"Unsupported control {type(c)}") - dFdm = -ufl.derivative(residual, c_rep, dc) - - # Safe return for empty sensitivities - if dFdm.empty(): - dFdm = ufl.ZeroBaseForm((dc,)) - - dFdm_adj = ufl.adjoint(dFdm) - sensitivity = ufl.action(dFdm_adj, self._adjoint_solutions) - compiled_sensitivity = dolfinx.fem.form( - sensitivity, - jit_options=self._jit_options, - form_compiler_options=self._form_compiler_options, - entity_maps=self._entity_maps, - ) - vec = _create_vector(compiled_sensitivity, sensitivity.arguments()[0].ufl_function_space()) - vec.array[:] = 0.0 - assemble_compiled_form(compiled_sensitivity, tensor=vec) - return vec - - def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_dependencies): - # First fetch all relevant values - outputs = self.get_outputs() - tlm_output = [output.tlm_value for output in outputs if output is not None] - if (hessian_inputs is None) or (len(tlm_output) == 0): - return - - # 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." - assert len(hessian_inputs) == 1, "Hessian computation only implemented for single hessian input blocks." - - # 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.") - 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[:] - - # 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, adjoint_solver.u): - adj_sol.x.array[:] = sol.x.array[:] - else: - 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( - self, - inputs, - hessian_inputs, - adj_inputs, - block_variable, - idx, - relevant_dependencies, - prepared=None, - ): - c = block_variable.output - - # 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." - - c_rep = block_variable.saved_output - - # If m = DirichletBC then d^2F(u,m)/dm^2 = 0 and d^2F(u,m)/dudm = 0, - # so we only have the term dF(u,m)/dm * adj_sol2 - if isinstance(c, dolfinx.fem.DirichletBC): - raise NotImplementedError("Hessian computation for DirichletBC control not implemented yet.") - - if isinstance(c_rep, dolfinx.fem.Constant): - raise NotImplementedError("Hessian computation for Constant control not implemented yet.") - # mesh = extract_mesh_from_form(F_form) - # W = c._ad_function_space(mesh) - - elif isinstance(c, dolfinx.mesh.Mesh): - raise NotImplementedError("Hessian computation for Mesh control not implemented yet.") - # X = dolfin.SpatialCoordinate(c) - # W = c._ad_function_space() - else: - assert isinstance(c, dolfinx.fem.Function) - W = c.function_space + F_form = ufl.replace(sum_form(self._rhs), replacement_map) + return F_form, replacement_map - # 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) + def _refresh_dFdu_state(self, problem: typing.Any) -> None: + """Refresh ``problem``'s "state" placeholder(s) from this block's own saved outputs. - 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 + Unlike ``LinearProblem`` (where ``dF/du`` never references the state), ``dF/du`` here + genuinely depends on ``u``'s current value, so the shared adjoint solver's compiled LHS + must be evaluated at *this* block's checkpointed output before it is used -- another + block sharing the same solver may have left a different value there. + """ + state_placeholder = problem._residual_state_placeholder + 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() diff --git a/src/dolfinx_adjoint/petsc_utils.py b/src/dolfinx_adjoint/petsc_utils.py index e1c7f18..3a27e50 100644 --- a/src/dolfinx_adjoint/petsc_utils.py +++ b/src/dolfinx_adjoint/petsc_utils.py @@ -11,8 +11,8 @@ def solve_linear_problem( A: PETSc.Mat, # type: ignore [name-defined] x: dolfinx.la.Vector, b: dolfinx.la.Vector, - petsc_options: typing.Optional[dict] = None, - P: typing.Optional[PETSc.Mat] = None, # type: ignore [name-defined] + petsc_options: dict | None = None, + P: PETSc.Mat | None = None, # type: ignore [name-defined] ): """Solve a linear problem :math:`Ax = b`. @@ -57,13 +57,29 @@ def solve_linear_problem( x.scatter_forward() -class LinearAdjointProblem(dolfinx.fem.petsc.LinearProblem): - """Linear problem helper class that homogenizes the boundary conditions, meaning that no lifting is applied.""" +class HomogeneousBCLinearProblem(dolfinx.fem.petsc.LinearProblem): + """Linear problem helper class that homogenizes the boundary conditions, meaning that no lifting is applied. + + Used for both the adjoint and the tangent-linear solve -- neither is "the adjoint problem" + specifically, they are both just a linear solve against a zero-lifted right-hand side, so the + name no longer singles out one of the two callers. + """ def solve( self, ) -> typing.Union[dolfinx.fem.Function, typing.Sequence[dolfinx.fem.Function]]: - """Solve the problem.""" + """Solve the problem. + + Unlike the base class, ``self.b`` is never (re-)assembled from the compiled + right-hand-side form here. This solver is shared, and reused verbatim, across every + block a Problem records (see dolfinx-adjoint-knowledge's solver-reuse note), so the + caller (see ``_ProblemBlockBase.prepare_evaluate_adj``/``prepare_evaluate_hessian``/ + ``prepare_evaluate_tlm``) has already assembled its own right-hand side directly into + ``self.b`` before calling ``solve()``. The only thing this method does to ``self.b`` is + modify it in place -- zeroing every Dirichlet-BC dof (``alpha=0.0``) -- rather than + lifting, since a caller's right-hand side is never the original problem's own Dirichlet + data. + """ # Assemble lhs self._A.zeroEntries() diff --git a/src/dolfinx_adjoint/solvers.py b/src/dolfinx_adjoint/solvers.py index 7e312e9..c4709f3 100644 --- a/src/dolfinx_adjoint/solvers.py +++ b/src/dolfinx_adjoint/solvers.py @@ -10,33 +10,16 @@ from .blocks.solvers import ( LinearProblemBlock, NonlinearProblemBlock, + _collect_coefficients, + _ProblemBlockBase, assign_mixed_parts, get_sorted_arguments, sum_form, ) -from .petsc_utils import LinearAdjointProblem +from .petsc_utils import HomogeneousBCLinearProblem 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: @@ -89,196 +72,173 @@ def resolve_u( return [pyadjoint.create_overloaded_object(ui) for ui in u] -class LinearProblem(dolfinx.fem.petsc.LinearProblem): - """A linear problem that can be used with adjoint methods. - - This class extends the `dolfinx.fem.petsc.LinearProblem` to support adjoint methods. - - Args: - a: The bilinear form representing the left-hand side of the equation. - L: The linear form representing the right-hand side of the equation. - bcs: Boundary conditions to apply to the problem. - u: Solution vector. - P: Preconditioner for the linear problem. - kind: Kind of PETSc Matrix to assemble the system into. - petsc_options: Options dictionary for the PETSc krylov supspace solver. - form_compiler_options: Form compiler options for generating assembly kernels. - jit_options: Options for just-in-time compilation of the forms. - entity_maps: Mapping from meshes that coefficients and arguments are defined on to the - integration domain of the forms. - ad_block_tag: Tag for adjoint blocks in the tape. - adjoint_petsc_options: PETSc options for adjoint problems. - tlm_petsc_options: Optional PETSc options for TLM problems. +class HessianTemplates(typing.NamedTuple): + """Per-dependency compiled templates used to assemble a Hessian action. + + Shared shape for ``LinearProblem``/``NonlinearProblem``, so callers in + ``blocks/solvers.py`` unpack by name rather than by position -- the two + classes used to return differently-shaped tuples here, which was a latent + footgun for any code touching both. + + Attributes: + soa_self: The SOA right-hand-side's contribution from ``dF/du``'s own + second derivative w.r.t. ``u`` (``d2Fdu2``) -- a + ``ufl.ZeroBaseForm`` if that term is structurally zero, so callers + can always assemble it unconditionally. Built by the same code + (``_build_soa_self_template``) for both classes; it simply always + compiles to zero for a linear problem, since ``dF/du`` doesn't + reference ``u`` there -- not a hardcoded special case. + soa_cross: The SOA right-hand-side's contribution, per dependency, + from that dependency's tangent-linear direction. + fixed: The part of each dependency's own Hessian-action output that + does not depend on any *other* dependency's tangent-linear value. + cross: Each dependency's Hessian-action contribution from *another* + dependency's tangent-linear direction, keyed by ``(c, c2)``. """ - @typing.overload - def __init__( - self, - a: ufl.Form, - L: ufl.Form, - *, - 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, - petsc_options_prefix: str = "dxa_linear_problem_", - form_compiler_options: dict | None = None, - jit_options: dict | None = None, - entity_maps: typing.Sequence[dolfinx.mesh.EntityMap] | None = None, - ad_block_tag: typing.Optional[str] = None, - adjoint_petsc_options: typing.Optional[dict] = None, - tlm_petsc_options: typing.Optional[dict] = None, - ) -> None: ... - @typing.overload - def __init__( - self, - a: typing.Sequence[typing.Sequence[ufl.Form]], - L: typing.Sequence[ufl.Form], - *, - 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, - petsc_options_prefix: str = "dxa_linear_problem_", - form_compiler_options: dict | None = None, - jit_options: dict | None = None, - entity_maps: typing.Sequence[dolfinx.mesh.EntityMap] | None = None, - ad_block_tag: typing.Optional[str] = None, - adjoint_petsc_options: typing.Optional[dict] = None, - tlm_petsc_options: typing.Optional[dict] = None, - ) -> None: ... - def __init__( - self, - a: ufl.Form | typing.Sequence[typing.Sequence[ufl.Form]], - L: ufl.Form | typing.Sequence[ufl.Form], - *, - 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, - petsc_options_prefix: str = "dxa_linear_problem_", - form_compiler_options: dict | None = None, - jit_options: dict | None = None, - entity_maps: typing.Sequence[dolfinx.mesh.EntityMap] | None = None, - ad_block_tag: typing.Optional[str] = None, - adjoint_petsc_options: typing.Optional[dict] = None, - tlm_petsc_options: typing.Optional[dict] = None, - ) -> None: - 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 - self._lhs = a - self._rhs = L - self._jit_options = jit_options - 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=_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=_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, - form_compiler_options=form_compiler_options, - jit_options=jit_options, - 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 + soa_self: dolfinx.fem.Form + soa_cross: dict + fixed: dict + cross: dict + + +def _build_soa_self_template( + dFdu_template: ufl.Form, + state_placeholder: dolfinx.fem.Function, + hessian_u_seed: dolfinx.fem.Function, + adjoint_solution_placeholder: dolfinx.fem.Function, + *, + jit_options: dict | None, + form_compiler_options: dict | None, + entity_maps: typing.Sequence[dolfinx.mesh.EntityMap] | None, +) -> dolfinx.fem.Form: + """Build the SOA self-term ``adjoint(d2F/du2) . adjoint_solution``. + + The same computation for both ``LinearProblem`` and ``NonlinearProblem``: + ``d2F/du2`` is structurally zero for a linear residual (``dF/du`` doesn't + reference ``u``), so this compiles to a ``ufl.ZeroBaseForm`` for + ``LinearProblem`` as a *result* of running the same code, not as a special + case -- callers can always assemble the returned form unconditionally, + mirroring how an inactive TLM/fixed template is represented elsewhere in + this module (e.g. ``_get_or_build_tlm_rhs_templates``). + """ + d2Fdu2 = ufl.algorithms.expand_derivatives(ufl.derivative(dFdu_template, state_placeholder, hessian_u_seed)) + if d2Fdu2.empty(): + soa_self_form = ufl.ZeroBaseForm((dFdu_template.arguments()[0],)) + else: + soa_self_form = ufl.action(ufl.adjoint(d2Fdu2), adjoint_solution_placeholder) + return dolfinx.fem.form( + soa_self_form, + jit_options=jit_options, + form_compiler_options=form_compiler_options, + entity_maps=entity_maps, + ) + + +class _ProblemBase: + """Shared lazy adjoint/TLM solver machinery for ``LinearProblem``/``NonlinearProblem``. + + A plain mixin -- it does not inherit from any ``dolfinx.fem.petsc`` class, + so ``LinearProblem(_ProblemBase, dolfinx.fem.petsc.LinearProblem)`` (and + the ``NonlinearProblem`` equivalent) has an unambiguous MRO for + ``solve()``: each subclass keeps defining its own ``solve()``, delegating + the shared middle to ``_record_and_solve`` below via the ``_make_block``/ + ``_dolfinx_solve`` hooks it implements. + + Every attribute referenced here is set by the concrete subclass's own + ``__init__`` (either directly, or inherited from the ``dolfinx.fem.petsc`` + base it also derives from) before any of these methods run. + """ - # 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. Each block holds a plain (strong) reference back - # to this Problem (see LinearProblemBlock._problem/__init__ for why a - # strong reference is safe here and doesn't reintroduce the MPI - # collective-destruction hazard documented in - # dolfinx-adjoint-knowledge's mpi-collective-destruction-hazard note), - # so this Problem -- and hence its solvers' PETSc objects -- is - # released deterministically via ordinary refcounting once every block - # referencing it is unreachable (e.g. after clearing the tape), rather - # than being left to pyadjoint's tape/cyclic-GC schedule. 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 + ad_block_tag: str | None + bcs: typing.Sequence[dolfinx.fem.DirichletBC] + _u: _Function | typing.Sequence[_Function] + _rhs: typing.Any + _preconditioner: typing.Any + _value_placeholders: dict[dolfinx.fem.Function, dolfinx.fem.Function] + _jit_options: dict | None + _form_compiler_options: dict | None + _entity_maps: typing.Sequence[dolfinx.mesh.EntityMap] | None + _adj_options: dict | None + _tlm_options: dict | None + _petsc_options_prefix: str + _kind: typing.Any + + def _init_adjoint_state(self) -> None: + """Initialize the lazily-built adjoint/TLM solver state. + + Called once, at the end of ``__init__``, after the base + ``dolfinx.fem.petsc`` solver has been constructed. Built lazily (on + first use, see ``_get_or_build_adjoint_solver``/ + ``_get_or_build_tlm_solver``/``_get_or_build_hessian_templates``) and + shared by every block this Problem records, rather than one per + block/solve() call. Each block holds a plain (strong) reference back + to this Problem (see ``*ProblemBlock._problem``/``__init__`` for why a + strong reference is safe here and doesn't reintroduce the MPI + collective-destruction hazard documented in + dolfinx-adjoint-knowledge's mpi-collective-destruction-hazard note), + so this Problem -- and hence its solvers' PETSc objects -- is released + deterministically via ordinary refcounting once every block + referencing it is unreachable, rather than being left to pyadjoint's + tape/cyclic-GC schedule. Laziness keeps pure forward (non-annotated) + use from paying for a symbolic adjoint form it never needs. + """ + self._adjoint_solver: HomogeneousBCLinearProblem | None = None + self._tlm_solver: HomogeneousBCLinearProblem | None = None + self._residual_template: ufl.Form | None = None + self._residual_state_placeholder: dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function] | None = None + self._dFdu_template: ufl.Form | typing.Sequence | None = None + self._dFdu_adj_template: ufl.Form | typing.Sequence | None = None + self._tlm_rhs_templates: dict | None = 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 + self._hessian_templates: HessianTemplates | None = None + self._adjoint_solution_placeholder: dolfinx.fem.Function | None = None + self._second_adjoint_solution_placeholder: dolfinx.fem.Function | None = None + self._hessian_u_seed: dolfinx.fem.Function | None = None + + def _get_or_build_residual_template( + self, + ) -> tuple[ufl.Form, dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function]]: + """Build (once) and return F with every coefficient replaced by its placeholder, and u + replaced by a dedicated "state" placeholder standing in for "u at this evaluation point". + + The one genuinely irreducible difference between the two Problem kinds -- LinearProblem + builds it from ``a``/``L`` via ``ufl.action``, NonlinearProblem already has ``F`` + directly -- everything built on top of it below (``dF/du``, the TLM right-hand side, the + Hessian templates) is derived from this one template by the same shared symbolic + differentiation, since that costs nothing extra at compile time: there is no reason to + keep a separate "dF/du is just a, never differentiated" shortcut for LinearProblem. + Not implemented on the base; each subclass overrides it. + """ + raise NotImplementedError 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``). + """Build (once) and return dF/du, evaluated at the residual template's state placeholder. + + Shared by both classes: derived from ``_get_or_build_residual_template`` by symbolic + differentiation, which is free at compile time, rather than special-cased per class. + This is the shared basis for 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 + ``*ProblemBlock.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] + F_template, state_placeholder = self._get_or_build_residual_template() + if isinstance(self._u, list): + assert isinstance(state_placeholder, typing.Sequence) + test_functions = get_sorted_arguments(F_template.arguments(), 0) + state_list = list(state_placeholder) + trial_functions = [ + ufl.TrialFunction(state.function_space, part=arg.part()) + for arg, state in zip(test_functions, state_list, strict=True) + ] + dFdu = ufl.derivative(F_template, state_list, trial_functions) + else: + assert isinstance(state_placeholder, dolfinx.fem.Function) + trial_function = ufl.TrialFunction(state_placeholder.function_space) + dFdu = ufl.derivative(F_template, state_placeholder, trial_function) + self._dFdu_template = ufl.algorithms.expand_derivatives(dFdu) return self._dFdu_template def _get_or_build_dFdu_adj_template(self) -> ufl.Form | typing.Sequence: @@ -287,121 +247,52 @@ def _get_or_build_dFdu_adj_template(self) -> ufl.Form | typing.Sequence: 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. + The same computation for both classes: ``_ProblemBlockBase._compute_adjoint`` + swaps argument numbers while preserving mixed-space ``part()`` tags (via + ``compute_form_adjoint``) and decomposes the result back into blocks (via + ``ufl.extract_blocks``) -- a no-op decomposition for a scalar, non-blocked + form. Kept exactly as that returns it (a nested list of forms for a blocked + problem) since that structure is what ``HomogeneousBCLinearProblem``/ + ``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._dFdu_adj_template = _ProblemBlockBase._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]], + 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. + Shared by both classes: built purely from the residual template + (``_get_or_build_residual_template``), since ``dF/dm`` genuinely depends on the state + for either a linear or nonlinear residual -- differentiating w.r.t. a coefficient + embedded in the residual while holding ``u`` fixed leaves ``u`` in the result even when + the residual is linear in ``u`` itself. + + 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 + ``*ProblemBlock.prepare_evaluate_tlm``), avoids that entirely. """ + F_template, state_placeholder = self._get_or_build_residual_template() 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: @@ -430,75 +321,61 @@ def _get_or_build_tlm_rhs_templates( ) 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] + return self._tlm_rhs_templates, self._tlm_seed_placeholders, state_placeholder - def _get_or_build_hessian_templates(self) -> tuple[dict, dict, dict]: + def _get_or_build_hessian_templates(self) -> HessianTemplates: """Build (once) and return the per-dependency Hessian templates used by - ``LinearProblemBlock.prepare_evaluate_hessian``'s SOA right-hand side and + ``*ProblemBlock.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. + Shared by both classes -- built purely from ``_get_or_build_residual_template``/ + ``_get_or_build_dFdu_template``/``_get_or_build_dFdu_adj_template``/ + ``_get_or_build_tlm_rhs_templates``, all themselves shared. ``soa_self`` (see + ``_build_soa_self_template``) comes out a ``ufl.ZeroBaseForm`` for ``LinearProblem`` + (``dF/du`` doesn't depend on ``u``) and generally nonzero for ``NonlinearProblem``, as a + *result* of running the same code, not a per-class branch. + + Scalar (non-blocked) problems only -- the blocked Hessian cross-term stays on the + pre-templating, per-call ``ufl.replace`` + recompile path (see + ``_ProblemBlockBase._evaluate_hessian_blocked_rhs``/``_evaluate_hessian_component_blocked``). + + Each of ``soa_cross``/``fixed``/``cross`` 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() + F_template, _ = self._get_or_build_residual_template() 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] + self._adjoint_solution_placeholder = dolfinx.fem.Function(state_placeholder.function_space) + self._second_adjoint_solution_placeholder = dolfinx.fem.Function(state_placeholder.function_space) + self._hessian_u_seed = dolfinx.fem.Function(state_placeholder.function_space) - 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) + soa_self = _build_soa_self_template( + dFdu_template, + state_placeholder, + self._hessian_u_seed, + self._adjoint_solution_placeholder, + jit_options=self._jit_options, + form_compiler_options=self._form_compiler_options, + entity_maps=self._entity_maps, ) - 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 = {} + soa_cross_templates: dict = {} fixed_templates: dict = {} cross_templates: dict = {} for c, c_placeholder in self._value_placeholders.items(): @@ -506,7 +383,7 @@ def _get_or_build_hessian_templates(self) -> tuple[dict, dict, dict]: 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_cross_templates[c] = dolfinx.fem.form( soa_form, jit_options=self._jit_options, form_compiler_options=self._form_compiler_options, @@ -540,40 +417,85 @@ def _get_or_build_hessian_templates(self) -> tuple[dict, dict, dict]: form_compiler_options=self._form_compiler_options, entity_maps=self._entity_maps, ) - self._hessian_templates = (soa_templates, fixed_templates, cross_templates) + self._hessian_templates = HessianTemplates(soa_self, soa_cross_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. + def _get_or_build_adjoint_solver(self) -> HomogeneousBCLinearProblem: + """Build (once) and return the adjoint solver shared by every block this Problem records.""" + if self._adjoint_solver is None: + self._adjoint_solver = HomogeneousBCLinearProblem( + self._get_or_build_dFdu_adj_template(), # type: ignore[arg-type] + self._rhs, + bcs=self.bcs, + P=self._preconditioner, + 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, + entity_maps=self._entity_maps, + ) # type: ignore[misc] + return self._adjoint_solver + + def _get_or_build_tlm_solver(self) -> HomogeneousBCLinearProblem: + """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 ``*ProblemBlock.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. """ - annotate = pyadjoint.annotate_tape({"annotate": annotate}) - if annotate: - block = LinearProblemBlock( - self._lhs, # type: ignore[arg-type] - self._rhs, # type: ignore[arg-type] + if self._tlm_solver is None: + dFdu_template = self._get_or_build_dFdu_template() # type: ignore[attr-defined] + if isinstance(self._u, list): + dFdu_template = ufl.extract_blocks(dFdu_template) + self._tlm_solver = HomogeneousBCLinearProblem( + dFdu_template, + self._rhs, bcs=self.bcs, - u=self.u, # type: ignore[arg-type] - P=self._preconditioner, # type: ignore[arg-type] + P=self._preconditioner, 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, entity_maps=self._entity_maps, - ad_block_tag=self.ad_block_tag, - problem=self, ) # type: ignore[misc] + return self._tlm_solver + + def _record_and_solve(self, annotate: bool) -> _Function | typing.Sequence[_Function]: + """Shared ``solve()`` skeleton for both classes. + + Records a tape block (via the subclass's ``_make_block`` hook) when + annotating, refreshes the forward solver's placeholder coefficients + from the user's own current values (a prior recompute -- see + ``*ProblemBlock.prepare_recompute_component`` -- may have left them + holding a checkpointed/candidate value instead), solves (via the + subclass's ``_dolfinx_solve`` hook), and records the block's outputs. + """ + annotate = pyadjoint.annotate_tape({"annotate": annotate}) + block = self._make_block() if annotate else None # type: ignore[attr-defined] + if annotate: + assert block is not None 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) + out = self._dolfinx_solve() # type: ignore[attr-defined] if annotate: + assert block is not None if isinstance(out, Function): block.add_output(out.create_block_variable()) else: @@ -583,7 +505,7 @@ def solve(self, annotate: bool = True) -> typing.Union[dolfinx.fem.Function, typ return out -class NonlinearProblem(dolfinx.fem.petsc.NonlinearProblem): +class LinearProblem(_ProblemBase, dolfinx.fem.petsc.LinearProblem): """A linear problem that can be used with adjoint methods. This class extends the `dolfinx.fem.petsc.LinearProblem` to support adjoint methods. @@ -608,21 +530,226 @@ class NonlinearProblem(dolfinx.fem.petsc.NonlinearProblem): @typing.overload def __init__( self, - F: ufl.form.Form, - u: _Function, + a: ufl.Form, + L: ufl.Form, *, bcs: typing.Sequence[dolfinx.fem.DirichletBC] | None = None, - J: ufl.form.Form | None = None, - P: ufl.form.Form | None = None, + u: _Function | None = None, + P: ufl.Form | None = None, kind: str | None = None, petsc_options: dict | None = None, - petsc_options_prefix: str = "dxa_nonlinear_problem_", + petsc_options_prefix: str = "dxa_linear_problem_", form_compiler_options: dict | None = None, jit_options: dict | None = None, entity_maps: typing.Sequence[dolfinx.mesh.EntityMap] | None = None, - ad_block_tag: typing.Optional[str] = None, - adjoint_petsc_options: typing.Optional[dict] = None, - tlm_petsc_options: typing.Optional[dict] = None, + ad_block_tag: str | None = None, + adjoint_petsc_options: dict | None = None, + tlm_petsc_options: dict | None = None, + ) -> None: ... + @typing.overload + def __init__( + self, + a: typing.Sequence[typing.Sequence[ufl.Form]], + L: typing.Sequence[ufl.Form], + *, + 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, + petsc_options_prefix: str = "dxa_linear_problem_", + 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, + ) -> None: ... + def __init__( + self, + a: ufl.Form | typing.Sequence[typing.Sequence[ufl.Form]], + L: ufl.Form | typing.Sequence[ufl.Form], + *, + 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, + petsc_options_prefix: str = "dxa_linear_problem_", + 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, + ) -> None: + 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 + self._lhs = a + self._rhs = L + self._jit_options = jit_options + 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=_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=_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, + form_compiler_options=form_compiler_options, + jit_options=jit_options, + 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 solver state: shared lazy-init machinery + # lives in _ProblemBase._init_adjoint_state -- see its docstring for + # why laziness and Problem-owned (not block-owned) solvers matter. + self._init_adjoint_state() + + def _get_or_build_residual_template( + self, + ) -> tuple[ufl.Form, dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function]]: + """Build (once) and return F = action(a, state) - L, placeholder-substituted, with a + dedicated "state" placeholder standing in for "u at this evaluation point", distinct + from the live ``self._u`` the forward solve owns. + + The shared basis for ``dF/du`` (``_ProblemBase._get_or_build_dFdu_template``) and the + TLM right-hand side (``_get_or_build_tlm_rhs_templates``): ``dF/dm`` genuinely depends + on the state 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), and + ``dF/du`` itself is now derived from this template by the same symbolic differentiation + ``NonlinearProblem`` uses, rather than a shortcut that skips differentiating altogether. + """ + if self._residual_template 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 = ufl.replace(sum_form(self._lhs), self._value_placeholders) # type: ignore[arg-type] + L_template = ufl.replace(sum_form(self._rhs), self._value_placeholders) # type: ignore[arg-type] + self._residual_template = ufl.action(a_template, state_arg) - L_template # type: ignore[arg-type] + return self._residual_template, self._residual_state_placeholder # type: ignore[return-value] + + def _make_block(self) -> LinearProblemBlock: + return LinearProblemBlock( + self._lhs, # type: ignore[arg-type] + self._rhs, # type: ignore[arg-type] + bcs=self.bcs, + u=self.u, # type: ignore[arg-type] + P=self._preconditioner, # type: ignore[arg-type] + form_compiler_options=self._form_compiler_options, + jit_options=self._jit_options, + entity_maps=self._entity_maps, + ad_block_tag=self.ad_block_tag, + problem=self, + ) # type: ignore[misc] + + def _dolfinx_solve(self) -> _Function | typing.Sequence[_Function]: + return dolfinx.fem.petsc.LinearProblem.solve(self) + + def solve(self, annotate: bool = True) -> typing.Union[dolfinx.fem.Function, typing.Sequence[dolfinx.fem.Function]]: + """ + Solve the linear problem and return the solution. + """ + return self._record_and_solve(annotate) + + +class NonlinearProblem(_ProblemBase, dolfinx.fem.petsc.NonlinearProblem): + """A linear problem that can be used with adjoint methods. + + This class extends the `dolfinx.fem.petsc.LinearProblem` to support adjoint methods. + + Args: + a: The bilinear form representing the left-hand side of the equation. + L: The linear form representing the right-hand side of the equation. + bcs: Boundary conditions to apply to the problem. + u: Solution vector. + P: Preconditioner for the linear problem. + kind: Kind of PETSc Matrix to assemble the system into. + petsc_options: Options dictionary for the PETSc krylov supspace solver. + form_compiler_options: Form compiler options for generating assembly kernels. + jit_options: Options for just-in-time compilation of the forms. + entity_maps: Mapping from meshes that coefficients and arguments are defined on to the + integration domain of the forms. + ad_block_tag: Tag for adjoint blocks in the tape. + adjoint_petsc_options: PETSc options for adjoint problems. + tlm_petsc_options: Optional PETSc options for TLM problems. + """ + + @typing.overload + def __init__( + self, + F: ufl.form.Form, + u: _Function, + *, + bcs: typing.Sequence[dolfinx.fem.DirichletBC] | None = None, + J: ufl.form.Form | None = None, + P: ufl.form.Form | None = None, + kind: str | None = None, + petsc_options: dict | None = None, + petsc_options_prefix: str = "dxa_nonlinear_problem_", + 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, ) -> None: ... @typing.overload def __init__( @@ -639,9 +766,9 @@ def __init__( form_compiler_options: dict | None = None, jit_options: dict | None = None, entity_maps: typing.Sequence[dolfinx.mesh.EntityMap] | None = None, - ad_block_tag: typing.Optional[str] = None, - adjoint_petsc_options: typing.Optional[dict] = None, - tlm_petsc_options: typing.Optional[dict] = None, + ad_block_tag: str | None = None, + adjoint_petsc_options: dict | None = None, + tlm_petsc_options: dict | None = None, ) -> None: ... def __init__( self, @@ -657,17 +784,32 @@ def __init__( form_compiler_options: dict | None = None, jit_options: dict | None = None, entity_maps: typing.Sequence[dolfinx.mesh.EntityMap] | None = None, - ad_block_tag: typing.Optional[str] = None, - adjoint_petsc_options: typing.Optional[dict] = None, - tlm_petsc_options: typing.Optional[dict] = None, + ad_block_tag: str | None = None, + adjoint_petsc_options: dict | None = None, + tlm_petsc_options: dict | None = None, ) -> None: 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 the test functions in a blocked + # residual once, here, mirroring LinearProblem: needed so a blocked + # residual's per-block forms can be safely combined into one + # whole-system form (via sum_form) inside _get_or_build_residual_template. + if not isinstance(F, ufl.Form): + F = assign_mixed_parts(F) # type: ignore[arg-type] + self._u = resolve_u(u, F) # type: ignore[arg-type] self._bcs = [] if bcs is None else bcs - self._lhs = dolfinx.fem.forms.derivative_block(F, self._u) # type: ignore[arg-type] + # The user's own J, kept only to scan for dependency coefficients that + # might appear in a hand-supplied Jacobian but not in F itself (e.g. a + # stabilization term); the Jacobian _get_or_build_dFdu_template uses + # for adjoint/TLM/Hessian purposes is always derived symbolically from + # F, never from this. Named distinctly from dolfinx.fem.petsc.NonlinearProblem's + # own `_J` (its compiled Jacobian, set by super().__init__() below) to avoid + # colliding with it. + self._user_J = J self._rhs = F self._jit_options = jit_options self._form_compiler_options = form_compiler_options @@ -720,276 +862,81 @@ 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). + # Adjoint and tangent-linear solver state: shared lazy-init machinery + # lives in _ProblemBase._init_adjoint_state -- see LinearProblem's use + # of it, and its docstring, for the rationale (same for both classes). + self._init_adjoint_state() + + @property + def bcs(self) -> typing.Sequence[dolfinx.fem.DirichletBC]: + """Dirichlet boundary conditions applied to the residual and Jacobian. + + ``dolfinx.fem.petsc.NonlinearProblem`` has no ``bcs`` attribute of its + own (its SNES callbacks close over a fixed ``bcs`` list at + construction, see the note in ``__init__``); this property exposes + ``self._bcs`` under the same name ``LinearProblem`` uses (there, it is + the base class's own attribute), so ``_ProblemBase``'s shared methods + can read/write ``self.bcs`` uniformly across both classes. """ - 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] + return self._bcs - 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] + @bcs.setter + def bcs(self, value: typing.Sequence[dolfinx.fem.DirichletBC]) -> None: + self._bcs = value - def _get_or_build_hessian_templates( + def _get_or_build_residual_template( 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. + ) -> tuple[ufl.Form, dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function]]: + """Build (once) and return F 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", distinct from the live ``self._u`` the forward SNES path owns. + + The shared basis for ``dF/du`` (``_ProblemBase._get_or_build_dFdu_template``) and the + TLM right-hand side (``_get_or_build_tlm_rhs_templates``): refreshed from a block's own + checkpointed output before each adjoint/TLM/Hessian solve (see + ``NonlinearProblemBlock._refresh_dFdu_state``/``prepare_evaluate_tlm``), keeping this + template fixed for the life of the Problem, exactly like the non-u dependencies already + routed through ``self._value_placeholders``. """ - 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, - ) + if self._residual_template 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_list = self._residual_state_placeholder + else: + self._residual_state_placeholder = dolfinx.fem.Function(self._u.function_space) # type: ignore[union-attr] + state_list = [self._residual_state_placeholder] + replace_map: dict = dict(self._value_placeholders) + replace_map.update(zip(u_list, state_list)) - 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, - ) + if isinstance(self._rhs, ufl.Form): + self._residual_template = ufl.replace(self._rhs, replace_map) + else: + self._residual_template = sum_form([ufl.replace(Fi, replace_map) for Fi in self._rhs]) + return self._residual_template, self._residual_state_placeholder # type: ignore[return-value] + + def _make_block(self) -> NonlinearProblemBlock: + return NonlinearProblemBlock( + J=self._user_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] + form_compiler_options=self._form_compiler_options, + jit_options=self._jit_options, + entity_maps=self._entity_maps, + ad_block_tag=self.ad_block_tag, + problem=self, + ) # type: ignore[misc] - 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 _dolfinx_solve(self) -> _Function | typing.Sequence[_Function]: + return dolfinx.fem.petsc.NonlinearProblem.solve(self) def solve(self, annotate: bool = True) -> typing.Union[dolfinx.fem.Function, typing.Sequence[dolfinx.fem.Function]]: """ - Solve the linear problem and return the solution. + Solve the nonlinear problem and return the solution. """ - annotate = pyadjoint.annotate_tape({"annotate": annotate}) - if annotate: - block = NonlinearProblemBlock( - J=self._lhs, # 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] - form_compiler_options=self._form_compiler_options, - jit_options=self._jit_options, - entity_maps=self._entity_maps, - ad_block_tag=self.ad_block_tag, - 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): - block.add_output(out.create_block_variable()) - else: - for ui in out: - assert isinstance(ui, Function) - block.add_output(ui.create_block_variable()) - return out + return self._record_and_solve(annotate) diff --git a/tests/test_tlm_update.py b/tests/test_tlm_update.py index 45ee2cb..f618440 100644 --- a/tests/test_tlm_update.py +++ b/tests/test_tlm_update.py @@ -18,7 +18,7 @@ import ufl from dolfinx_adjoint import Function, assemble_scalar -from dolfinx_adjoint.solvers import LinearProblem +from dolfinx_adjoint.solvers import LinearProblem, NonlinearProblem direct_solve = { "ksp_type": "preonly", @@ -121,6 +121,110 @@ 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 _navier_stokes(mesh): + """``_viscous_stokes`` with a ``u . grad(u)`` convective term added, making the residual + genuinely nonlinear in the state -- a ``NonlinearProblem`` sibling of ``_viscous_stokes``, + exercising the blocked (multi-output) Hessian path for ``NonlinearProblem``, which used to + raise ``NotImplementedError`` unconditionally and is now the same shared code + ``LinearProblemBlock`` already used for its own blocked Hessian. + """ + el_u = basix.ufl.element("P", mesh.basix_cell(), 2, shape=(mesh.geometry.dim,)) + el_p = basix.ufl.element("P", mesh.basix_cell(), 1) + V = dolfinx.fem.functionspace(mesh, el_u) + Q = dolfinx.fem.functionspace(mesh, el_p) + Z = dolfinx.fem.functionspace(mesh, ("DG", 0)) + dx = ufl.Measure("dx", domain=mesh) + + mu = Function(Z, name="viscosity") + mu.interpolate(lambda x: 1.0 + 0.5 * np.sin(np.pi * x[0])) + + uh, ph = Function(V, name="velocity"), Function(Q, name="pressure") + v, q = ufl.TestFunction(V), ufl.TestFunction(Q) + + # A moderate, non-conservative body force: large enough that the Taylor remainders stay + # clear of round-off, small enough that the Newton solve below converges reliably (the + # convective term is quadratic in the state, so scaling it up the way _viscous_stokes + # scales its linear counterpart would make the residual far stiffer). + x = ufl.SpatialCoordinate(mesh) + f = 10.0 * ufl.as_vector((ufl.sin(ufl.pi * x[1]), ufl.cos(ufl.pi * x[0]))) + + F0 = ( + ufl.inner(mu * ufl.grad(uh), ufl.grad(v)) * dx + + ufl.inner(ufl.dot(ufl.grad(uh), uh), v) * dx + + ufl.inner(ph, ufl.div(v)) * dx + - ufl.inner(f, v) * dx + ) + F1 = ufl.inner(q, ufl.div(uh)) * 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) + zero = dolfinx.fem.Constant(mesh, np.zeros(mesh.geometry.dim, dtype=dolfinx.default_scalar_type)) + bc = dolfinx.fem.dirichletbc(zero, dofs, V) + + forward_options = { + "snes_type": "newtonls", + "snes_error_if_not_converged": True, + # A warm-started recompute (see NonlinearProblemBlock._refresh_dFdu_state / + # *ProblemBlockBase.prepare_recompute_component) starts SNES already at (or + # extremely close to) a converged point when the control hasn't moved. With + # only the default rtol, SNES keeps trying to shrink a residual that is + # already at floating-point noise, and the resulting near-singular Newton + # step can make backtracking line search report DIVERGED_LINE_SEARCH. An + # explicit absolute tolerance lets it recognize "already converged" and + # exit immediately instead. + "snes_atol": 1e-9, + "snes_rtol": 1e-9, + "snes_stol": 1e-12, + "ksp_type": "preonly", + "pc_type": "lu", + "pc_factor_mat_solver_type": "mumps", + } + problem = NonlinearProblem( + [F0, F1], + u=[uh, ph], + bcs=[bc], + # A prefix distinct from the default ("dxa_nonlinear_problem_"): PETSc's options + # database is process-global and keyed by prefix, and this fixture's snes_atol/rtol/stol + # (needed for the warm-started recompute above) must not leak into -- or collide with -- + # another NonlinearProblem elsewhere in the suite that happens to use the default prefix. + petsc_options_prefix="dxa_navier_stokes_test_", + petsc_options=forward_options, + adjoint_petsc_options=direct_solve, + tlm_petsc_options=direct_solve, + ) + problem.solve() + + # Quartic in the state and with no constant offset, for the same round-off-avoidance + # reason as _viscous_stokes's objective. + J = assemble_scalar(ufl.inner(uh, uh) ** 2 * dx) + return pyadjoint.ReducedFunctional(J, pyadjoint.Control(mu)), Z + + +def test_hessian_is_independent_of_previous_evaluation_points_navier_stokes(mesh_2D): + """``test_hessian_is_independent_of_previous_evaluation_points``'s ``NonlinearProblem`` + sibling: the same second-order Taylor test, but on a genuinely nonlinear, blocked + (multi-output) residual, exercising the blocked Hessian path + ``_ProblemBlockBase._evaluate_hessian_blocked_rhs``/``_evaluate_hessian_component_blocked`` + for ``NonlinearProblem`` for the first time -- previously this path only ever ran for + ``LinearProblem``. + """ + pyadjoint.get_working_tape().clear_tape() + Jh, Z = _navier_stokes(mesh_2D) + + m2 = Function(Z) + 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])) + + 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 _diffusive_poisson(mesh): """Scalar Poisson problem whose control ``m`` sits inside ``a`` (the diffusivity). From d3b56df0972a8c9f72248c06cdb8775bb8ff173b Mon Sep 17 00:00:00 2001 From: Joergen Schartum Dokken Date: Tue, 1 Sep 2026 10:35:13 +0000 Subject: [PATCH 2/8] Another round of hassling Claude with cleanup commands. --- src/dolfinx_adjoint/blocks/solvers.py | 1072 +++++++++++++------------ src/dolfinx_adjoint/solvers.py | 661 ++++++++++----- src/dolfinx_adjoint/typing_utils.py | 6 + src/dolfinx_adjoint/ufl_utils.py | 221 +++++ tests/test_solver_reuse.py | 159 +++- tests/test_tlm_update.py | 4 +- 6 files changed, 1400 insertions(+), 723 deletions(-) create mode 100644 src/dolfinx_adjoint/typing_utils.py create mode 100644 src/dolfinx_adjoint/ufl_utils.py diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index 4fba1ba..62924b3 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -1,6 +1,9 @@ from __future__ import annotations +import abc import typing +import warnings +import weakref from petsc4py import PETSc @@ -10,147 +13,14 @@ import ufl from dolfinx.fem.function import Function as _Function -from ..compat import compute_form_adjoint from ..types import Function +from ..typing_utils import NestedMutableSequence +from ..ufl_utils import assign_mixed_parts, collect_coefficients, sum_form 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]"] - - -@typing.overload -def assign_mixed_parts[T: NestedSequence[ufl.Form]](form1: T, /) -> T: ... -@typing.overload -def assign_mixed_parts[T: NestedSequence[ufl.Form], S: NestedSequence[ufl.Form]]( - form1: T, form2: S, / -) -> tuple[T, S]: ... -def assign_mixed_parts( - *form_structs: NestedSequence[ufl.Form], -) -> NestedSequence[ufl.Form] | tuple[NestedSequence[ufl.Form], ...]: - """ - Recursively assigns mixed-space `part` indices to {py:class}`ufl.Argument` - (test and trial functions), within nested iterables of forms. - - When solving monolithic block systems in FEniCSx, the UFL arguments must have the - method {py:meth}`ufl.Argument.part` return the index corresponding to their block position. - For a block matrix (list of lists), the TestFunction corresponds to the row index, and the - TrialFunction corresponds to the column index. - - This utility traverses arbitrary nested structures (e.g., a 2D list for the LHS - matrix `a` and a 1D list for the RHS vector `L` simultaneously), extracts arguments - that lack a part index, builds a unified replacement map, and applies it. - - Args: - *form_structs: One or more UFL forms, or nested iterables (lists/tuples) of - UFL forms. Passing multiple structures (like `a` and `L`) ensures they - share the same replacement map, preventing mismatched compilation. - - Returns: - 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 {py:func}`ufl.TestFunctions` - and {py:func}`ufl.TrialFunctions` of a single - {py:class}`ufl.MixedFunctionSpace` built from the row/column function spaces - discovered while walking the structure. - """ - 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. - - `indices` will be `(row,)` for vectors and `(row, col)` for matrices. - """ - if isinstance(obj, ufl.Form): - for arg in obj.arguments(): - if arg.part() is None: - # The argument number corresponds to the index of the row/column - # in the nested structure - num = arg.number() - if num < len(indices): - 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,)) - else: - raise TypeError(f"Expected ufl.Form or iterable, got {type(obj)}") - - 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: - _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, - strictly preserving original sequence types (lists vs. tuples). - """ - if isinstance(obj, ufl.Form): - return ufl.replace(obj, replace_map) - elif isinstance(obj, (list, tuple)): - return type(obj)(_replace(item) for item in obj) - return obj - - # 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] - - -def get_sorted_arguments(arguments: typing.Iterable[ufl.Argument], number: int) -> typing.Iterable[ufl.Argument]: - """Extract all arguments of a given number, sorted by part.""" - return sorted(filter(lambda x: x.number() == number, arguments), key=lambda a: a.part()) - - -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 a blocked ``NonlinearProblem``'s forms are not - required to be before ``assign_mixed_parts`` runs. - """ - 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 _map_block_variables_to_form( form: ufl.Form | NestedMutableSequence[ufl.Form] | None, @@ -178,54 +48,32 @@ def _map_block_variables_to_form( return replace_map -def sum_form(form: NestedSequence[ufl.Form | None]) -> ufl.Form | None: - """Sum a blocked form into a single form.""" - # Handle top-level None - if form is None: - return None - - if isinstance(form, ufl.Form): - return form - - elif isinstance(form, typing.Iterable): - # Recursively sum items, filtering out Nones - valid_forms: list[ufl.Form] = [] - for fi in form: - summed_fi = sum_form(fi) - if summed_fi is not None: - valid_forms.append(summed_fi) - - # Handle empty case safely - if not valid_forms: - return None - - # Safely sum without defaulting to integer 0, removing the need for type: ignore - return sum(valid_forms[1:], start=valid_forms[0]) - - else: - raise TypeError(f"Cannot sum form of type {type(form)}") - - -class _ProblemBlockBase(pyadjoint.Block): - """Shared tape-block machinery for ``LinearProblemBlock``/``NonlinearProblemBlock``. +class _ProblemBlockBase(pyadjoint.Block, abc.ABC): + """Shared tape-block machinery for + {py:class}`~dolfinx_adjoint.blocks.solvers.LinearProblemBlock`/{py:class}`~dolfinx_adjoint.blocks.solvers.NonlinearProblemBlock`. Holds what is unconditionally identical or shares one implementation between the two Problem kinds: fetching the owning Problem, recovering Dirichlet BC dependencies, detecting a boundary-condition-only adjoint, transposing a (possibly blocked) bilinear form, the shared warm-started - recompute flow, and -- since ``LinearProblem._compute_residual`` and - ``NonlinearProblem._compute_residual`` both settle on the same output - shape (a single summed ``ufl.Form`` plus its dependency replacement map, - see each subclass's docstring) -- every first-order adjoint, TLM and + recompute flow, and -- since + {py:meth}`LinearProblemBlock._compute_residual` + and + {py:meth}`NonlinearProblemBlock._compute_residual` + both settle on the same output shape (a single summed {py:class}`ufl.Form` + plus its dependency replacement map, see each subclass's docstring) -- + every first-order adjoint, TLM and Hessian method built *on top of* that residual, both scalar and blocked. Each concrete subclass still implements its own ``__init__`` (constructor - kwargs differ: ``a``/``L`` vs ``J``/``F``) and ``_compute_residual`` + kwargs differ: ``a``/``L`` vs ``J``/``F``) and + {py:meth}`~dolfinx_adjoint.blocks.solvers._ProblemBlockBase._compute_residual` itself, since building the residual is the one place a shared algorithm isn't possible -- the two classes start from different user-supplied data. """ - _problem_obj: typing.Any + _problem_ref: weakref.ReferenceType["LinearProblem | NonlinearProblem"] + _rebuilt_problem: "LinearProblem | NonlinearProblem | None" = None _bcs: typing.Sequence[dolfinx.fem.DirichletBC] _u: _Function | typing.Sequence[_Function] _adjoint_solutions: dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function] @@ -235,113 +83,140 @@ class _ProblemBlockBase(pyadjoint.Block): _form_compiler_options: dict | None _entity_maps: typing.Sequence[dolfinx.mesh.EntityMap] | None - @property - def _problem(self) -> typing.Any: - """Return this block's owning Problem, which owns the shared solvers.""" - return self._problem_obj + def get_reference_problem(self) -> "LinearProblem | NonlinearProblem": + """Return this block's owning Problem, which owns the shared solvers. + + Held via a weakref, not a strong reference: a throwaway Problem (solve it, then + only touch the resulting {py:class}`~pyadjoint.ReducedFunctional`) must not be + kept alive for as long as this Block stays reachable from the tape. If the + Problem has already been collected, rebuild an equivalent one from the + fields this Block stored at construction time (see + {py:meth}`~dolfinx_adjoint.blocks.solvers._ProblemBlockBase._rebuild_problem`) + and keep a strong reference to the rebuilt Problem on this Block from then + on -- this is a rare, costly fallback + (a fresh Problem means fresh, uncompiled forms and solvers), not the common + path, so it warns; the strong reference amortizes that cost across every later + call this Block makes rather than rebuilding again on every call. + + Returns: + This block's owning Problem (rebuilt, if necessary). + """ + problem = self._problem_ref() + if problem is None: + problem = self._rebuild_problem() + self._problem_ref = weakref.ref(problem) + self._rebuilt_problem = problem + return problem + + @abc.abstractmethod + def _rebuild_problem(self) -> "LinearProblem | NonlinearProblem": + """Reconstruct an equivalent Problem from the fields this Block stored at + construction time, after the original Problem has been garbage collected. + + Each subclass overrides this to call its own Problem constructor + ({py:class}`~dolfinx_adjoint.LinearProblem`/{py:class}`~dolfinx_adjoint.NonlinearProblem` + take different arguments -- ``a``/``L`` vs ``F``/``J``). + + Returns: + A freshly constructed Problem, equivalent to the one this Block was + originally given. + """ + @abc.abstractmethod def _compute_residual(self) -> tuple[ufl.Form, dict[Function, Function]]: """Build this block's residual ``F(u, v) = 0`` at its checkpointed dependency values. - The one genuinely irreducible difference between the two Problem kinds: ``LinearProblem`` - derives it from ``a``/``L`` via ``ufl.action``, ``NonlinearProblem`` already has ``F`` - directly. Both settle on the same output shape -- a single summed ``ufl.Form`` plus its - dependency replacement map -- which is what lets every method built on top of this one - (below) be shared. Not implemented on the base; each subclass overrides it. + The one genuinely irreducible difference between the two Problem kinds: + {py:class}`~dolfinx_adjoint.LinearProblem` derives it from ``a``/``L`` via + {py:func}`ufl.action`, {py:class}`~dolfinx_adjoint.NonlinearProblem` + already has ``F`` directly. Both settle on the same output shape -- a + single summed {py:class}`ufl.Form` plus its dependency replacement map -- + which is what lets every method built on top of this one (below) be + shared. Each subclass overrides this. + + Returns: + A ``(F_form, replacement_map)`` pair: a single summed {py:class}`ufl.Form` for the + residual, and the dependency-to-checkpoint replacement map used to build it. """ - raise NotImplementedError - def _refresh_dFdu_state(self, problem: typing.Any) -> None: + def _refresh_dFdu_state(self, problem: "LinearProblem | NonlinearProblem") -> None: """Refresh whichever coefficient stands in for "the state" in ``dF/du``, if any. - A no-op by default: ``LinearProblem``'s ``dF/du`` (``a``) is bilinear and never - references the state at all, so there is nothing to refresh before using the shared - adjoint solver. ``NonlinearProblemBlock`` overrides this, since ``dF/du`` genuinely - depends on ``u``'s current value there -- and, unlike the TLM path (which already loops - over ``problem._get_or_build_tlm_rhs_templates()``'s state placeholder(s) generically), - neither ``prepare_evaluate_adj`` nor ``prepare_evaluate_hessian`` otherwise touches it. + A no-op by default: {py:class}`~dolfinx_adjoint.LinearProblem`'s ``dF/du`` + (``a``) is bilinear and never references the state at all, so there is + nothing to refresh before using the shared adjoint solver. + {py:class}`~dolfinx_adjoint.blocks.solvers.NonlinearProblemBlock` + overrides this, since ``dF/du`` genuinely depends on ``u``'s current + value there -- and, unlike the TLM path (which already loops over + ``problem.residual_state_placeholder`` generically), neither + {py:meth}`~dolfinx_adjoint.blocks.solvers._ProblemBlockBase.prepare_evaluate_adj` + nor + {py:meth}`~dolfinx_adjoint.blocks.solvers._ProblemBlockBase.prepare_evaluate_hessian` + otherwise touches it. + + Args: + problem: This block's owning Problem (see ``self.get_reference_problem()``). """ - - def _recover_bcs(self): - bcs = [] - for block_variable in self.get_dependencies(): - c = block_variable.output - c_rep = block_variable.saved_output - - if isinstance(c, dolfinx.fem.DirichletBC): - bcs.append(c_rep) - return bcs - - def _should_compute_boundary_adjoint( - self, relevant_dependencies: typing.List[tuple[int, pyadjoint.block_variable.BlockVariable]] - ) -> bool: - """Determine if the adjoint should be computed with respect to the boundary conditions.""" - bdy = False - for _, dep in relevant_dependencies: - if isinstance(dep.output, dolfinx.fem.DirichletBC): - bdy = True - break - return bdy - - @classmethod - def _compute_adjoint(cls, form: ufl.Form) -> typing.Sequence[typing.Sequence[ufl.Form]] | ufl.Form: - """ - Compute adjoint of a bilinear form :math:`a(u, v)`, which could be written as a blocked system. - """ - return ufl.extract_blocks(compute_form_adjoint(form)) + pass def _create_replace_map(self, form: ufl.Form | NestedMutableSequence[ufl.Form] | None) -> dict[Function, Function]: - """Replace dependencies with latest checkpoint.""" + """Map each dependency and output to its checkpointed value, wherever it appears in ``form``. + + Args: + form: This block's residual (or nested block structure thereof). + + Returns: + A dict from each coefficient (dependency or output) appearing in ``form`` + to the {py:class}`~dolfinx_adjoint.Function` it should be replaced by (its + checkpointed value). + """ replace_map: dict = {} replace_map.update(_map_block_variables_to_form(form, self.get_dependencies())) replace_map.update(_map_block_variables_to_form(form, self.get_outputs())) return replace_map - def _compute_residual_derivative(self) -> ufl.Form | list[list[ufl.Form]]: - """Compute the derivative of the residual with respect to the outputs. - - Shared by both Problem kinds: built purely from ``self._compute_residual()``'s output - (a single summed form, the same shape for both classes), so no per-class override is - needed even though the two classes construct that residual very differently. - """ - F_form, _ = self._compute_residual() - - outputs = self.get_outputs() - # Use r.saved_output directly; no lookup in replacement_map needed! - r_funcs = [r.saved_output for r in outputs] - - test_functions = get_sorted_arguments(F_form.arguments(), 0) - trial_functions = [ - ufl.TrialFunction(output.function_space, part=arg.part()) - for arg, output in zip(test_functions, r_funcs, strict=True) - ] - - dFdu = ufl.derivative(F_form, r_funcs, trial_functions) - - if isinstance(self._u, list): - return ufl.extract_blocks(dFdu) - return dFdu - def prepare_evaluate_tlm( self, inputs, tlm_inputs, relevant_outputs ) -> typing.Sequence[Function] | dolfinx.fem.Function: - - # The TLM solver -- and the compiled LHS it solves with, shared - # verbatim with dF/du (see *Problem._get_or_build_dFdu_template) -- - # are shared across every block this Problem records; likewise the - # per-dependency TLM right-hand-side templates (see - # *Problem._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 + """Assemble and solve the tangent-linear (TLM) system for this block. + + The TLM solver -- and the compiled LHS it solves with, shared verbatim with + dF/du (see + {py:meth}`*Problem._get_or_build_dFdu_template`) + -- are shared across every block this Problem records; likewise the + per-dependency TLM right-hand-side templates (see + {py:meth}`*Problem._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. Only dependencies that actually have a + tangent-linear value this call are assembled into the right-hand side (see + {py:meth}`*Problem._get_or_build_tlm_rhs_templates` + for why an inactive dependency's term must be skipped entirely rather than + evaluated with a zeroed direction). + + Args: + inputs: The dependencies' current values. Unused: each dependency is + visited via ``self.get_dependencies()`` instead. + tlm_inputs: The dependencies' tangent-linear values, parallel to + ``inputs``. Unused for the same reason. + relevant_outputs: The output block variables relevant to + {py:meth}`~dolfinx_adjoint.blocks.solvers._ProblemBlockBase.evaluate_tlm_component`. + Unused: this block always solves for every output at once (see + ``self.get_outputs()``). + + Returns: + ``self._tlm_solutions``: this block's own tangent-linear solution(s), + already solved for. Passed through unchanged as ``prepared`` to every + subsequent {py:meth}`~dolfinx_adjoint.blocks.solvers._ProblemBlockBase.evaluate_tlm_component` call. + """ + problem = self.get_reference_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) + 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() @@ -388,8 +263,22 @@ def prepare_evaluate_tlm( return self._tlm_solutions 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. + """Return this output's share of the tangent-linear solution already computed + by {py:meth}`~dolfinx_adjoint.blocks.solvers._ProblemBlockBase.prepare_evaluate_tlm`. + + Args: + inputs: The dependencies' current values. Unused. + tlm_inputs: The dependencies' tangent-linear values. Unused. + block_variable: The output block variable corresponding to ``idx``. + Unused: the result is read from ``prepared`` instead. + idx: Index of the output to return, into ``self._tlm_solutions`` if it is + a list (a blocked problem); ignored for a scalar problem. + prepared: ``self._tlm_solutions``, as returned by + {py:meth}`~dolfinx_adjoint.blocks.solvers._ProblemBlockBase.prepare_evaluate_tlm`. + + Returns: + This output's tangent-linear solution. + """ if isinstance(self._tlm_solutions, list): return self._tlm_solutions[idx] else: @@ -402,21 +291,40 @@ def prepare_evaluate_adj( adj_inputs: typing.Sequence[dolfinx.la.Vector], relevant_dependencies: typing.List[tuple[int, pyadjoint.block_variable.BlockVariable]], ) -> 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. Refresh this - # block's own checkpointed values into the placeholders, refresh - # whatever "state" dF/du is evaluated at if it depends on one (a - # no-op for LinearProblem, see _refresh_dFdu_state), 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 + """Assemble and solve the first-order adjoint equation for this block. + + The adjoint solver -- and the compiled LHS it solves with -- are shared + across every block this Problem records. Refresh this block's own + checkpointed values into the placeholders, refresh whatever "state" dF/du is + evaluated at if it depends on one (a no-op for + {py:class}`~dolfinx_adjoint.LinearProblem`, see + {py:meth}`~dolfinx_adjoint.blocks.solvers._ProblemBlockBase._refresh_dFdu_state`), + 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. + + Args: + inputs: The dependencies' current values. Unused: each dependency's + checkpointed value is read directly from ``self.get_dependencies()``. + adj_inputs: The adjoint seed(s) received from this block's output(s) -- + one entry per output for a blocked problem -- assembled into the + adjoint equation's right-hand side. + relevant_dependencies: The dependency block variables relevant to + {py:meth}`~dolfinx_adjoint.blocks.solvers._ProblemBlockBase.evaluate_adj_component`. + Unused: every dependency is visited via ``self.get_dependencies()`` + instead. + + Returns: + A ``(F_form, replacement_map)`` pair -- this block's residual (from + ``self._compute_residual()``) and its dependency replacement map -- + passed through unchanged as ``prepared`` to every subsequent + {py:meth}`~dolfinx_adjoint.blocks.solvers._ProblemBlockBase.evaluate_adj_component` call. + """ + problem = self.get_reference_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) + 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() @@ -468,8 +376,25 @@ def evaluate_adj_component( idx: int, prepared: tuple[ufl.Form, dict[Function, Function]], ) -> _SpecialVector: - """Evaluate the adjoint component, i.e. :math:`\\frac{\\partial F}{\\partial m}`.""" - + r"""Return this dependency's contribution to the adjoint action, + :math:`\left(\partial F/\partial m\right)^{*}\lambda`. + + Args: + inputs: The dependencies' current values. Unused. + adj_inputs: The adjoint seed(s) already consumed by + {py:meth}`~dolfinx_adjoint.blocks.solvers._ProblemBlockBase.prepare_evaluate_adj`. Unused here. + block_variable: The dependency block variable corresponding to ``idx``. + idx: Index of the dependency to compute the contribution for, used to + select the correct mixed-space {py:meth}`ufl.Argument.part` for a blocked problem's + trial function. + prepared: The ``(F_form, replacement_map)`` pair returned by + {py:meth}`~dolfinx_adjoint.blocks.solvers._ProblemBlockBase.prepare_evaluate_adj`. + + Returns: + The assembled sensitivity vector for this dependency, where + :math:`\lambda` is the first-order adjoint solution + (``self._adjoint_solutions``). + """ residual, replacement_map = prepared c = block_variable.output c_rep = block_variable.saved_output @@ -508,10 +433,10 @@ def prepare_recompute_component( ) -> _Function | typing.Sequence[_Function]: """Prepare for recomputing the block with different control inputs, and solve. - The forward solver (``self._problem``) is bound, forever, to + The forward solver (``self.get_reference_problem()``) is bound, forever, to compiled forms referencing dedicated placeholder coefficients rather - than the user's own dependency objects (see ``LinearProblem``/ - ``NonlinearProblem._value_placeholders``): writing this call's + than the user's own dependency objects (see + {py:class}`~dolfinx_adjoint.solvers._ProblemBase`'s ``value_placeholders``): 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 @@ -519,7 +444,7 @@ def prepare_recompute_component( The Problem's own unknown(s) are warm-started from this block's own saved outputs rather than zeroed: required for SNES/Newton - convergence, and applied to a KSP-based ``LinearProblem`` the same + convergence, and applied to a KSP-based {py:class}`~dolfinx_adjoint.LinearProblem` the same way, so an iterative solver configured with a nonzero initial guess benefits from it too -- both classes now solve by refreshing placeholder/unknown values and calling an unchanging, already-built @@ -527,33 +452,48 @@ def prepare_recompute_component( form. Solving happens once, here -- not once per output in - ``recompute_component`` -- which matters for a multi-output (blocked) + {py:meth}`~dolfinx_adjoint.blocks.solvers._ProblemBlockBase.recompute_component` + -- which matters for a multi-output (blocked) problem. The base ``dolfinx`` ``solve()`` is called directly (via ``problem._dolfinx_solve()``), not ``problem.solve()``, which would record another block onto the tape. + + Args: + inputs: The dependencies' current (candidate/checkpointed) values. + Unused: each dependency is visited via ``self.get_dependencies()`` + instead. + relevant_outputs: The ``(idx, block_variable)`` pairs identifying which + of this block's own outputs the Problem's unknown(s) should be + warm-started from. + + Returns: + The Problem's own (just-recomputed) unknown(s), ``problem.u`` -- shared + storage, not yet isolated per block; + {py:meth}`~dolfinx_adjoint.blocks.solvers._ProblemBlockBase.recompute_component` + below copies out this block's own share. """ - problem = self._problem + problem = self.get_reference_problem() for block_variable in self.get_dependencies(): - placeholder = problem._value_placeholders.get(block_variable.output) + 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 + # Problem itself -- see _problem), since another block may have # used it with different bcs in between. problem.bcs = self._bcs # Warm-start the Problem's own unknown(s) from this block's own saved # outputs. - u_list = problem._u if isinstance(problem._u, list) else [problem._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() with pyadjoint.stop_annotating(): problem._dolfinx_solve() - return problem._u + return problem.u def recompute_component( self, @@ -562,7 +502,23 @@ def recompute_component( idx: int, prepared: _Function | typing.Sequence[_Function], ) -> Function: - """Return an isolated copy of this block's own share of the already-recomputed state.""" + """Return an isolated copy of this block's own share of the already-recomputed state. + + Args: + inputs: The dependencies' current values. Unused. + block_variable: The output block variable corresponding to ``idx``. + Unused: the result is read from ``prepared`` instead. + idx: Index of the output to return, into ``prepared`` if it is a + sequence (a blocked problem); ignored for a scalar problem. + prepared: The Problem's own unknown(s), as returned by + {py:meth}`~dolfinx_adjoint.blocks.solvers._ProblemBlockBase.prepare_recompute_component` + -- shared storage the Problem itself still owns. + + Returns: + An isolated copy of this output, so this tape block's own checkpoint + stays stable even if the shared Problem's unknown is later overwritten + by another block's recompute. + """ if isinstance(prepared, dolfinx.fem.Function): assert idx == 0 # Return an explicit copy so each tape block gets an isolated state snapshot @@ -574,17 +530,37 @@ def recompute_component( def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_dependencies): """Assemble and solve the second-order-adjoint (SOA) equation. - Scalar (single-output) problems share one implementation for both - Problem kinds, built entirely from the Problem's cached Hessian/TLM - templates (see ``HessianTemplates``) -- no block-specific residual - construction needed, since the SOA self-term is already correctly - zero/nonzero per class (see ``_build_soa_self_template``). A blocked - (multi-output) problem has no templated fast path; its right-hand - side is assembled by ``_evaluate_hessian_blocked_rhs``, also shared - by both Problem kinds, built from ``_compute_residual_derivative``/ - ``_compute_adjoint`` -- themselves built only from the per-class - ``_compute_residual`` hook (the two Problem kinds start from - different user-supplied data there). + Shared by both Problem kinds and both scalar and blocked problems, built + entirely from the Problem's cached Hessian/TLM templates (see + {py:class}`~dolfinx_adjoint.solvers.HessianTemplates`) -- no + block-specific residual construction needed, since the SOA self-term is + already correctly zero/nonzero per class (see + {py:func}`~dolfinx_adjoint.solvers._build_soa_self_template`) and, for a + blocked problem, already split into one compiled form per output row (see + {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_hessian_templates`). + + Args: + inputs: The dependencies' current values. Unused. + hessian_inputs: The Hessian seed(s) received from this block's output(s) + -- one entry per output for a blocked problem -- added into the + second-order-adjoint equation's right-hand side. + adj_inputs: The first-order adjoint seed(s). Unused: the first-order + adjoint solution needed here was already solved for and cached in + ``self._adjoint_solutions`` by + {py:meth}`~dolfinx_adjoint.blocks.solvers._ProblemBlockBase.prepare_evaluate_adj`. + relevant_dependencies: The dependency block variables relevant to + {py:meth}`~dolfinx_adjoint.blocks.solvers._ProblemBlockBase.evaluate_hessian_component`. + Unused: every dependency is visited via ``self.get_dependencies()`` + instead. + + Returns: + A ``(residual, adjoint_solution, second_adjoint_solution)`` tuple -- this + block's residual (from ``self._compute_residual()``) and both adjoint + solutions -- passed through unchanged as ``prepared`` to every + subsequent + {py:meth}`~dolfinx_adjoint.blocks.solvers._ProblemBlockBase.evaluate_hessian_component` + call. ``None`` if there is nothing to do (no Hessian input, or no + dependency has a tangent-linear value). """ outputs = self.get_outputs() tlm_output = [output.tlm_value for output in outputs if output is not None] @@ -598,31 +574,48 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ # 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 + problem = self.get_reference_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) + 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() - self._refresh_dFdu_state(problem) - if len(outputs) == 1: - # Use the cached per-dependency Hessian templates (see - # *Problem._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() - hessian_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] + # Use the cached per-dependency Hessian templates (see + # *Problem._get_or_build_hessian_templates) instead of rebuilding and + # recompiling the SOA right-hand side symbolically on every call. + # Refreshing the state placeholder here subsumes _refresh_dFdu_state's + # no-op-for-Linear/state-only-for-Nonlinear distinction: the Hessian's + # own F_template/L1 depend on the state's value for *both* classes + # (Linear's F_template is action(a, state) - L, genuinely state-valued), + # unlike dF/du itself, which only Nonlinear's depends on. + _, seed_placeholders, state_placeholder = problem._get_or_build_tlm_rhs_templates() + hessian_templates = problem._get_or_build_hessian_templates() + + state_list = state_placeholder if isinstance(state_placeholder, list) else [state_placeholder] + adj_sol_placeholders = problem.adjoint_solution_placeholder + adj_sol_placeholder_list = ( + adj_sol_placeholders if isinstance(adj_sol_placeholders, list) else [adj_sol_placeholders] + ) + adjoint_solutions_list = ( + self._adjoint_solutions if isinstance(self._adjoint_solutions, list) else [self._adjoint_solutions] + ) + hessian_u_seed = problem.hessian_u_seed + hessian_u_seed_list = hessian_u_seed if isinstance(hessian_u_seed, list) else [hessian_u_seed] + + for placeholder, out_bv in zip(state_list, outputs, strict=True): + placeholder.x.array[:] = out_bv.saved_output.x.array[:] + placeholder.x.scatter_forward() + for placeholder, adj_sol in zip(adj_sol_placeholder_list, adjoint_solutions_list, strict=True): + placeholder.x.array[:] = adj_sol.x.array[:] + placeholder.x.scatter_forward() + for placeholder, tlm_val in zip(hessian_u_seed_list, tlm_output, strict=True): + placeholder.x.array[:] = tlm_val.x.array[:] + placeholder.x.scatter_forward() + if len(outputs) == 1: b = adjoint_solver.b with b.localForm() as b_loc: b_loc.set(0.0) @@ -648,7 +641,44 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ b_loc.array[:] += hessian_inputs[0].array[:] b.ghostUpdate(PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) else: - self._evaluate_hessian_blocked_rhs(adjoint_solver, hessian_inputs, tlm_output, relevant_dependencies) + assert isinstance(hessian_templates.soa_self, list) + bs = [] + for i, soa_self_i in enumerate(hessian_templates.soa_self): + out_i = outputs[i].saved_output + bi = dolfinx.la.vector(out_i.function_space.dofmap.index_map, out_i.function_space.dofmap.index_map_bs) + bi.array[:] = 0.0 + dolfinx.fem.assemble_vector(bi.array, soa_self_i) + bs.append(bi) + + 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.") + templates = hessian_templates.soa_cross.get(c) + if templates is None: + continue + seed = seed_placeholders[c] + seed.x.array[:] = tlm_input.x.array[:] + seed.x.scatter_forward() + for bi, template_i in zip(bs, templates, strict=True): + dolfinx.fem.assemble_vector(bi.array, template_i) + + for i, bi in enumerate(bs): + bi.scatter_reverse(dolfinx.la.InsertMode.add) + bi.scatter_forward() + bi.array[:] *= -1 + hess_input = hessian_inputs[i] + if hess_input is not None: + bi.array[:] += hess_input.array + bi.scatter_forward() + + 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) # The SOA (second-order-adjoint) equation shares its LHS verbatim with # the first-order adjoint equation (both are adjoint(dF/du)) -- @@ -660,83 +690,24 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ adj_sol.x.array[:] = sol.x.array[:] else: 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 - - def _evaluate_hessian_blocked_rhs(self, adjoint_solver, hessian_inputs, tlm_output, relevant_dependencies): - """Assemble the SOA right-hand side for a blocked (multi-output) problem into ``adjoint_solver.b``. - - No templated fast path exists for a blocked problem (see - ``*Problem._get_or_build_hessian_templates``): the right-hand side is built symbolically, - per call, from ``_compute_residual_derivative``/``_compute_adjoint`` -- the same two - hooks the scalar path's templates are themselves built from -- so this one implementation - is shared by both Problem kinds. ``d2Fdu2`` (the SOA self-term) is always included - unconditionally rather than asserted zero: it comes out zero for a linear residual as a - result of the maths, exactly like ``_build_soa_self_template``'s scalar-path self-term -- - and, mirroring that helper, is reduced from bilinear (test and trial) to linear (test - only) by adjointing and acting on the first-order adjoint solution before use, since - ``ufl.derivative`` w.r.t. a *coefficient* (``unknowns``) leaves ``dF/du``'s own trial - argument untouched. - """ - 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)) - - if d2Fdu2.empty(): - b_form = d2Fdu2 - else: - b_form = 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) - - bs = [] - b_form = ufl.extract_blocks(b_form) - for i, hess_input in enumerate(hessian_inputs): - if hess_input is not None: - bi = dolfinx.la.vector(hess_input.index_map, hess_input.block_size) - else: - out_i = self.get_outputs()[i].saved_output - bi = dolfinx.la.vector(out_i.function_space.dofmap.index_map, out_i.function_space.dofmap.index_map_bs) - bs.append(bi) - bi.array[:] = 0.0 - form_i = ufl.algorithms.apply_derivatives.apply_derivatives(b_form[i]) - 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.assemble_vector(bi.array, compiled_soa_rhs) - bi.scatter_reverse(dolfinx.la.InsertMode.add) - bi.scatter_forward() - bi.array[:] *= -1 - if hess_input is not None: - bi.array[:] += hess_input.array + # fixed/cross (used by evaluate_hessian_component below) depend on + # problem.second_adjoint_solution_placeholder's current value for + # both scalar and blocked problems -- refresh it here regardless. + second_adjoint_solutions_list = ( + self._second_adjoint_solutions + if isinstance(self._second_adjoint_solutions, list) + else [self._second_adjoint_solutions] + ) + second_adj_placeholders = problem.second_adjoint_solution_placeholder + second_adj_placeholder_list = ( + second_adj_placeholders if isinstance(second_adj_placeholders, list) else [second_adj_placeholders] + ) + for placeholder, sol in zip(second_adj_placeholder_list, second_adjoint_solutions_list, strict=True): + placeholder.x.array[:] = sol.x.array[:] + placeholder.x.scatter_forward() - bi.scatter_forward() - 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) + return self._compute_residual(), self._adjoint_solutions, self._second_adjoint_solutions def evaluate_hessian_component( self, @@ -750,15 +721,33 @@ def evaluate_hessian_component( ): """Return this dependency's contribution to the Hessian action. - Scalar (single-output) problems share one implementation for both - Problem kinds (see ``prepare_evaluate_hessian``); a blocked - (multi-output) problem falls back to - ``_evaluate_hessian_component_blocked``. + Shared by both Problem kinds and both scalar and blocked problems, built + entirely from the Problem's cached Hessian templates -- ``fixed``/``cross`` + have the same per-dependency (not per-output-row) shape regardless of + blocking, since they live on the *control's* own test space, not the + (possibly blocked) state's (see {py:class}`~dolfinx_adjoint.solvers.HessianTemplates`). + + Args: + inputs: The dependencies' current values. Unused. + hessian_inputs: The Hessian seed(s) already consumed by + {py:meth}`~dolfinx_adjoint.blocks.solvers._ProblemBlockBase.prepare_evaluate_hessian`. Unused here. + adj_inputs: The first-order adjoint seed(s). Unused. + block_variable: The dependency block variable corresponding to ``idx``. + idx: Index of the dependency to compute the contribution for. Unused + directly: ``block_variable.output``/``.saved_output`` identify the + dependency instead. + relevant_dependencies: The ``(idx, block_variable)`` pairs for every + other dependency that may contribute a cross term. + prepared: The ``(residual, adjoint_solution, second_adjoint_solution)`` + tuple returned by + {py:meth}`~dolfinx_adjoint.blocks.solvers._ProblemBlockBase.prepare_evaluate_hessian`. + Unused here: every quantity it carries was already used to + refresh the cached templates' placeholders there. + + Returns: + The assembled Hessian-action contribution for this dependency. """ c = block_variable.output - residual_prepared, adj_sol, adj_sol2 = prepared - outputs = self.get_outputs() - tlm_output = [output.tlm_value for output in outputs] c_rep = block_variable.saved_output # If m = DirichletBC then d^2F(u,m)/dm^2 = 0 and d^2F(u,m)/dudm = 0, @@ -777,115 +766,40 @@ def evaluate_hessian_component( assert isinstance(c, dolfinx.fem.Function) W = c.function_space - if len(outputs) == 1: - # Use the cached per-dependency 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 - # *Problem._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 - hessian_templates = problem._get_or_build_hessian_templates() - _, seed_placeholders, _ = problem._get_or_build_tlm_rhs_templates() - - fixed_template = hessian_templates.fixed[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 = hessian_templates.cross.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 - - return self._evaluate_hessian_component_blocked( - residual_prepared, adj_sol, adj_sol2, c, c_rep, W, tlm_output, relevant_dependencies - ) - - def _evaluate_hessian_component_blocked( - self, residual_prepared, adj_sol, adj_sol2, c, c_rep, W, tlm_output, relevant_dependencies - ): - """Return this dependency's Hessian-action contribution for a blocked (multi-output) problem. - - No templated fast path exists for a blocked problem; this builds it symbolically, per - call, from the residual ``prepare_evaluate_hessian`` already recorded - (``residual_prepared``) -- shared by both Problem kinds, since both settle on the same - ``_compute_residual`` output shape. - - 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 - mixed. Therefore, we instead we compute it as dL/dm = d(lambda_i^T F(m))/dm, - which is equivalent. - """ - F_form, replacement_map = residual_prepared - outputs = self.get_outputs() - - F_summed = sum_form(F_form) - L1 = ufl.action(F_summed, adj_sol) - L2 = ufl.action(F_summed, adj_sol2) - - # Compute first derivatives (1-forms tested exactly against the single 'dc' object) - dc = ufl.TestFunction(W) - assert c_rep in replacement_map.values() - dL1dm = ufl.derivative(L1, c_rep, dc) - dL2dm = ufl.derivative(L2, c_rep, dc) - - sa = [output.saved_output for output in outputs] - d2Fdudm = ufl.algorithms.expand_derivatives(ufl.derivative(dL1dm, sa, tlm_output)) + # Use the cached per-dependency 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 *Problem._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.get_reference_problem() + hessian_templates = problem._get_or_build_hessian_templates() + _, seed_placeholders, _ = problem._get_or_build_tlm_rhs_templates() + + fixed_template = hessian_templates.fixed[c] + hessian_output = _create_vector(fixed_template, W) + hessian_output.array[:] = 0.0 + assemble_compiled_form(fixed_template, hessian_output) - d2Fdm2 = ufl.ZeroBaseForm((dc,)) # Initialize the second derivative form - # 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 + template = hessian_templates.cross.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) - if isinstance(c2_rep, dolfinx.mesh.Mesh): - X = ufl.SpatialCoordinate(c2_rep) - d2Fdm2 += ufl.algorithms.expand_derivatives(ufl.derivative(dL1dm, X, tlm_input)) - else: - d2Fdm2 += ufl.algorithms.expand_derivatives(ufl.derivative(dL1dm, c2_rep, tlm_input)) - - hessian_form = ufl.algorithms.expand_derivatives(d2Fdm2 + dL2dm + 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, - ) - - test_functions = get_sorted_arguments(hessian_form.arguments(), 0) - assert len(test_functions) == 1 - hessian_output = _create_vector(compiled_hessian, test_functions[0].ufl_function_space()) - - hessian_output.array[:] = 0.0 - assemble_compiled_form(compiled_hessian, hessian_output) hessian_output.array[:] *= -1.0 return hessian_output @@ -913,6 +827,11 @@ def __init__( form_compiler_options: dict | None = None, jit_options: dict | None = None, entity_maps: typing.Sequence[dolfinx.mesh.EntityMap] | None = None, + kind: typing.Any = None, + petsc_options: dict | None = None, + petsc_options_prefix: str | None = None, + adjoint_petsc_options: dict | None = None, + tlm_petsc_options: dict | None = None, ad_block_tag: str | None = None, problem: "LinearProblem" = ..., ) -> None: ... @@ -929,6 +848,11 @@ def __init__( form_compiler_options: dict | None = None, jit_options: dict | None = None, entity_maps: typing.Sequence[dolfinx.mesh.EntityMap] | None = None, + kind: typing.Any = None, + petsc_options: dict | None = None, + petsc_options_prefix: str | None = None, + adjoint_petsc_options: dict | None = None, + tlm_petsc_options: dict | None = None, ad_block_tag: str | None = None, problem: "LinearProblem" = ..., ) -> None: ... @@ -944,20 +868,30 @@ def __init__( form_compiler_options: dict | None = None, jit_options: dict | None = None, entity_maps: typing.Sequence[dolfinx.mesh.EntityMap] | None = None, + kind: typing.Any = None, + petsc_options: dict | None = None, + petsc_options_prefix: str | None = None, + adjoint_petsc_options: dict | None = None, + tlm_petsc_options: dict | None = None, ad_block_tag: str | None = None, problem: "LinearProblem" = None, # type: ignore[assignment] ) -> None: assert problem is not None, "problem must be provided." - # Strong reference, deliberately: a throwaway LinearProblem (solve it, - # then only touch the ReducedFunctional) is a common pattern, so the - # block must keep the Problem -- and its shared solvers -- alive for - # as long as the block itself is reachable. Not cyclic garbage on its - # own (verified), so this doesn't reintroduce the MPI - # collective-destruction hazard from dolfinx-adjoint-knowledge's - # mpi-collective-destruction-hazard note -- that needs an unmerged - # checkpoint schedule making the *tape* cyclic. - self._problem_obj = problem + # Held via a weakref, not a strong reference: a throwaway LinearProblem + # (solve it, then only touch the ReducedFunctional) must not be kept + # alive for as long as this Block is reachable from the tape. The + # remaining constructor arguments cached below (kind/petsc_options/...) + # are not needed for this Block's own forward/adjoint/TLM/Hessian math + # -- that all lives on `problem` -- they exist solely to let + # _rebuild_problem() reconstruct an equivalent LinearProblem if the + # original one is ever collected while this Block still needs it. + self._problem_ref = weakref.ref(problem) + self._kind = kind + self._petsc_options = petsc_options + self._petsc_options_prefix = petsc_options_prefix + self._adjoint_petsc_options = adjoint_petsc_options + self._tlm_petsc_options = tlm_petsc_options super().__init__(ad_block_tag=ad_block_tag) # Collect all arguments in variational forms and replace them with similar @@ -1032,7 +966,7 @@ def __init__( self.add_dependency(bc, no_duplicates=True) # No forward/adjoint/TLM solver is built here: this block shares the - # ones owned by self._problem (see LinearProblem in ../solvers.py), + # ones owned by self.get_reference_problem() (see LinearProblem in ../solvers.py), # built once and reused across every block that Problem records # instead of once per solve() call. @@ -1049,6 +983,10 @@ def __init__( def _compute_residual(self) -> tuple[ufl.Form, dict[Function, Function]]: """Convert the formulation :math:`a(u, v)=L(v)` into a residual :math:`F(u_b, v) = 0` where :math:`u_b` is the solution of the forward problem at the current time and all coefficients are updated. + + Returns: + A ``(F_form, replacement_map)`` pair, per + {py:meth}`~dolfinx_adjoint.blocks.solvers._ProblemBlockBase._compute_residual`. """ # NOTE: Should probably be possible to compile this form once. replacement_functions = self.get_outputs() @@ -1063,11 +1001,43 @@ def _compute_residual(self) -> tuple[ufl.Form, dict[Function, Function]]: F_form = ufl.replace(F_form, replacement_map) return F_form, replacement_map + def _rebuild_problem(self) -> "LinearProblem": + """See {py:meth}`~dolfinx_adjoint.blocks.solvers._ProblemBlockBase._rebuild_problem`. + + Returns: + A freshly constructed {py:class}`~dolfinx_adjoint.LinearProblem`, built from the ``a``/``L``/bcs/ + options this block itself stored at construction time. + """ + from ..solvers import LinearProblem + + warnings.warn( + "This block's LinearProblem was garbage collected before being " + "recomputed/differentiated; rebuilding an equivalent one. Keep the " + "LinearProblem object alive for as long as its blocks may need " + "replay to avoid this cost.", + stacklevel=4, + ) + return LinearProblem( + self._lhs, # type: ignore[arg-type] + self._rhs, # type: ignore[arg-type] + bcs=self._bcs, + u=self._u, # type: ignore[arg-type] + P=self._preconditioner, # type: ignore[arg-type] + kind=self._kind, + petsc_options=self._petsc_options, + petsc_options_prefix=self._petsc_options_prefix, + adjoint_petsc_options=self._adjoint_petsc_options, + tlm_petsc_options=self._tlm_petsc_options, + form_compiler_options=self._form_compiler_options, + jit_options=self._jit_options, + entity_maps=self._entity_maps, + ) # type: ignore[misc] + class NonlinearProblemBlock(_ProblemBlockBase): - """A linear problem that can be used with adjoint methods. + """A nonlinear problem that can be used with adjoint methods. - This class extends the `dolfinx.fem.petsc.LinearProblem` to support adjoint methods. + This class extends the `dolfinx.fem.petsc.NonlinearProblem` to support adjoint methods. """ _adjoint_solutions: dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function] @@ -1086,6 +1056,11 @@ def __init__( form_compiler_options: dict | None = None, jit_options: dict | None = None, entity_maps: typing.Sequence[dolfinx.mesh.EntityMap] | None = None, + kind: typing.Any = None, + petsc_options: dict | None = None, + petsc_options_prefix: str | None = None, + adjoint_petsc_options: dict | None = None, + tlm_petsc_options: dict | None = None, ad_block_tag: str | None = None, problem: "NonlinearProblem" = ..., ) -> None: ... @@ -1101,6 +1076,11 @@ def __init__( form_compiler_options: dict | None = None, jit_options: dict | None = None, entity_maps: typing.Sequence[dolfinx.mesh.EntityMap] | None = None, + kind: typing.Any = None, + petsc_options: dict | None = None, + petsc_options_prefix: str | None = None, + adjoint_petsc_options: dict | None = None, + tlm_petsc_options: dict | None = None, ad_block_tag: str | None = None, problem: "NonlinearProblem" = ..., ) -> None: ... @@ -1115,14 +1095,28 @@ def __init__( form_compiler_options: dict | None = None, jit_options: dict | None = None, entity_maps: typing.Sequence[dolfinx.mesh.EntityMap] | None = None, + kind: typing.Any = None, + petsc_options: dict | None = None, + petsc_options_prefix: str | None = None, + adjoint_petsc_options: dict | None = None, + tlm_petsc_options: dict | None = None, ad_block_tag: str | None = None, problem: "NonlinearProblem" = None, # type: ignore[assignment] ) -> None: 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 + # Held via a weakref -- see LinearProblemBlock.__init__ for the + # rationale. The remaining constructor arguments cached below are not + # needed for this Block's own math (that all lives on `problem`); + # they exist solely for _rebuild_problem() to reconstruct an + # equivalent NonlinearProblem if the original one is ever collected + # while this Block still needs it. + self._problem_ref = weakref.ref(problem) + self._kind = kind + self._petsc_options = petsc_options + self._petsc_options_prefix = petsc_options_prefix + self._adjoint_petsc_options = adjoint_petsc_options + self._tlm_petsc_options = tlm_petsc_options super().__init__(ad_block_tag=ad_block_tag) self._preconditioner = P @@ -1140,12 +1134,15 @@ def __init__( assert isinstance(F, typing.Iterable) replace_dict = {ui: _ui for ui, _ui in zip(u, self._u)} self._rhs = [ufl.replace(Fi, replace_dict) for Fi in F] + # Kept only for _rebuild_problem() -- see NonlinearProblem.__init__'s + # own self._user_J for why this must never be read for anything else. + self._user_J = J # NOTE: Add mesh and constants as dependencies later on u_list = self._u if isinstance(self._u, list) else [self._u] - for c in _collect_coefficients(J) - set(u_list): + for c in collect_coefficients(J) - set(u_list): self.add_dependency(c, no_duplicates=True) - for c in _collect_coefficients(self._rhs) - set(u_list): + for c in collect_coefficients(self._rhs) - set(u_list): self.add_dependency(c, no_duplicates=True) # Cache form parameters for later @@ -1156,7 +1153,7 @@ def __init__( self._bcs = bcs if bcs is not None else [] # No forward/adjoint solver is built here: this block shares the ones - # owned by self._problem (see NonlinearProblem in ../solvers.py), + # owned by self.get_reference_problem() (see NonlinearProblem in ../solvers.py), # built once and reused across every block that Problem records # instead of once per solve() call. @@ -1173,11 +1170,17 @@ def __init__( def _compute_residual(self) -> tuple[ufl.Form, dict[Function, Function]]: """Build the residual :math:`F(u_b, v) = 0` at the current checkpointed dependency values. - Unlike ``LinearProblemBlock`` (which derives ``F`` from ``a``/``L`` via ``ufl.action``), - ``F`` is already the residual here -- this only needs to substitute in the checkpointed + Unlike {py:class}`~dolfinx_adjoint.blocks.solvers.LinearProblemBlock` (which + derives ``F`` from ``a``/``L`` via {py:func}`ufl.action`), ``F`` is already + the residual here -- this only needs to substitute in the checkpointed dependency and output values. Settles on the same output shape as - ``LinearProblemBlock._compute_residual`` (a single summed form plus its replacement map) - so every method built on top of it (adjoint, TLM, Hessian) is shared on the base class. + {py:meth}`~dolfinx_adjoint.blocks.solvers.LinearProblemBlock._compute_residual` + (a single summed form plus its replacement map) so every method built on + top of it (adjoint, TLM, Hessian) is shared on the base class. + + Returns: + A ``(F_form, replacement_map)`` pair, per + {py:meth}`~dolfinx_adjoint.blocks.solvers._ProblemBlockBase._compute_residual`. """ replacement_functions = self.get_outputs() replacement_map = self._create_replace_map(self._rhs) @@ -1189,15 +1192,50 @@ def _compute_residual(self) -> tuple[ufl.Form, dict[Function, Function]]: F_form = ufl.replace(sum_form(self._rhs), replacement_map) return F_form, replacement_map - def _refresh_dFdu_state(self, problem: typing.Any) -> None: + def _rebuild_problem(self) -> "NonlinearProblem": + """See {py:meth}`~dolfinx_adjoint.blocks.solvers._ProblemBlockBase._rebuild_problem`. + + Returns: + A freshly constructed {py:class}`~dolfinx_adjoint.NonlinearProblem`, built from the ``F``/``J``/ + bcs/options this block itself stored at construction time. + """ + from ..solvers import NonlinearProblem + + warnings.warn( + "This block's NonlinearProblem was garbage collected before being " + "recomputed/differentiated; rebuilding an equivalent one. Keep the " + "NonlinearProblem object alive for as long as its blocks may need " + "replay to avoid this cost.", + stacklevel=4, + ) + return NonlinearProblem( + self._rhs, # type: ignore[arg-type] + u=self._u, # type: ignore[arg-type] + bcs=self._bcs, + J=self._user_J, # type: ignore[arg-type] + P=self._preconditioner, # type: ignore[arg-type] + kind=self._kind, + petsc_options=self._petsc_options, + petsc_options_prefix=self._petsc_options_prefix, + adjoint_petsc_options=self._adjoint_petsc_options, + tlm_petsc_options=self._tlm_petsc_options, + form_compiler_options=self._form_compiler_options, + jit_options=self._jit_options, + entity_maps=self._entity_maps, + ) # type: ignore[misc] + + def _refresh_dFdu_state(self, problem: "LinearProblem | NonlinearProblem") -> None: """Refresh ``problem``'s "state" placeholder(s) from this block's own saved outputs. - Unlike ``LinearProblem`` (where ``dF/du`` never references the state), ``dF/du`` here + Unlike {py:class}`~dolfinx_adjoint.LinearProblem` (where ``dF/du`` never references the state), ``dF/du`` here genuinely depends on ``u``'s current value, so the shared adjoint solver's compiled LHS must be evaluated at *this* block's checkpointed output before it is used -- another block sharing the same solver may have left a different value there. + + Args: + problem: This block's owning Problem (see ``self.get_reference_problem()``). """ - state_placeholder = problem._residual_state_placeholder + state_placeholder = problem.residual_state_placeholder 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[:] diff --git a/src/dolfinx_adjoint/solvers.py b/src/dolfinx_adjoint/solvers.py index c4709f3..88f2567 100644 --- a/src/dolfinx_adjoint/solvers.py +++ b/src/dolfinx_adjoint/solvers.py @@ -1,5 +1,7 @@ from __future__ import annotations +import abc +import itertools import typing import dolfinx.fem.petsc @@ -7,44 +9,26 @@ import ufl from dolfinx.fem.function import Function as _Function -from .blocks.solvers import ( - LinearProblemBlock, - NonlinearProblemBlock, - _collect_coefficients, - _ProblemBlockBase, +from .blocks.solvers import LinearProblemBlock, NonlinearProblemBlock, _ProblemBlockBase +from .petsc_utils import HomogeneousBCLinearProblem +from .types import Function +from .ufl_utils import ( assign_mixed_parts, + collect_coefficients, + compute_adjoint, get_sorted_arguments, + recursive_replace, sum_form, ) -from .petsc_utils import HomogeneousBCLinearProblem -from .types import Function - -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] +# Backs each Problem's default PETSc options prefix (see LinearProblem/NonlinearProblem +# __init__). A plain incrementing counter, not id(self)/uuid.uuid4()/time.time(): those +# can differ across MPI ranks for the same logical Problem (memory layout, clock skew), +# and PETSc's options-database handling around SNESSetFromOptions/KSPSetFromOptions is +# collective, so every rank must resolve the same prefix for the same Problem. A counter +# incremented once per Problem construction is deterministic and identical on every rank, +# since construction happens in lock-step in a well-formed SPMD program. +_PROBLEM_PREFIX_COUNTER = itertools.count() @typing.overload @@ -56,6 +40,26 @@ def resolve_u(u: typing.Sequence[_Function] | None, L: typing.Sequence[ufl.Form] def resolve_u( u: _Function | typing.Sequence[_Function] | None, L: ufl.Form | typing.Sequence[ufl.Form] ) -> _Function | typing.Sequence[_Function]: + """Resolve the unknown {py:class}`~dolfinx_adjoint.Function` ``u`` for a `*Problem`. + + If ``u`` was not supplied by the caller, a fresh {py:class}`~dolfinx_adjoint.Function` + is created per block, using the function space of the corresponding + {py:class}`ufl.Argument` in ``L``. If ``u`` was supplied, it is wrapped with + {py:func}`pyadjoint.create_overloaded_object` so the tape can record operations on it, + regardless of whether it arrived as a plain {py:class}`~dolfinx_adjoint.Function` or + already-overloaded one. + + Args: + u: The unknown Function (or, for a blocked problem, a sequence of them), or + ``None`` to have one created per block. + L: The right-hand-side form (or, for a blocked problem, a sequence of forms) whose + test-function space determines the space of a newly created ``u``. Only + consulted when ``u`` is ``None``. + + Returns: + A single {py:class}`~dolfinx_adjoint.Function`, or a list of them for a blocked + problem, matching the shape of ``L``. + """ if u is None: try: # Extract function space for unknown from the right hand @@ -75,7 +79,8 @@ def resolve_u( class HessianTemplates(typing.NamedTuple): """Per-dependency compiled templates used to assemble a Hessian action. - Shared shape for ``LinearProblem``/``NonlinearProblem``, so callers in + Shared shape for {py:class}`~dolfinx_adjoint.LinearProblem`/ + {py:class}`~dolfinx_adjoint.NonlinearProblem`, so callers in ``blocks/solvers.py`` unpack by name rather than by position -- the two classes used to return differently-shaped tuples here, which was a latent footgun for any code touching both. @@ -83,25 +88,57 @@ class HessianTemplates(typing.NamedTuple): Attributes: soa_self: The SOA right-hand-side's contribution from ``dF/du``'s own second derivative w.r.t. ``u`` (``d2Fdu2``) -- a - ``ufl.ZeroBaseForm`` if that term is structurally zero, so callers + {py:class}`ufl.ZeroBaseForm` if that term is structurally zero, so callers can always assemble it unconditionally. Built by the same code - (``_build_soa_self_template``) for both classes; it simply always + ({py:func}`~dolfinx_adjoint.solvers._build_soa_self_template`) for both classes; it simply always compiles to zero for a linear problem, since ``dF/du`` doesn't - reference ``u`` there -- not a hardcoded special case. + reference ``u`` there -- not a hardcoded special case. A list of + one compiled form per output row (not a single form) for a + blocked problem, since the SOA right-hand side is itself block + structured then. soa_cross: The SOA right-hand-side's contribution, per dependency, - from that dependency's tangent-linear direction. + from that dependency's tangent-linear direction. Also a list of + one form per output row, per dependency, for a blocked problem. fixed: The part of each dependency's own Hessian-action output that does not depend on any *other* dependency's tangent-linear value. + One form per dependency regardless of blocking -- it lives on the + *control's* own test space, not the (possibly blocked) state's. cross: Each dependency's Hessian-action contribution from *another* - dependency's tangent-linear direction, keyed by ``(c, c2)``. + dependency's tangent-linear direction, keyed by ``(c, c2)``. Same + per-dependency (not per-row) shape as ``fixed``. """ - soa_self: dolfinx.fem.Form + soa_self: dolfinx.fem.Form | list[dolfinx.fem.Form] soa_cross: dict fixed: dict cross: dict +def _pad_blocks_by_part(form: ufl.form.BaseForm, test_funcs: typing.Sequence[ufl.Argument]) -> list[ufl.form.BaseForm]: + """Split a blocked one-form into one entry per ``test_funcs`` part, in part order. + + {py:func}`ufl.extract_blocks` only returns an entry for a part that actually appears in + ``form`` -- a part a differentiation happened to eliminate entirely (e.g. a + dependency that only appears in one output block's equation) is simply absent from + its result, not returned as an explicit zero. Padding with {py:class}`ufl.ZeroBaseForm` for + every missing part keeps every caller's per-row list the same length and order as + ``test_funcs``, so it can always be assembled/indexed positionally. + + ``form.empty()`` is checked before calling {py:func}`ufl.extract_blocks`: a form that is + structurally empty (no arguments at all, e.g. ``d2Fdu2`` for a linear residual) has + no parts for {py:func}`ufl.extract_blocks` to find, so every row is padded to zero directly + instead. + """ + padded: list[ufl.form.BaseForm] = [ufl.ZeroBaseForm((test,)) for test in test_funcs] + if form.empty(): + return padded + for block in ufl.extract_blocks(form): + args = block.arguments() + assert len(args) == 1, "Expected a single test function in the block." + padded[args[0].part()] = block + return padded + + def _build_soa_self_template( dFdu_template: ufl.Form, state_placeholder: dolfinx.fem.Function, @@ -114,13 +151,14 @@ def _build_soa_self_template( ) -> dolfinx.fem.Form: """Build the SOA self-term ``adjoint(d2F/du2) . adjoint_solution``. - The same computation for both ``LinearProblem`` and ``NonlinearProblem``: - ``d2F/du2`` is structurally zero for a linear residual (``dF/du`` doesn't - reference ``u``), so this compiles to a ``ufl.ZeroBaseForm`` for - ``LinearProblem`` as a *result* of running the same code, not as a special - case -- callers can always assemble the returned form unconditionally, - mirroring how an inactive TLM/fixed template is represented elsewhere in - this module (e.g. ``_get_or_build_tlm_rhs_templates``). + The same computation for both {py:class}`~dolfinx_adjoint.LinearProblem` and + {py:class}`~dolfinx_adjoint.NonlinearProblem`: ``d2F/du2`` is structurally zero + for a linear residual (``dF/du`` doesn't reference ``u``), so this compiles to + a {py:class}`ufl.ZeroBaseForm` for {py:class}`~dolfinx_adjoint.LinearProblem` + as a *result* of running the same code, not as a special case -- callers can + always assemble the returned form unconditionally, mirroring how an inactive + TLM/fixed template is represented elsewhere in this module (e.g. + {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_tlm_rhs_templates`). """ d2Fdu2 = ufl.algorithms.expand_derivatives(ufl.derivative(dFdu_template, state_placeholder, hessian_u_seed)) if d2Fdu2.empty(): @@ -135,15 +173,18 @@ def _build_soa_self_template( ) -class _ProblemBase: - """Shared lazy adjoint/TLM solver machinery for ``LinearProblem``/``NonlinearProblem``. +class _ProblemBase(abc.ABC): + """Shared lazy adjoint/TLM solver machinery for + {py:class}`~dolfinx_adjoint.LinearProblem`/{py:class}`~dolfinx_adjoint.NonlinearProblem`. A plain mixin -- it does not inherit from any ``dolfinx.fem.petsc`` class, so ``LinearProblem(_ProblemBase, dolfinx.fem.petsc.LinearProblem)`` (and - the ``NonlinearProblem`` equivalent) has an unambiguous MRO for - ``solve()``: each subclass keeps defining its own ``solve()``, delegating - the shared middle to ``_record_and_solve`` below via the ``_make_block``/ - ``_dolfinx_solve`` hooks it implements. + the {py:class}`~dolfinx_adjoint.NonlinearProblem` equivalent) has an + unambiguous MRO for ``solve()``: each subclass keeps defining its own + ``solve()``, delegating the shared middle to + {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._record_and_solve` below via + the {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._make_block`/ + {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._dolfinx_solve` hooks it implements. Every attribute referenced here is set by the concrete subclass's own ``__init__`` (either directly, or inherited from the ``dolfinx.fem.petsc`` @@ -164,24 +205,96 @@ class _ProblemBase: _petsc_options_prefix: str _kind: typing.Any + @property + def value_placeholders(self) -> dict[dolfinx.fem.Function, dolfinx.fem.Function]: + """Map from each non-``u`` dependency the compiled forward/adjoint/TLM/Hessian + forms were built against to its dedicated placeholder coefficient. + + Public so + {py:class}`*ProblemBlock` + methods (a different module) can refresh a dependency's placeholder value + ahead of a solve/replay without reaching into this Problem's own private + state. + + Returns: + The placeholder dictionary, keyed by the user's original dependency. + """ + return self._value_placeholders + + @property + def residual_state_placeholder(self) -> dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function]: + """The dedicated "state" placeholder(s) standing in for ``u`` in every + compiled template built from the residual (see + {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_residual_template`). + + Returns: + A single {py:class}`~dolfinx_adjoint.Function` for a scalar problem, or one + per output block for a blocked problem. + """ + assert self._residual_state_placeholder is not None + return self._residual_state_placeholder + + @property + def adjoint_solution_placeholder(self) -> dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function]: + """The placeholder(s) standing in for the first-order adjoint solution in the + cached Hessian templates (see {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_hessian_templates`). + + Returns: + A single {py:class}`~dolfinx_adjoint.Function` for a scalar problem, or one + per output block for a blocked problem. ``None`` until the Hessian templates have been built. + """ + assert self._adjoint_solution_placeholder is not None + return self._adjoint_solution_placeholder + + @property + def second_adjoint_solution_placeholder(self) -> dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function]: + """The placeholder(s) standing in for the second-order adjoint (SOA) solution + in the cached Hessian templates (see + {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_hessian_templates`). + + Returns: + A single {py:class}`~dolfinx_adjoint.Function` for a scalar problem, or one + per output block for a blocked problem. + """ + assert self._second_adjoint_solution_placeholder is not None + return self._second_adjoint_solution_placeholder + + @property + def hessian_u_seed(self) -> dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function]: + """The placeholder(s) standing in for the state's own tangent-linear direction + in the cached Hessian self-term (see + {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_hessian_templates`). + + Returns: + A single {py:class}`~dolfinx_adjoint.Function` for a scalar problem, or one + per output block for a blocked problem. + """ + assert self._hessian_u_seed is not None + return self._hessian_u_seed + def _init_adjoint_state(self) -> None: """Initialize the lazily-built adjoint/TLM solver state. Called once, at the end of ``__init__``, after the base ``dolfinx.fem.petsc`` solver has been constructed. Built lazily (on - first use, see ``_get_or_build_adjoint_solver``/ - ``_get_or_build_tlm_solver``/``_get_or_build_hessian_templates``) and - shared by every block this Problem records, rather than one per - block/solve() call. Each block holds a plain (strong) reference back - to this Problem (see ``*ProblemBlock._problem``/``__init__`` for why a - strong reference is safe here and doesn't reintroduce the MPI - collective-destruction hazard documented in - dolfinx-adjoint-knowledge's mpi-collective-destruction-hazard note), - so this Problem -- and hence its solvers' PETSc objects -- is released - deterministically via ordinary refcounting once every block - referencing it is unreachable, rather than being left to pyadjoint's - tape/cyclic-GC schedule. Laziness keeps pure forward (non-annotated) - use from paying for a symbolic adjoint form it never needs. + first use, see + {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_adjoint_solver`/ + {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_tlm_solver`/ + {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_hessian_templates`) + and shared by every block this Problem records, rather than one per + block/solve() call. Each block holds only a ``weakref`` back to this + Problem (see + {py:meth}`*ProblemBlock.get_reference_problem`), + not a strong reference: a throwaway Problem (solve it, then only touch + the resulting {py:class}`~pyadjoint.ReducedFunctional`) must not be + kept alive for as long as its blocks remain reachable on the tape. If + this Problem is collected while a block still needs it, that block + rebuilds an equivalent one on demand + ({py:meth}`*ProblemBlock._rebuild_problem`, + with a {py:class}`UserWarning` since it is a costly fallback) rather + than silently failing. Laziness keeps pure forward (non-annotated) + use from paying for a symbolic + adjoint form it never needs. """ self._adjoint_solver: HomogeneousBCLinearProblem | None = None self._tlm_solver: HomogeneousBCLinearProblem | None = None @@ -192,10 +305,13 @@ def _init_adjoint_state(self) -> None: self._tlm_rhs_templates: dict | None = None self._tlm_seed_placeholders: dict[dolfinx.fem.Function, dolfinx.fem.Function] = {} self._hessian_templates: HessianTemplates | None = None - self._adjoint_solution_placeholder: dolfinx.fem.Function | None = None - self._second_adjoint_solution_placeholder: dolfinx.fem.Function | None = None - self._hessian_u_seed: dolfinx.fem.Function | None = None + self._adjoint_solution_placeholder: dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function] | None = None + self._second_adjoint_solution_placeholder: ( + dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function] | None + ) = None + self._hessian_u_seed: dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function] | None = None + @abc.abstractmethod def _get_or_build_residual_template( self, ) -> tuple[ufl.Form, dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function]]: @@ -203,25 +319,37 @@ def _get_or_build_residual_template( replaced by a dedicated "state" placeholder standing in for "u at this evaluation point". The one genuinely irreducible difference between the two Problem kinds -- LinearProblem - builds it from ``a``/``L`` via ``ufl.action``, NonlinearProblem already has ``F`` + builds it from ``a``/``L`` via {py:func}`ufl.action`, NonlinearProblem already has ``F`` directly -- everything built on top of it below (``dF/du``, the TLM right-hand side, the Hessian templates) is derived from this one template by the same shared symbolic differentiation, since that costs nothing extra at compile time: there is no reason to keep a separate "dF/du is just a, never differentiated" shortcut for LinearProblem. - Not implemented on the base; each subclass overrides it. + Each subclass overrides this. + + Returns: + A ``(F_template, state_placeholder)`` pair: the residual with every + coefficient (including ``u``) substituted by its dedicated + placeholder, and that state placeholder itself (a single + {py:class}`~dolfinx_adjoint.Function`, or one per output block for a + blocked problem). """ - raise NotImplementedError def _get_or_build_dFdu_template(self) -> ufl.Form | typing.Sequence: """Build (once) and return dF/du, evaluated at the residual template's state placeholder. - Shared by both classes: derived from ``_get_or_build_residual_template`` by symbolic - differentiation, which is free at compile time, rather than special-cased per class. - This is the shared basis for 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 - ``*ProblemBlock.prepare_evaluate_adj``/``prepare_evaluate_hessian``/``prepare_evaluate_tlm``). + Shared by both classes: derived from + {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_residual_template` + by symbolic differentiation, which is free at compile time, rather than + special-cased per class. This is the shared basis for the adjoint operator + ({py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_adjoint_solver`, + which just adjoints it) and the TLM operator + ({py:meth}`~dolfinx_adjoint.solvers._ProblemBase._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 + {py:meth}`*ProblemBlock.prepare_evaluate_adj`/ + {py:meth}`~dolfinx_adjoint.blocks.solvers._ProblemBlockBase.prepare_evaluate_hessian`/ + {py:meth}`~dolfinx_adjoint.blocks.solvers._ProblemBlockBase.prepare_evaluate_tlm`). """ if self._dFdu_template is None: F_template, state_placeholder = self._get_or_build_residual_template() @@ -243,22 +371,26 @@ def _get_or_build_dFdu_template(self) -> ufl.Form | typing.Sequence: 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``). - - The same computation for both classes: ``_ProblemBlockBase._compute_adjoint`` - swaps argument numbers while preserving mixed-space ``part()`` tags (via - ``compute_form_adjoint``) and decomposes the result back into blocks (via - ``ufl.extract_blocks``) -- a no-op decomposition for a scalar, non-blocked - form. Kept exactly as that returns it (a nested list of forms for a blocked - problem) since that structure is what ``HomogeneousBCLinearProblem``/ - ``dolfinx.fem.petsc.LinearProblem`` needs for block matrix assembly; callers - that need a single summed form (Hessian templating, scalar-only) apply - ``sum_form`` themselves. + ({py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_adjoint_solver`) + and, for scalar problems, the Hessian SOA right-hand-side's + cross-dependency templates + ({py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_hessian_templates`). + + The same computation for both classes: + {py:func}`~dolfinx_adjoint.ufl_utils.compute_adjoint` swaps argument + numbers while preserving mixed-space {py:meth}`ufl.Argument.part` tags + (via {py:func}`~dolfinx_adjoint.compat.compute_form_adjoint`) and + decomposes the result back into blocks (via + {py:func}`ufl.extract_blocks`) -- a no-op decomposition for a scalar, + non-blocked form. Kept exactly as that returns it (a nested list of + forms for a blocked problem) since that structure is what + {py:class}`~dolfinx_adjoint.petsc_utils.HomogeneousBCLinearProblem`/ + {py:class}`dolfinx.fem.petsc.LinearProblem` needs for block matrix + assembly; callers that need a single summed form (Hessian templating, + scalar-only) apply {py:func}`~dolfinx_adjoint.ufl_utils.sum_form` themselves. """ if self._dFdu_adj_template is None: - self._dFdu_adj_template = _ProblemBlockBase._compute_adjoint( + self._dFdu_adj_template = compute_adjoint( self._get_or_build_dFdu_template() # type: ignore[arg-type] ) return self._dFdu_adj_template @@ -273,10 +405,11 @@ def _get_or_build_tlm_rhs_templates( """Build (once) and return the per-dependency TLM right-hand-side templates. Shared by both classes: built purely from the residual template - (``_get_or_build_residual_template``), since ``dF/dm`` genuinely depends on the state - for either a linear or nonlinear residual -- differentiating w.r.t. a coefficient - embedded in the residual while holding ``u`` fixed leaves ``u`` in the result even when - the residual is linear in ``u`` itself. + ({py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_residual_template`), + since ``dF/dm`` genuinely depends on the state for either a linear or + nonlinear residual -- differentiating w.r.t. a coefficient embedded in + the residual while holding ``u`` fixed leaves ``u`` in the result even + when the residual is linear in ``u`` itself. One compiled one-form is built per dependency, using a dedicated "direction" placeholder for that dependency (``_tlm_seed_placeholders``) rather than a single @@ -289,7 +422,8 @@ def _get_or_build_tlm_rhs_templates( `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 - ``*ProblemBlock.prepare_evaluate_tlm``), avoids that entirely. + {py:meth}`*ProblemBlock.prepare_evaluate_tlm`), + avoids that entirely. """ F_template, state_placeholder = self._get_or_build_residual_template() if self._tlm_rhs_templates is None: @@ -303,13 +437,7 @@ def _get_or_build_tlm_rhs_templates( 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 + dFdm_c = _pad_blocks_by_part(dFdm_c, test_funcs) else: if dFdm_c == 0 or dFdm_c.empty(): dFdm_c = ufl.ZeroBaseForm((test_funcs[0],)) @@ -325,51 +453,96 @@ def _get_or_build_tlm_rhs_templates( def _get_or_build_hessian_templates(self) -> HessianTemplates: """Build (once) and return the per-dependency Hessian templates used by - ``*ProblemBlock.prepare_evaluate_hessian``'s SOA right-hand side and - ``evaluate_hessian_component``'s own Hessian-action output. - - Shared by both classes -- built purely from ``_get_or_build_residual_template``/ - ``_get_or_build_dFdu_template``/``_get_or_build_dFdu_adj_template``/ - ``_get_or_build_tlm_rhs_templates``, all themselves shared. ``soa_self`` (see - ``_build_soa_self_template``) comes out a ``ufl.ZeroBaseForm`` for ``LinearProblem`` - (``dF/du`` doesn't depend on ``u``) and generally nonzero for ``NonlinearProblem``, as a - *result* of running the same code, not a per-class branch. - - Scalar (non-blocked) problems only -- the blocked Hessian cross-term stays on the - pre-templating, per-call ``ufl.replace`` + recompile path (see - ``_ProblemBlockBase._evaluate_hessian_blocked_rhs``/``_evaluate_hessian_component_blocked``). - - Each of ``soa_cross``/``fixed``/``cross`` 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. + {py:meth}`*ProblemBlock.prepare_evaluate_hessian`'s + SOA right-hand side and + {py:meth}`~dolfinx_adjoint.blocks.solvers._ProblemBlockBase.evaluate_hessian_component`'s + own Hessian-action output. + + Shared by both classes and by both scalar and blocked problems -- built + purely from + {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_residual_template`/ + {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_dFdu_template`/ + {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_dFdu_adj_template`/ + {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_tlm_rhs_templates`, + all themselves shared. ``soa_self`` (see + {py:func}`~dolfinx_adjoint.solvers._build_soa_self_template`) comes out a + {py:class}`ufl.ZeroBaseForm` for {py:class}`~dolfinx_adjoint.LinearProblem` + (``dF/du`` doesn't depend on ``u``) and generally nonzero for + {py:class}`~dolfinx_adjoint.NonlinearProblem`, as a *result* of running the + same code, not a per-class branch. + + ``fixed``/``cross`` are one-forms over a *control's own* test space + (``c.function_space``), so their shape doesn't depend on how many output + blocks ``u`` has -- {py:func}`ufl.action` already reduces the (possibly + part-tagged, blocked) state/adjoint-solution arguments away before + ``fixed``/``cross`` are built. Only ``soa_self``/``soa_cross`` feed the + (possibly blocked) SOA right-hand-side vector, so for a blocked problem + they become a list of one compiled form per output row (padded via + {py:func}`~dolfinx_adjoint.solvers._pad_blocks_by_part` for any row a + differentiation happened to eliminate entirely) instead of a single form. + + Each of ``soa_cross``/``fixed``/``cross`` is kept as its own compiled + one-form (or list of one-forms), 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 + {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._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() F_template, _ = self._get_or_build_residual_template() 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(state_placeholder.function_space) - self._second_adjoint_solution_placeholder = dolfinx.fem.Function(state_placeholder.function_space) - self._hessian_u_seed = dolfinx.fem.Function(state_placeholder.function_space) + blocked = isinstance(self._u, list) + soa_self: dolfinx.fem.Form | list[dolfinx.fem.Form] + if blocked: + assert isinstance(state_placeholder, typing.Sequence) + state_list = list(state_placeholder) + test_funcs = list(get_sorted_arguments(F_template.arguments(), 0)) + self._adjoint_solution_placeholder = [dolfinx.fem.Function(s.function_space) for s in state_list] + self._second_adjoint_solution_placeholder = [dolfinx.fem.Function(s.function_space) for s in state_list] + self._hessian_u_seed = [dolfinx.fem.Function(s.function_space) for s in state_list] + state_arg: typing.Any = state_list - soa_self = _build_soa_self_template( - dFdu_template, - state_placeholder, - self._hessian_u_seed, - self._adjoint_solution_placeholder, - jit_options=self._jit_options, - form_compiler_options=self._form_compiler_options, - entity_maps=self._entity_maps, - ) + d2Fdu2 = ufl.algorithms.expand_derivatives( + ufl.derivative(dFdu_template, state_list, self._hessian_u_seed) + ) + if d2Fdu2.empty(): + soa_self_form = d2Fdu2 + else: + soa_self_form = ufl.action(ufl.adjoint(d2Fdu2), self._adjoint_solution_placeholder) + soa_self = [ + dolfinx.fem.form( + form_i, + jit_options=self._jit_options, + form_compiler_options=self._form_compiler_options, + entity_maps=self._entity_maps, + ) + for form_i in _pad_blocks_by_part(soa_self_form, test_funcs) + ] + else: + assert isinstance(state_placeholder, dolfinx.fem.Function) + self._adjoint_solution_placeholder = dolfinx.fem.Function(state_placeholder.function_space) + self._second_adjoint_solution_placeholder = dolfinx.fem.Function(state_placeholder.function_space) + self._hessian_u_seed = dolfinx.fem.Function(state_placeholder.function_space) + state_arg = state_placeholder + + soa_self = _build_soa_self_template( + dFdu_template, + state_placeholder, + self._hessian_u_seed, + self._adjoint_solution_placeholder, + 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) L1 = ufl.action(F_template, self._adjoint_solution_placeholder) @@ -383,19 +556,28 @@ def _get_or_build_hessian_templates(self) -> HessianTemplates: 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, - ) + if blocked: + soa_cross_templates[c] = [ + dolfinx.fem.form( + form_i, + jit_options=self._jit_options, + form_compiler_options=self._form_compiler_options, + entity_maps=self._entity_maps, + ) + for form_i in _pad_blocks_by_part(soa_form, test_funcs) + ] + else: + 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) - ) + d2Fdudm = ufl.algorithms.expand_derivatives(ufl.derivative(dL1dm, state_arg, 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,)) @@ -442,14 +624,16 @@ def _get_or_build_tlm_solver(self) -> HomogeneousBCLinearProblem: 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 ``*ProblemBlock.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. + result out (see + {py:meth}`*ProblemBlock.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 + itself, inside {py:func}`~dolfinx_adjoint.compat.compute_form_adjoint`/ + {py:func}`~dolfinx_adjoint.ufl_utils.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 + with {py:func}`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. @@ -472,18 +656,58 @@ def _get_or_build_tlm_solver(self) -> HomogeneousBCLinearProblem: ) # type: ignore[misc] return self._tlm_solver + @abc.abstractmethod + def _make_block(self) -> _ProblemBlockBase: + """Construct the tape block this Problem records for its forward solve. + + Each subclass overrides this to instantiate its own Block kind + ({py:class}`~dolfinx_adjoint.blocks.solvers.LinearProblemBlock`/ + {py:class}`~dolfinx_adjoint.blocks.solvers.NonlinearProblemBlock`, + constructor kwargs differing the same way the two Problem kinds' own + constructors do), passing ``self`` so the Block can reach back into + this Problem's shared solvers (see + {py:meth}`*ProblemBlock.get_reference_problem`). + + Returns: + A newly constructed, not-yet-recorded Block for this solve. + """ + + @abc.abstractmethod + def _dolfinx_solve(self) -> _Function | typing.Sequence[_Function]: + """Perform the actual forward solve, via the base ``dolfinx.fem.petsc`` class. + + Each subclass overrides this to call its own base class's ``solve()`` + directly ({py:meth}`dolfinx.fem.petsc.LinearProblem.solve` / + {py:meth}`NonlinearProblem.solve`) + rather than ``self.solve()``, which would recurse back into this + Problem's own overridden, tape-recording ``solve()``. + + Returns: + The solution {py:class}`~dolfinx_adjoint.Function`, or one per + output block for a blocked problem. + """ + def _record_and_solve(self, annotate: bool) -> _Function | typing.Sequence[_Function]: """Shared ``solve()`` skeleton for both classes. - Records a tape block (via the subclass's ``_make_block`` hook) when + Records a tape block (via the subclass's + {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._make_block` hook) when annotating, refreshes the forward solver's placeholder coefficients from the user's own current values (a prior recompute -- see - ``*ProblemBlock.prepare_recompute_component`` -- may have left them - holding a checkpointed/candidate value instead), solves (via the - subclass's ``_dolfinx_solve`` hook), and records the block's outputs. + {py:meth}`*ProblemBlock.prepare_recompute_component` + -- may have left them holding a checkpointed/candidate value instead), + solves (via the subclass's + {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._dolfinx_solve` hook), + and records the block's outputs. + + Args: + annotate: Whether to record this solve as a block on the working tape. + + Returns: + The solution {py:class}`~dolfinx_adjoint.Function`, or one per output block for a blocked problem. """ annotate = pyadjoint.annotate_tape({"annotate": annotate}) - block = self._make_block() if annotate else None # type: ignore[attr-defined] + block = self._make_block() if annotate else None if annotate: assert block is not None tape = pyadjoint.get_working_tape() @@ -493,7 +717,7 @@ def _record_and_solve(self, annotate: bool) -> _Function | typing.Sequence[_Func placeholder.x.array[:] = original.x.array[:] placeholder.x.scatter_forward() - out = self._dolfinx_solve() # type: ignore[attr-defined] + out = self._dolfinx_solve() if annotate: assert block is not None if isinstance(out, Function): @@ -538,7 +762,7 @@ def __init__( P: ufl.Form | None = None, kind: str | None = None, petsc_options: dict | None = None, - petsc_options_prefix: str = "dxa_linear_problem_", + petsc_options_prefix: str | None = None, form_compiler_options: dict | None = None, jit_options: dict | None = None, entity_maps: typing.Sequence[dolfinx.mesh.EntityMap] | None = None, @@ -557,7 +781,7 @@ def __init__( P: typing.Sequence[typing.Sequence[ufl.Form]] | None = None, kind: str | typing.Sequence[typing.Sequence[str]] | None = None, petsc_options: dict | None = None, - petsc_options_prefix: str = "dxa_linear_problem_", + petsc_options_prefix: str | None = None, form_compiler_options: dict | None = None, jit_options: dict | None = None, entity_maps: typing.Sequence[dolfinx.mesh.EntityMap] | None = None, @@ -575,7 +799,7 @@ def __init__( 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, - petsc_options_prefix: str = "dxa_linear_problem_", + petsc_options_prefix: str | None = None, form_compiler_options: dict | None = None, jit_options: dict | None = None, entity_maps: typing.Sequence[dolfinx.mesh.EntityMap] | None = None, @@ -599,13 +823,19 @@ def __init__( self._u = resolve_u(u, L) # type: ignore[arg-type] + # A caller-omitted prefix must still be unique per Problem (see + # _PROBLEM_PREFIX_COUNTER) so PETSc's process-global options database + # never lets two Problems' SNES/KSP options bleed into each other. + if petsc_options_prefix is None: + petsc_options_prefix = f"dxa_linear_problem_{next(_PROBLEM_PREFIX_COUNTER)}_" + # Cache some objects self._lhs = a self._rhs = L + self._petsc_options = petsc_options self._jit_options = jit_options 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 @@ -622,9 +852,9 @@ def __init__( # 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) + coefficients = collect_coefficients(a) | collect_coefficients(L) if P is not None: - coefficients |= _collect_coefficients(P) + 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 @@ -632,11 +862,11 @@ def __init__( # Initialize linear solver super().__init__( - a=_replace_with_placeholders(a, self._value_placeholders), # type: ignore[arg-type] - L=_replace_with_placeholders(L, self._value_placeholders), # type: ignore[arg-type] + a=recursive_replace(a, self._value_placeholders), # type: ignore[arg-type] + L=recursive_replace(L, self._value_placeholders), # type: ignore[arg-type] bcs=bcs, u=self._u, # type: ignore[arg-type] - P=_replace_with_placeholders(P, self._value_placeholders), # type: ignore[arg-type] + P=recursive_replace(P, self._value_placeholders), # type: ignore[arg-type] kind=kind, # type: ignore[arg-type] petsc_options_prefix=petsc_options_prefix, petsc_options=petsc_options, @@ -662,12 +892,16 @@ def _get_or_build_residual_template( dedicated "state" placeholder standing in for "u at this evaluation point", distinct from the live ``self._u`` the forward solve owns. - The shared basis for ``dF/du`` (``_ProblemBase._get_or_build_dFdu_template``) and the - TLM right-hand side (``_get_or_build_tlm_rhs_templates``): ``dF/dm`` genuinely depends - on the state 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), and - ``dF/du`` itself is now derived from this template by the same symbolic differentiation - ``NonlinearProblem`` uses, rather than a shortcut that skips differentiating altogether. + The shared basis for ``dF/du`` + ({py:meth}`_ProblemBase._get_or_build_dFdu_template`) + and the TLM right-hand side + ({py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_tlm_rhs_templates`): + ``dF/dm`` genuinely depends on the state 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), and + ``dF/du`` itself is now derived from this template by the same symbolic + differentiation {py:class}`~dolfinx_adjoint.NonlinearProblem` uses, + rather than a shortcut that skips differentiating altogether. """ if self._residual_template is None: u_list = self._u if isinstance(self._u, list) else [self._u] @@ -696,6 +930,11 @@ def _make_block(self) -> LinearProblemBlock: form_compiler_options=self._form_compiler_options, jit_options=self._jit_options, entity_maps=self._entity_maps, + kind=self._kind, + petsc_options=self._petsc_options, + petsc_options_prefix=self._petsc_options_prefix, + adjoint_petsc_options=self._adj_options, + tlm_petsc_options=self._tlm_options, ad_block_tag=self.ad_block_tag, problem=self, ) # type: ignore[misc] @@ -703,7 +942,7 @@ def _make_block(self) -> LinearProblemBlock: def _dolfinx_solve(self) -> _Function | typing.Sequence[_Function]: return dolfinx.fem.petsc.LinearProblem.solve(self) - def solve(self, annotate: bool = True) -> typing.Union[dolfinx.fem.Function, typing.Sequence[dolfinx.fem.Function]]: + def solve(self, annotate: bool = True) -> dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function]: """ Solve the linear problem and return the solution. """ @@ -711,18 +950,18 @@ def solve(self, annotate: bool = True) -> typing.Union[dolfinx.fem.Function, typ class NonlinearProblem(_ProblemBase, dolfinx.fem.petsc.NonlinearProblem): - """A linear problem that can be used with adjoint methods. + """A nonlinear problem that can be used with adjoint methods. - This class extends the `dolfinx.fem.petsc.LinearProblem` to support adjoint methods. + This class extends the `dolfinx.fem.petsc.NonlinearProblem` to support adjoint methods. Args: - a: The bilinear form representing the left-hand side of the equation. - L: The linear form representing the right-hand side of the equation. - bcs: Boundary conditions to apply to the problem. + F: The residual form. u: Solution vector. - P: Preconditioner for the linear problem. + bcs: Boundary conditions to apply to the problem. + J: The Jacobian form. Computed from ``F`` if not supplied. + P: Preconditioner for the nonlinear problem. kind: Kind of PETSc Matrix to assemble the system into. - petsc_options: Options dictionary for the PETSc krylov supspace solver. + petsc_options: Options dictionary for the PETSc SNES solver. form_compiler_options: Form compiler options for generating assembly kernels. jit_options: Options for just-in-time compilation of the forms. entity_maps: Mapping from meshes that coefficients and arguments are defined on to the @@ -743,7 +982,7 @@ def __init__( P: ufl.form.Form | None = None, kind: str | None = None, petsc_options: dict | None = None, - petsc_options_prefix: str = "dxa_nonlinear_problem_", + petsc_options_prefix: str | None = None, form_compiler_options: dict | None = None, jit_options: dict | None = None, entity_maps: typing.Sequence[dolfinx.mesh.EntityMap] | None = None, @@ -762,7 +1001,7 @@ def __init__( P: typing.Sequence[typing.Sequence[ufl.form.Form]] | None = None, kind: str | typing.Sequence[typing.Sequence[str]] | None = None, petsc_options: dict | None = None, - petsc_options_prefix: str = "dxa_nonlinear_problem_", + petsc_options_prefix: str | None = None, form_compiler_options: dict | None = None, jit_options: dict | None = None, entity_maps: typing.Sequence[dolfinx.mesh.EntityMap] | None = None, @@ -780,7 +1019,7 @@ def __init__( P: ufl.form.Form | typing.Sequence[typing.Sequence[ufl.form.Form]] | None = None, kind: str | typing.Sequence[typing.Sequence[str]] | None = None, petsc_options: dict | None = None, - petsc_options_prefix: str = "dxa_nonlinear_problem_", + petsc_options_prefix: str | None = None, form_compiler_options: dict | None = None, jit_options: dict | None = None, entity_maps: typing.Sequence[dolfinx.mesh.EntityMap] | None = None, @@ -802,6 +1041,13 @@ def __init__( self._u = resolve_u(u, F) # type: ignore[arg-type] self._bcs = [] if bcs is None else bcs + + # A caller-omitted prefix must still be unique per Problem (see + # _PROBLEM_PREFIX_COUNTER) so PETSc's process-global options database + # never lets two Problems' SNES/KSP options bleed into each other. + if petsc_options_prefix is None: + petsc_options_prefix = f"dxa_nonlinear_problem_{next(_PROBLEM_PREFIX_COUNTER)}_" + # The user's own J, kept only to scan for dependency coefficients that # might appear in a hand-supplied Jacobian but not in F itself (e.g. a # stabilization term); the Jacobian _get_or_build_dFdu_template uses @@ -811,10 +1057,10 @@ def __init__( # colliding with it. self._user_J = J self._rhs = F + self._petsc_options = petsc_options self._jit_options = jit_options 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 @@ -840,18 +1086,18 @@ def __init__( # 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) + coefficients = collect_coefficients(F) - set(u_list) if J is not None: - coefficients |= _collect_coefficients(J) - set(u_list) + 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=_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] + F=recursive_replace(F, self._value_placeholders), # type: ignore[arg-type] + J=recursive_replace(J, self._value_placeholders), # type: ignore[arg-type] + P=recursive_replace(P, self._value_placeholders), # type: ignore[arg-type] bcs=self._bcs, u=self._u, # type: ignore[arg-type] kind=kind, # type: ignore[arg-type] @@ -871,11 +1117,11 @@ def __init__( def bcs(self) -> typing.Sequence[dolfinx.fem.DirichletBC]: """Dirichlet boundary conditions applied to the residual and Jacobian. - ``dolfinx.fem.petsc.NonlinearProblem`` has no ``bcs`` attribute of its + {py:class}`dolfinx.fem.petsc.NonlinearProblem` has no ``bcs`` attribute of its own (its SNES callbacks close over a fixed ``bcs`` list at - construction, see the note in ``__init__``); this property exposes - ``self._bcs`` under the same name ``LinearProblem`` uses (there, it is - the base class's own attribute), so ``_ProblemBase``'s shared methods + construction, see the note in {py:meth}`~dolfinx_adjoint.NonlinearProblem.__init__`); this property exposes + ``self._bcs`` under the same name {py:class}`~dolfinx_adjoint.LinearProblem` uses (there, it is + the base class's own attribute), so {py:class}`~dolfinx_adjoint.solvers._ProblemBase`'s shared methods can read/write ``self.bcs`` uniformly across both classes. """ return self._bcs @@ -891,12 +1137,16 @@ def _get_or_build_residual_template( and u itself replaced by a dedicated "state" placeholder standing in for "u at this evaluation point", distinct from the live ``self._u`` the forward SNES path owns. - The shared basis for ``dF/du`` (``_ProblemBase._get_or_build_dFdu_template``) and the - TLM right-hand side (``_get_or_build_tlm_rhs_templates``): refreshed from a block's own - checkpointed output before each adjoint/TLM/Hessian solve (see - ``NonlinearProblemBlock._refresh_dFdu_state``/``prepare_evaluate_tlm``), keeping this - template fixed for the life of the Problem, exactly like the non-u dependencies already - routed through ``self._value_placeholders``. + The shared basis for ``dF/du`` + ({py:meth}`_ProblemBase._get_or_build_dFdu_template`) + and the TLM right-hand side + ({py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_tlm_rhs_templates`): + refreshed from a block's own checkpointed output before each + adjoint/TLM/Hessian solve (see + {py:meth}`NonlinearProblemBlock._refresh_dFdu_state`/ + {py:meth}`~dolfinx_adjoint.blocks.solvers._ProblemBlockBase.prepare_evaluate_tlm`), + keeping this template fixed for the life of the Problem, exactly like + the non-u dependencies already routed through ``self._value_placeholders``. """ if self._residual_template is None: u_list = self._u if isinstance(self._u, list) else [self._u] @@ -928,6 +1178,11 @@ def _make_block(self) -> NonlinearProblemBlock: form_compiler_options=self._form_compiler_options, jit_options=self._jit_options, entity_maps=self._entity_maps, + kind=self._kind, + petsc_options=self._petsc_options, + petsc_options_prefix=self._petsc_options_prefix, + adjoint_petsc_options=self._adj_options, + tlm_petsc_options=self._tlm_options, ad_block_tag=self.ad_block_tag, problem=self, ) # type: ignore[misc] @@ -935,7 +1190,7 @@ def _make_block(self) -> NonlinearProblemBlock: def _dolfinx_solve(self) -> _Function | typing.Sequence[_Function]: return dolfinx.fem.petsc.NonlinearProblem.solve(self) - def solve(self, annotate: bool = True) -> typing.Union[dolfinx.fem.Function, typing.Sequence[dolfinx.fem.Function]]: + def solve(self, annotate: bool = True) -> dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function]: """ Solve the nonlinear problem and return the solution. """ diff --git a/src/dolfinx_adjoint/typing_utils.py b/src/dolfinx_adjoint/typing_utils.py new file mode 100644 index 0000000..1ec194e --- /dev/null +++ b/src/dolfinx_adjoint/typing_utils.py @@ -0,0 +1,6 @@ +from __future__ import annotations + +import typing + +type NestedSequence[T] = T | typing.Sequence["NestedSequence[T]"] +type NestedMutableSequence[T] = T | typing.MutableSequence["NestedMutableSequence[T]"] diff --git a/src/dolfinx_adjoint/ufl_utils.py b/src/dolfinx_adjoint/ufl_utils.py new file mode 100644 index 0000000..27e6be9 --- /dev/null +++ b/src/dolfinx_adjoint/ufl_utils.py @@ -0,0 +1,221 @@ +from __future__ import annotations + +import typing + +import ufl + +from .compat import compute_form_adjoint +from .typing_utils import NestedSequence + + +@typing.overload +def assign_mixed_parts[T: NestedSequence[ufl.Form]](form1: T, /) -> T: ... +@typing.overload +def assign_mixed_parts[T: NestedSequence[ufl.Form], S: NestedSequence[ufl.Form]]( + form1: T, form2: S, / +) -> tuple[T, S]: ... +def assign_mixed_parts( + *form_structs: NestedSequence[ufl.Form], +) -> NestedSequence[ufl.Form] | tuple[NestedSequence[ufl.Form], ...]: + """ + Recursively assigns mixed-space `part` indices to {py:class}`ufl.Argument` + (test and trial functions), within nested iterables of forms. + + When solving monolithic block systems in FEniCSx, the UFL arguments must have the + method {py:meth}`ufl.Argument.part` return the index corresponding to their block position. + For a block matrix (list of lists), the TestFunction corresponds to the row index, and the + TrialFunction corresponds to the column index. + + This utility traverses arbitrary nested structures (e.g., a 2D list for the LHS + matrix `a` and a 1D list for the RHS vector `L` simultaneously), extracts arguments + that lack a part index, builds a unified replacement map, and applies it. + + Args: + *form_structs: One or more UFL forms, or nested iterables (lists/tuples) of + UFL forms. Passing multiple structures (like `a` and `L`) ensures they + share the same replacement map, preventing mismatched compilation. + + Returns: + 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 {py:func}`ufl.TestFunctions` + and {py:func}`ufl.TrialFunctions` of a single + {py:class}`ufl.MixedFunctionSpace` built from the row/column function spaces + discovered while walking the structure. + """ + 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. + + `indices` will be `(row,)` for vectors and `(row, col)` for matrices. + """ + if isinstance(obj, ufl.Form): + for arg in obj.arguments(): + if arg.part() is None: + # The argument number corresponds to the index of the row/column + # in the nested structure + num = arg.number() + if num < len(indices): + 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,)) + else: + raise TypeError(f"Expected ufl.Form or iterable, got {type(obj)}") + + 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: + _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, + strictly preserving original sequence types (lists vs. tuples). + """ + if isinstance(obj, ufl.Form): + return ufl.replace(obj, replace_map) + elif isinstance(obj, (list, tuple)): + return type(obj)(_replace(item) for item in obj) + return obj + + # 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] + + +def get_sorted_arguments(arguments: typing.Iterable[ufl.Argument], number: int) -> typing.Iterable[ufl.Argument]: + """Extract all arguments of a given number, sorted by part.""" + return sorted(filter(lambda x: x.number() == number, arguments), key=lambda a: a.part()) + + +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 a blocked ``NonlinearProblem``'s forms are not + required to be before ``assign_mixed_parts`` runs. + """ + 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 sum_form(form: NestedSequence[ufl.Form | None]) -> ufl.Form | None: + """Sum a blocked form into a single form.""" + # Handle top-level None + if form is None: + return None + + if isinstance(form, ufl.Form): + return form + + elif isinstance(form, typing.Iterable): + # Recursively sum items, filtering out Nones + valid_forms: list[ufl.Form] = [] + for fi in form: + summed_fi = sum_form(fi) + if summed_fi is not None: + valid_forms.append(summed_fi) + + # Handle empty case safely + if not valid_forms: + return None + + # Safely sum without defaulting to integer 0, removing the need for type: ignore + return sum(valid_forms[1:], start=valid_forms[0]) + + else: + raise TypeError(f"Cannot sum form of type {type(form)}") + + +def compute_adjoint(form: ufl.Form) -> typing.Sequence[typing.Sequence[ufl.Form]] | ufl.Form: + """Compute the adjoint of a (possibly blocked) bilinear form. + + A module-level function, not a method: it needs no ``Block``/``Problem`` state, + just ``form`` itself, so both ``_ProblemBlockBase`` (``blocks/solvers.py``) and + ``_ProblemBase`` (``solvers.py``) can call it directly rather than one reaching + into a "private" method defined on the other. + + Args: + form: A bilinear form :math:`a(u, v)`, either a single ``ufl.Form`` or a + blocked (nested list) system. + + Returns: + The transposed form :math:`a(v, u)`, decomposed back into blocks (via + ``ufl.extract_blocks``) -- a no-op decomposition for a scalar form. + """ + return ufl.extract_blocks(compute_form_adjoint(form)) + + +def recursive_replace(form: ufl.Form | typing.Sequence | None, placeholders: dict) -> ufl.Form | typing.Sequence | None: + """Recursively apply {py:func}`ufl.replace` 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 {py:class}`~pyadjoint.Block`) exists to avoid: a + self-referential ``_replace`` closure inside a + {py:class}`~dolfinx_adjoint.LinearProblem`/{py:class}`~dolfinx_adjoint.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. + + Args: + form: A single form, ``None``, or an arbitrarily nested sequence of + forms/``None`` (e.g. a blocked system). + placeholders: Map passed straight through to {py:func}`ufl.replace` at + each form encountered. + + Returns: + A structure with the same nesting as ``form``, each form replaced via + {py:func}`ufl.replace`; ``None`` in, ``None`` out. + """ + if form is None: + return None + if isinstance(form, ufl.Form): + return ufl.replace(form, placeholders) + return [recursive_replace(f, placeholders) for f in form] diff --git a/tests/test_solver_reuse.py b/tests/test_solver_reuse.py index ce91e34..b81738f 100644 --- a/tests/test_solver_reuse.py +++ b/tests/test_solver_reuse.py @@ -16,6 +16,7 @@ import dolfinx import numpy as np import pyadjoint +import pytest import ufl from dolfinx_adjoint import Function, LinearProblem, NonlinearProblem, assemble_scalar, assign @@ -557,7 +558,7 @@ def test_linear_problem_released_by_refcounting_not_gc(): 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 + (``dolfinx_adjoint.ufl_utils.recursive_replace``) a plain module-level function taking the placeholder dict as an explicit argument, so it does not need to capture itself or ``self``. """ @@ -647,3 +648,159 @@ def test_nonlinear_problem_released_by_refcounting_not_gc(): ) finally: gc.enable() + + +def test_linear_problem_rebuilt_after_garbage_collection(): + """A block only holds a ``weakref`` to its owning LinearProblem (see + LinearProblemBlock._problem/_rebuild_problem), precisely so dropping every + external reference to the Problem releases it immediately, even while blocks + that reference it are still on the tape. Replaying the tape afterwards (e.g. to + differentiate) must still work: the block rebuilds an equivalent LinearProblem on + demand, warning since that is a costly fallback, and the rebuilt Problem must + give the same answer the original would have. + """ + 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", + } + + problem = LinearProblem(a, L, bcs=[bc], u=uh, petsc_options=petsc_options, petsc_options_prefix="dxa_rebuild_test_") + problem.solve() + + J = assemble_scalar(uh * uh * ufl.dx) + control = pyadjoint.Control(m) + Jh = pyadjoint.ReducedFunctional(J, control) + + problem_ref = weakref.ref(problem) + del problem + assert problem_ref() is None, ( + "LinearProblem should be released the instant its last external reference is " + "dropped, since blocks only hold a weakref to it" + ) + + pert = Function(V) + pert.interpolate(lambda x: np.cos(x[1])) + + with pytest.warns(UserWarning, match="LinearProblem was garbage collected"): + min_rate = pyadjoint.taylor_test(Jh, m, pert) + assert np.isclose(min_rate, 2.0, rtol=1e-2, atol=1e-2), f"Expected convergence rate close to 2.0, got {min_rate}" + + +def test_nonlinear_problem_rebuilt_after_garbage_collection(): + """As ``test_linear_problem_rebuilt_after_garbage_collection``, 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", + } + adjoint_options = { + "ksp_type": "preonly", + "pc_type": "lu", + "ksp_error_if_not_converged": True, + "pc_factor_mat_solver_type": "mumps", + } + + problem = NonlinearProblem( + F1, + u=u1, + bcs=[bc], + petsc_options=options, + adjoint_petsc_options=adjoint_options, + petsc_options_prefix="dxa_nonlinear_rebuild_test_", + ) + problem.solve() + + d = pyadjoint.AdjFloat(0.2) + J = assemble_scalar((u1 - d) * (u1 - d) * ufl.dx) + control = pyadjoint.Control(f) + Jh = pyadjoint.ReducedFunctional(J, control) + + problem_ref = weakref.ref(problem) + del problem + assert problem_ref() is None, ( + "NonlinearProblem should be released the instant its last external reference " + "is dropped, since blocks only hold a weakref to it" + ) + + pert = Function(V) + pert.interpolate(lambda x: np.cos(x[1])) + + with pytest.warns(UserWarning, match="NonlinearProblem was garbage collected"): + min_rate = pyadjoint.taylor_test(Jh, f, pert) + assert np.isclose(min_rate, 2.0, rtol=1e-2, atol=1e-2), f"Expected convergence rate close to 2.0, got {min_rate}" + + +def test_default_petsc_options_prefix_is_unique_per_problem(): + """Two Problems of the same kind that don't pass ``petsc_options_prefix`` + explicitly must not resolve to the same default prefix. + + PETSc's options database is process-global and keyed by prefix: two Problems + sharing one would let one instance's SNES/KSP options silently leak into (or get + overwritten by) the other's, exactly the kind of collision this default-prefix + scheme exists to prevent -- see ``dolfinx_adjoint.solvers._PROBLEM_PREFIX_COUNTER``. + """ + mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 2, 2) + V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) + + u, v = ufl.TrialFunction(V), ufl.TestFunction(V) + a = ufl.inner(u, v) * ufl.dx + L = ufl.inner(dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(1.0)), v) * ufl.dx + + linear_a = LinearProblem(a, L) + linear_b = LinearProblem(a, L) + assert linear_a._petsc_options_prefix != linear_b._petsc_options_prefix + + u1 = Function(V, name="state") + F = ufl.inner((1 + u1**2) * u1, v) * ufl.dx - ufl.inner(dolfinx.fem.Constant(mesh, 1.0), v) * ufl.dx + nonlinear_a = NonlinearProblem(F, u=u1) + nonlinear_b = NonlinearProblem(F, u=u1) + assert nonlinear_a._petsc_options_prefix != nonlinear_b._petsc_options_prefix + + # Also unique across the two classes, since they draw from one shared counter. + assert linear_a._petsc_options_prefix != nonlinear_a._petsc_options_prefix + + # An explicitly-passed prefix must be honored unchanged. + explicit = LinearProblem(a, L, petsc_options_prefix="my_custom_prefix_") + assert explicit._petsc_options_prefix == "my_custom_prefix_" diff --git a/tests/test_tlm_update.py b/tests/test_tlm_update.py index f618440..e9b6ed3 100644 --- a/tests/test_tlm_update.py +++ b/tests/test_tlm_update.py @@ -204,8 +204,8 @@ def _navier_stokes(mesh): def test_hessian_is_independent_of_previous_evaluation_points_navier_stokes(mesh_2D): """``test_hessian_is_independent_of_previous_evaluation_points``'s ``NonlinearProblem`` sibling: the same second-order Taylor test, but on a genuinely nonlinear, blocked - (multi-output) residual, exercising the blocked Hessian path - ``_ProblemBlockBase._evaluate_hessian_blocked_rhs``/``_evaluate_hessian_component_blocked`` + (multi-output) residual, exercising the blocked Hessian path in + ``_ProblemBase._get_or_build_hessian_templates``/``_ProblemBlockBase.prepare_evaluate_hessian`` for ``NonlinearProblem`` for the first time -- previously this path only ever ran for ``LinearProblem``. """ From 2b907d5452afea8e7b8cc349115762db1023e8b2 Mon Sep 17 00:00:00 2001 From: Joergen Schartum Dokken Date: Tue, 1 Sep 2026 11:04:51 +0000 Subject: [PATCH 3/8] More cleanup of ufl utils --- src/dolfinx_adjoint/ufl_utils.py | 122 +++++++++++++++++-------------- src/dolfinx_adjoint/utils.py | 16 +++- 2 files changed, 81 insertions(+), 57 deletions(-) diff --git a/src/dolfinx_adjoint/ufl_utils.py b/src/dolfinx_adjoint/ufl_utils.py index 27e6be9..59a28d9 100644 --- a/src/dolfinx_adjoint/ufl_utils.py +++ b/src/dolfinx_adjoint/ufl_utils.py @@ -8,6 +8,66 @@ from .typing_utils import NestedSequence +def recursive_space_discovery( + obj: NestedSequence[ufl.Form], indices: tuple[int, ...], spaces: dict[int, ufl.FunctionSpace] +) -> None: + """Recursively discover, for each row/column index, the function space of the + (as yet unassigned) argument occupying that position. + + `indices` will be `(row,)` for vectors and `(row, col)` for matrices. + + Arguments: + obj: A UFL form or nested iterable of forms. + indices: The current row/column indices in the nested structure. + spaces: A dictionary mapping row/column indices to discovered function spaces. + This dictionary is updated in-place as the function traverses the structure. + """ + if isinstance(obj, ufl.Form): + for arg in obj.arguments(): + if arg.part() is None: + # The argument number corresponds to the index of the row/column + # in the nested structure + num = arg.number() + if num < len(indices): + spaces.setdefault(indices[num], arg.ufl_function_space()) + elif isinstance(obj, typing.Iterable): + for i, item in enumerate(obj): + if item is not None: + recursive_space_discovery(item, indices + (i,), spaces) + else: + raise TypeError(f"Expected ufl.Form or iterable, got {type(obj)}") + + +def build_argument_replacement_map( + obj: NestedSequence[ufl.Form], + indices: tuple[int, ...], + test_functions: typing.Sequence[ufl.TestFunction], + trial_functions: typing.Sequence[ufl.TrialFunction], + replace_map: dict[ufl.Argument, ufl.Argument], +) -> None: + """ + Recursively build a mapping from ufl arguments that does not have a `part`-index + to their replacements in a {py:class}`ufl.MixedFunctionSpace`. + + Arguments: + obj: A UFL form or nested iterable of forms. + indices: The current row/column indices in the nested structure. + test_functions: A sequence of test functions used for replacement. + trial_functions: A sequence of trial functions used for replacement. + replace_map: A dictionary mapping old arguments to new arguments. + """ + 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: + build_argument_replacement_map(item, indices + (i,), test_functions, trial_functions, replace_map) + + @typing.overload def assign_mixed_parts[T: NestedSequence[ufl.Form]](form1: T, /) -> T: ... @typing.overload @@ -36,9 +96,9 @@ def assign_mixed_parts( share the same replacement map, preventing mismatched compilation. Returns: - 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. + The modified form structures with identical nesting, 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 {py:func}`ufl.TestFunctions` @@ -47,30 +107,8 @@ def assign_mixed_parts( discovered while walking the structure. """ 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. - - `indices` will be `(row,)` for vectors and `(row, col)` for matrices. - """ - if isinstance(obj, ufl.Form): - for arg in obj.arguments(): - if arg.part() is None: - # The argument number corresponds to the index of the row/column - # in the nested structure - num = arg.number() - if num < len(indices): - 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,)) - else: - raise TypeError(f"Expected ufl.Form or iterable, got {type(obj)}") - for struct in form_structs: - _discover_spaces(struct, ()) + recursive_space_discovery(struct, (), spaces) # If no replacements are needed, exit early to save computation if not spaces: @@ -81,36 +119,12 @@ def _discover_spaces(obj: NestedSequence[ufl.Form], indices: tuple[int, ...]) -> 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: - _build_map(item, indices + (i,)) - + replace_map: dict[ufl.Argument, ufl.Argument] = {} for struct in form_structs: - _build_map(struct, ()) - - def _replace(obj: typing.Any) -> typing.Any: - """ - Recursively rebuild the structure using the populated replace_map, - strictly preserving original sequence types (lists vs. tuples). - """ - if isinstance(obj, ufl.Form): - return ufl.replace(obj, replace_map) - elif isinstance(obj, (list, tuple)): - return type(obj)(_replace(item) for item in obj) - return obj + build_argument_replacement_map(struct, (), test_functions, trial_functions, replace_map) # Apply the replacements and unpack if necessary - replaced = tuple(_replace(struct) for struct in form_structs) + replaced = tuple(recursive_replace(struct, replace_map) for struct in form_structs) return replaced if len(replaced) > 1 else replaced[0] @@ -119,7 +133,7 @@ def get_sorted_arguments(arguments: typing.Iterable[ufl.Argument], number: int) return sorted(filter(lambda x: x.number() == number, arguments), key=lambda a: a.part()) -def collect_coefficients(form: ufl.Form | typing.Sequence | None) -> set: +def collect_coefficients(form: ufl.Form | typing.Sequence | None) -> set[ufl.Coefficient]: """Return the set of UFL coefficients appearing anywhere in ``form``. ``form`` may be a single form or an arbitrarily nested sequence of forms diff --git a/src/dolfinx_adjoint/utils.py b/src/dolfinx_adjoint/utils.py index 6b0b01f..a9bab8e 100644 --- a/src/dolfinx_adjoint/utils.py +++ b/src/dolfinx_adjoint/utils.py @@ -25,8 +25,12 @@ def function_from_vector( ) -> dolfinx.fem.Function: """Create a new Function from a vector. - :arg V: The function space - :arg vector: The vector data. + Arguments: + V: The function space + vector: The vector data. + Returns: + A new {py:class}`dolfinx.fem.Function` instance that has been assigned the + values from the vector (deep-copy) """ ret = dolfinx.fem.Function(V, dtype=vector.array.dtype) ret.x.array[:] = vector.array[:] @@ -34,7 +38,13 @@ def function_from_vector( def gather(vector: dolfinx.la.Vector) -> npt.NDArray[numpy.number]: - """Gather a vector on all processes.""" + """Gather a vector on all processes. + + Args: + vector: The vector to gather. + Returns: + A numpy array containing the gathered vector. + """ local_size = vector.index_map.size_local * vector.block_size comm = vector.index_map.comm data = comm.allgather(vector.array[:local_size]) From 99fabd1c666ef5347eb3cf44a7bb4cf84fcbf521 Mon Sep 17 00:00:00 2001 From: Joergen Schartum Dokken Date: Tue, 1 Sep 2026 11:46:01 +0000 Subject: [PATCH 4/8] Simplify docstrings. Clean up naming. --- src/dolfinx_adjoint/solvers.py | 502 ++++++++++++++------------------- 1 file changed, 209 insertions(+), 293 deletions(-) diff --git a/src/dolfinx_adjoint/solvers.py b/src/dolfinx_adjoint/solvers.py index 88f2567..b49a639 100644 --- a/src/dolfinx_adjoint/solvers.py +++ b/src/dolfinx_adjoint/solvers.py @@ -21,26 +21,25 @@ sum_form, ) -# Backs each Problem's default PETSc options prefix (see LinearProblem/NonlinearProblem -# __init__). A plain incrementing counter, not id(self)/uuid.uuid4()/time.time(): those -# can differ across MPI ranks for the same logical Problem (memory layout, clock skew), -# and PETSc's options-database handling around SNESSetFromOptions/KSPSetFromOptions is -# collective, so every rank must resolve the same prefix for the same Problem. A counter -# incremented once per Problem construction is deterministic and identical on every rank, -# since construction happens in lock-step in a well-formed SPMD program. +# A counter incremented once per Problem construction is deterministic +# and identical on every rank, since construction happens in lock-step +# in a well-formed SPMD program. _PROBLEM_PREFIX_COUNTER = itertools.count() @typing.overload -def resolve_u(u: _Function | None, L: ufl.Form) -> _Function: ... +def find_or_create_then_overload(u: _Function | None, L: ufl.Form) -> _Function: ... @typing.overload -def resolve_u(u: typing.Sequence[_Function] | None, L: typing.Sequence[ufl.Form]) -> typing.Sequence[_Function]: ... +def find_or_create_then_overload( + u: typing.Sequence[_Function] | None, L: typing.Sequence[ufl.Form] +) -> typing.Sequence[_Function]: ... -def resolve_u( +def find_or_create_then_overload( u: _Function | typing.Sequence[_Function] | None, L: ufl.Form | typing.Sequence[ufl.Form] ) -> _Function | typing.Sequence[_Function]: - """Resolve the unknown {py:class}`~dolfinx_adjoint.Function` ``u`` for a `*Problem`. + """Find or create, then overload the unknown + {py:class}`~dolfinx_adjoint.Function` ``u`` for a `*Problem`. If ``u`` was not supplied by the caller, a fresh {py:class}`~dolfinx_adjoint.Function` is created per block, using the function space of the corresponding @@ -176,21 +175,18 @@ def _build_soa_self_template( class _ProblemBase(abc.ABC): """Shared lazy adjoint/TLM solver machinery for {py:class}`~dolfinx_adjoint.LinearProblem`/{py:class}`~dolfinx_adjoint.NonlinearProblem`. - - A plain mixin -- it does not inherit from any ``dolfinx.fem.petsc`` class, - so ``LinearProblem(_ProblemBase, dolfinx.fem.petsc.LinearProblem)`` (and - the {py:class}`~dolfinx_adjoint.NonlinearProblem` equivalent) has an - unambiguous MRO for ``solve()``: each subclass keeps defining its own - ``solve()``, delegating the shared middle to - {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._record_and_solve` below via - the {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._make_block`/ - {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._dolfinx_solve` hooks it implements. - - Every attribute referenced here is set by the concrete subclass's own - ``__init__`` (either directly, or inherited from the ``dolfinx.fem.petsc`` - base it also derives from) before any of these methods run. """ + # A plain mixin -- not a subclass of any dolfinx.fem.petsc class -- so + # LinearProblem(_ProblemBase, dolfinx.fem.petsc.LinearProblem) (and the + # NonlinearProblem equivalent) gets an unambiguous MRO for solve(): each + # subclass keeps defining its own solve(), delegating the shared middle to + # _record_and_solve() below via the _make_block()/_dolfinx_solve() hooks it + # implements. + + # Every attribute below is set by the concrete subclass's own __init__ + # (either directly, or inherited from the dolfinx.fem.petsc base it also + # derives from) before any method on this class runs. ad_block_tag: str | None bcs: typing.Sequence[dolfinx.fem.DirichletBC] _u: _Function | typing.Sequence[_Function] @@ -207,18 +203,10 @@ class _ProblemBase(abc.ABC): @property def value_placeholders(self) -> dict[dolfinx.fem.Function, dolfinx.fem.Function]: - """Map from each non-``u`` dependency the compiled forward/adjoint/TLM/Hessian - forms were built against to its dedicated placeholder coefficient. - - Public so - {py:class}`*ProblemBlock` - methods (a different module) can refresh a dependency's placeholder value - ahead of a solve/replay without reaching into this Problem's own private - state. - - Returns: - The placeholder dictionary, keyed by the user's original dependency. - """ + """Map from each non-``u`` dependency to its dedicated placeholder coefficient.""" + # Public so *ProblemBlock (blocks/solvers.py) can refresh a dependency's + # placeholder value ahead of a solve/replay without reaching into this + # Problem's private state. return self._value_placeholders @property @@ -226,76 +214,52 @@ def residual_state_placeholder(self) -> dolfinx.fem.Function | typing.Sequence[d """The dedicated "state" placeholder(s) standing in for ``u`` in every compiled template built from the residual (see {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_residual_template`). - - Returns: - A single {py:class}`~dolfinx_adjoint.Function` for a scalar problem, or one - per output block for a blocked problem. + A single Function for a scalar problem, or one per output block for a + blocked problem. """ assert self._residual_state_placeholder is not None return self._residual_state_placeholder @property def adjoint_solution_placeholder(self) -> dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function]: - """The placeholder(s) standing in for the first-order adjoint solution in the - cached Hessian templates (see {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_hessian_templates`). - - Returns: - A single {py:class}`~dolfinx_adjoint.Function` for a scalar problem, or one - per output block for a blocked problem. ``None`` until the Hessian templates have been built. + """The placeholder(s) for the first-order adjoint solution in the cached + Hessian templates (see + {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_hessian_templates`). """ assert self._adjoint_solution_placeholder is not None return self._adjoint_solution_placeholder @property def second_adjoint_solution_placeholder(self) -> dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function]: - """The placeholder(s) standing in for the second-order adjoint (SOA) solution - in the cached Hessian templates (see + """The placeholder(s) for the second-order adjoint (SOA) solution in the + cached Hessian templates (see {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_hessian_templates`). - - Returns: - A single {py:class}`~dolfinx_adjoint.Function` for a scalar problem, or one - per output block for a blocked problem. """ assert self._second_adjoint_solution_placeholder is not None return self._second_adjoint_solution_placeholder @property def hessian_u_seed(self) -> dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function]: - """The placeholder(s) standing in for the state's own tangent-linear direction - in the cached Hessian self-term (see + """The placeholder(s) for the state's own tangent-linear direction in the + cached Hessian self-term (see {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_hessian_templates`). - - Returns: - A single {py:class}`~dolfinx_adjoint.Function` for a scalar problem, or one - per output block for a blocked problem. """ assert self._hessian_u_seed is not None return self._hessian_u_seed def _init_adjoint_state(self) -> None: - """Initialize the lazily-built adjoint/TLM solver state. - - Called once, at the end of ``__init__``, after the base - ``dolfinx.fem.petsc`` solver has been constructed. Built lazily (on - first use, see - {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_adjoint_solver`/ - {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_tlm_solver`/ - {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_hessian_templates`) - and shared by every block this Problem records, rather than one per - block/solve() call. Each block holds only a ``weakref`` back to this - Problem (see - {py:meth}`*ProblemBlock.get_reference_problem`), - not a strong reference: a throwaway Problem (solve it, then only touch - the resulting {py:class}`~pyadjoint.ReducedFunctional`) must not be - kept alive for as long as its blocks remain reachable on the tape. If - this Problem is collected while a block still needs it, that block - rebuilds an equivalent one on demand - ({py:meth}`*ProblemBlock._rebuild_problem`, - with a {py:class}`UserWarning` since it is a costly fallback) rather - than silently failing. Laziness keeps pure forward (non-annotated) - use from paying for a symbolic - adjoint form it never needs. - """ + """Initialize the lazily-built adjoint/TLM solver state.""" + # Called once, at the end of __init__, after the base dolfinx.fem.petsc + # solver is constructed. Everything below is built lazily, on first use + # (see _get_or_build_adjoint_solver/_get_or_build_tlm_solver/ + # _get_or_build_hessian_templates), so pure forward (non-annotated) use + # never pays for a symbolic adjoint form it doesn't need. + # + # Shared by every block this Problem records rather than rebuilt per + # block/solve() call. Each block holds only a weakref back here, not a + # strong reference -- see *ProblemBlock.get_reference_problem + # (blocks/solvers.py) for why, and how a block copes if this Problem is + # collected before it's done needing it. self._adjoint_solver: HomogeneousBCLinearProblem | None = None self._tlm_solver: HomogeneousBCLinearProblem | None = None self._residual_template: ufl.Form | None = None @@ -315,42 +279,33 @@ def _init_adjoint_state(self) -> None: def _get_or_build_residual_template( self, ) -> tuple[ufl.Form, dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function]]: - """Build (once) and return F with every coefficient replaced by its placeholder, and u - replaced by a dedicated "state" placeholder standing in for "u at this evaluation point". + """Build (once) and return F with every coefficient replaced by its placeholder, + and u replaced by a dedicated "state" placeholder for "u at this evaluation point". - The one genuinely irreducible difference between the two Problem kinds -- LinearProblem - builds it from ``a``/``L`` via {py:func}`ufl.action`, NonlinearProblem already has ``F`` - directly -- everything built on top of it below (``dF/du``, the TLM right-hand side, the - Hessian templates) is derived from this one template by the same shared symbolic - differentiation, since that costs nothing extra at compile time: there is no reason to - keep a separate "dF/du is just a, never differentiated" shortcut for LinearProblem. - Each subclass overrides this. + Each subclass overrides this -- the one genuinely irreducible difference + between the two Problem kinds (LinearProblem builds it from ``a``/``L`` + via {py:func}`ufl.action`; NonlinearProblem already has ``F`` directly). + Everything derived below (``dF/du``, TLM right-hand side, Hessian + templates) shares this one template, via the same symbolic + differentiation for both kinds. Returns: A ``(F_template, state_placeholder)`` pair: the residual with every - coefficient (including ``u``) substituted by its dedicated - placeholder, and that state placeholder itself (a single - {py:class}`~dolfinx_adjoint.Function`, or one per output block for a - blocked problem). + coefficient (including ``u``) substituted by its placeholder, and + that state placeholder itself (single, or one per output block for + a blocked problem). """ def _get_or_build_dFdu_template(self) -> ufl.Form | typing.Sequence: - """Build (once) and return dF/du, evaluated at the residual template's state placeholder. - - Shared by both classes: derived from - {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_residual_template` - by symbolic differentiation, which is free at compile time, rather than - special-cased per class. This is the shared basis for the adjoint operator - ({py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_adjoint_solver`, - which just adjoints it) and the TLM operator - ({py:meth}`~dolfinx_adjoint.solvers._ProblemBase._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 - {py:meth}`*ProblemBlock.prepare_evaluate_adj`/ - {py:meth}`~dolfinx_adjoint.blocks.solvers._ProblemBlockBase.prepare_evaluate_hessian`/ - {py:meth}`~dolfinx_adjoint.blocks.solvers._ProblemBlockBase.prepare_evaluate_tlm`). - """ + """Build (once) and return dF/du, evaluated at the residual template's state placeholder.""" + # Shared by both classes: derived from _get_or_build_residual_template by + # symbolic differentiation (free at compile time) rather than + # special-cased per class. Basis for 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 -- callers only ever refresh the placeholders' + # values afterwards (see *ProblemBlock.prepare_evaluate_adj/ + # prepare_evaluate_hessian/prepare_evaluate_tlm in blocks/solvers.py). if self._dFdu_template is None: F_template, state_placeholder = self._get_or_build_residual_template() if isinstance(self._u, list): @@ -370,26 +325,21 @@ def _get_or_build_dFdu_template(self) -> ufl.Form | typing.Sequence: 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 - ({py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_adjoint_solver`) - and, for scalar problems, the Hessian SOA right-hand-side's - cross-dependency templates - ({py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_hessian_templates`). - - The same computation for both classes: - {py:func}`~dolfinx_adjoint.ufl_utils.compute_adjoint` swaps argument - numbers while preserving mixed-space {py:meth}`ufl.Argument.part` tags - (via {py:func}`~dolfinx_adjoint.compat.compute_form_adjoint`) and - decomposes the result back into blocks (via - {py:func}`ufl.extract_blocks`) -- a no-op decomposition for a scalar, - non-blocked form. Kept exactly as that returns it (a nested list of - forms for a blocked problem) since that structure is what - {py:class}`~dolfinx_adjoint.petsc_utils.HomogeneousBCLinearProblem`/ - {py:class}`dolfinx.fem.petsc.LinearProblem` needs for block matrix - assembly; callers that need a single summed form (Hessian templating, - scalar-only) apply {py:func}`~dolfinx_adjoint.ufl_utils.sum_form` themselves. + """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`). """ if self._dFdu_adj_template is None: + # compute_adjoint() (ufl_utils.py) swaps argument numbers while + # preserving mixed-space ufl.Argument.part tags, then decomposes back + # into blocks via ufl.extract_blocks -- a no-op for a scalar form. + # Kept exactly as that returns it (a nested list of forms for a + # blocked problem), since that's the shape HomogeneousBCLinearProblem/ + # dolfinx.fem.petsc.LinearProblem need for block matrix assembly; + # callers wanting a single summed form (Hessian templating, + # scalar-only) apply ufl_utils.sum_form() themselves. self._dFdu_adj_template = compute_adjoint( self._get_or_build_dFdu_template() # type: ignore[arg-type] ) @@ -402,29 +352,12 @@ def _get_or_build_tlm_rhs_templates( dict[dolfinx.fem.Function, dolfinx.fem.Function], dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function], ]: - """Build (once) and return the per-dependency TLM right-hand-side templates. - - Shared by both classes: built purely from the residual template - ({py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_residual_template`), - since ``dF/dm`` genuinely depends on the state for either a linear or - nonlinear residual -- differentiating w.r.t. a coefficient embedded in - the residual while holding ``u`` fixed leaves ``u`` in the result even - when the residual is linear in ``u`` itself. - - 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 - {py:meth}`*ProblemBlock.prepare_evaluate_tlm`), - avoids that entirely. - """ + """Build (once) and return the per-dependency TLM right-hand-side templates.""" + # Shared by both classes: built purely from the residual template + # (_get_or_build_residual_template), since dF/dm genuinely depends on the + # state for either a linear or nonlinear residual -- differentiating + # w.r.t. a coefficient embedded in the residual while holding u fixed + # leaves u in the result even when the residual is linear in u itself. F_template, state_placeholder = self._get_or_build_residual_template() if self._tlm_rhs_templates is None: if isinstance(self._u, list): @@ -433,6 +366,20 @@ def _get_or_build_tlm_rhs_templates( test_funcs = [F_template.arguments()[0]] templates: dict[dolfinx.fem.Function, typing.Any] = {} + # One compiled one-form per dependency, using a dedicated "direction" + # placeholder (below, cached in _tlm_seed_placeholders) rather than a + # single form summed over every dependency: summing symbolically + # would require deciding, once and for all, which dependencies + # contribute, but that varies from call to call. Zeroing an unused + # dependency's seed and evaluating its term anyway isn't a safe + # substitute for skipping it -- if that dependency appears somewhere + # singular at its current value (e.g. a 1/c term with c legitimately + # zero somewhere), the assembled contribution is 0 * inf = NaN there + # even though the seed is zero, silently corrupting the sum. Keeping + # each dependency as its own compiled form, only ever assembled when + # it actually has a tangent-linear value (see + # *ProblemBlock.prepare_evaluate_tlm in blocks/solvers.py), avoids + # that entirely. 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)) @@ -452,45 +399,14 @@ def _get_or_build_tlm_rhs_templates( return self._tlm_rhs_templates, self._tlm_seed_placeholders, state_placeholder def _get_or_build_hessian_templates(self) -> HessianTemplates: - """Build (once) and return the per-dependency Hessian templates used by - {py:meth}`*ProblemBlock.prepare_evaluate_hessian`'s - SOA right-hand side and - {py:meth}`~dolfinx_adjoint.blocks.solvers._ProblemBlockBase.evaluate_hessian_component`'s - own Hessian-action output. - - Shared by both classes and by both scalar and blocked problems -- built - purely from - {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_residual_template`/ - {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_dFdu_template`/ - {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_dFdu_adj_template`/ - {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._get_or_build_tlm_rhs_templates`, - all themselves shared. ``soa_self`` (see - {py:func}`~dolfinx_adjoint.solvers._build_soa_self_template`) comes out a - {py:class}`ufl.ZeroBaseForm` for {py:class}`~dolfinx_adjoint.LinearProblem` - (``dF/du`` doesn't depend on ``u``) and generally nonzero for - {py:class}`~dolfinx_adjoint.NonlinearProblem`, as a *result* of running the - same code, not a per-class branch. - - ``fixed``/``cross`` are one-forms over a *control's own* test space - (``c.function_space``), so their shape doesn't depend on how many output - blocks ``u`` has -- {py:func}`ufl.action` already reduces the (possibly - part-tagged, blocked) state/adjoint-solution arguments away before - ``fixed``/``cross`` are built. Only ``soa_self``/``soa_cross`` feed the - (possibly blocked) SOA right-hand-side vector, so for a blocked problem - they become a list of one compiled form per output row (padded via - {py:func}`~dolfinx_adjoint.solvers._pad_blocks_by_part` for any row a - differentiation happened to eliminate entirely) instead of a single form. - - Each of ``soa_cross``/``fixed``/``cross`` is kept as its own compiled - one-form (or list of one-forms), 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 - {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._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. + """Build (once) and return the per-dependency Hessian templates. + + Feeds *ProblemBlock.prepare_evaluate_hessian's SOA right-hand side and + evaluate_hessian_component's Hessian-action output (blocks/solvers.py). + Shared by both classes and by scalar and blocked problems, built + entirely from the other cached templates (`_get_or_build_residual_template`, + `_get_or_build_dFdu_template`, `_get_or_build_dFdu_adj_template`, + `_get_or_build_tlm_rhs_templates`). """ if self._hessian_templates is None: _, seed_placeholders, state_placeholder = self._get_or_build_tlm_rhs_templates() @@ -503,6 +419,10 @@ def _get_or_build_hessian_templates(self) -> HessianTemplates: blocked = isinstance(self._u, list) soa_self: dolfinx.fem.Form | list[dolfinx.fem.Form] if blocked: + # One placeholder Function per output block, mirroring how + # _get_or_build_hessian_templates's scalar branch below uses a + # single one -- these back the adjoint_solution_placeholder/ + # second_adjoint_solution_placeholder/hessian_u_seed properties. assert isinstance(state_placeholder, typing.Sequence) state_list = list(state_placeholder) test_funcs = list(get_sorted_arguments(F_template.arguments(), 0)) @@ -511,6 +431,13 @@ def _get_or_build_hessian_templates(self) -> HessianTemplates: self._hessian_u_seed = [dolfinx.fem.Function(s.function_space) for s in state_list] state_arg: typing.Any = state_list + # soa_self = adjoint(d2F/du2) . adjoint_solution -- the SOA + # right-hand side's contribution from dF/du's own second + # derivative w.r.t. u. d2Fdu2 is structurally zero whenever the + # residual is linear in u (dF/du doesn't reference u), so this + # comes out a ufl.ZeroBaseForm for LinearProblem as a *result* + # of running the same code, not a per-class branch (see + # _build_soa_self_template, used below for the scalar case). d2Fdu2 = ufl.algorithms.expand_derivatives( ufl.derivative(dFdu_template, state_list, self._hessian_u_seed) ) @@ -518,6 +445,10 @@ def _get_or_build_hessian_templates(self) -> HessianTemplates: soa_self_form = d2Fdu2 else: soa_self_form = ufl.action(ufl.adjoint(d2Fdu2), self._adjoint_solution_placeholder) + # soa_self feeds the (blocked) SOA right-hand-side vector, so it + # becomes a list of one compiled form per output row here, + # padded via _pad_blocks_by_part for any row a differentiation + # happened to eliminate entirely. soa_self = [ dolfinx.fem.form( form_i, @@ -544,6 +475,12 @@ def _get_or_build_hessian_templates(self) -> HessianTemplates: entity_maps=self._entity_maps, ) + # dFdu_adj_applied/L1/L2 are shared building blocks for every + # dependency's soa_cross/fixed/cross templates below: dFdu_adj_applied + # is dF/du^T applied to the first-order adjoint solution (the base + # for each dependency's soa_cross term), L1/L2 are the residual + # applied to the first/second-order adjoint solutions respectively + # (the base for each dependency's fixed/cross terms). 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) @@ -551,9 +488,20 @@ def _get_or_build_hessian_templates(self) -> HessianTemplates: soa_cross_templates: dict = {} fixed_templates: dict = {} cross_templates: dict = {} + # fixed/cross live on each control's own test space, so their shape + # is unaffected by blocking. soa_self/soa_cross do feed the + # (possibly blocked) SOA right-hand side, so for a blocked problem + # each becomes a list of one compiled form per output row (padded + # via _pad_blocks_by_part for any row a differentiation eliminated). + # + # Each cross-term below uses its own dedicated seed_placeholders + # direction placeholder rather than one combined form, for the same + # 0 * inf = NaN reason as _get_or_build_tlm_rhs_templates. for c, c_placeholder in self._value_placeholders.items(): seed = seed_placeholders[c] + # soa_cross[c]: the SOA right-hand side's contribution from c's + # own tangent-linear direction, via dFdu_adj_applied. soa_form = ufl.algorithms.expand_derivatives(ufl.derivative(dFdu_adj_applied, c_placeholder, seed)) if not (soa_form == 0 or soa_form.empty()): if blocked: @@ -574,6 +522,10 @@ def _get_or_build_hessian_templates(self) -> HessianTemplates: entity_maps=self._entity_maps, ) + # fixed[c]: c's own Hessian-action contribution that does not + # depend on any *other* dependency's tangent-linear value -- + # the second-order-adjoint term (dL2dm, from L2) plus the + # mixed state/control second derivative (d2Fdudm, from L1). dc = ufl.TestFunction(c.function_space) dL1dm = ufl.derivative(L1, c_placeholder, dc) dL2dm = ufl.derivative(L2, c_placeholder, dc) @@ -588,6 +540,8 @@ def _get_or_build_hessian_templates(self) -> HessianTemplates: entity_maps=self._entity_maps, ) + # cross[(c, c2)]: c's Hessian-action contribution from another + # dependency c2's tangent-linear direction, reusing dL1dm. 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)) @@ -620,27 +574,23 @@ def _get_or_build_adjoint_solver(self) -> HomogeneousBCLinearProblem: return self._adjoint_solver def _get_or_build_tlm_solver(self) -> HomogeneousBCLinearProblem: - """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 - {py:meth}`*ProblemBlock.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 {py:func}`~dolfinx_adjoint.compat.compute_form_adjoint`/ - {py:func}`~dolfinx_adjoint.ufl_utils.compute_adjoint`), dF/du - is used here as-is, so for a blocked problem it must be decomposed - with {py:func}`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. - """ + """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 *ProblemBlock.prepare_evaluate_tlm, blocks/solvers.py) + # rather than relying on solver-owned storage identity, since that + # storage is now shared across every block instead of private to one. if self._tlm_solver is None: dFdu_template = self._get_or_build_dFdu_template() # type: ignore[attr-defined] if isinstance(self._u, list): + # Unlike the adjoint operator (which decomposes dF/du back into + # blocks itself, inside compat.compute_form_adjoint/ + # ufl_utils.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 + # not, on its own, a compilable one -- the parts must be split + # apart first. dFdu_template = ufl.extract_blocks(dFdu_template) self._tlm_solver = HomogeneousBCLinearProblem( dFdu_template, @@ -661,12 +611,10 @@ def _make_block(self) -> _ProblemBlockBase: """Construct the tape block this Problem records for its forward solve. Each subclass overrides this to instantiate its own Block kind - ({py:class}`~dolfinx_adjoint.blocks.solvers.LinearProblemBlock`/ - {py:class}`~dolfinx_adjoint.blocks.solvers.NonlinearProblemBlock`, - constructor kwargs differing the same way the two Problem kinds' own - constructors do), passing ``self`` so the Block can reach back into - this Problem's shared solvers (see - {py:meth}`*ProblemBlock.get_reference_problem`). + (LinearProblemBlock/NonlinearProblemBlock, blocks/solvers.py; constructor + kwargs differ the same way the two Problem kinds' own constructors do), + passing ``self`` so the Block can reach back into this Problem's shared + solvers (see *ProblemBlock.get_reference_problem). Returns: A newly constructed, not-yet-recorded Block for this solve. @@ -677,34 +625,23 @@ def _dolfinx_solve(self) -> _Function | typing.Sequence[_Function]: """Perform the actual forward solve, via the base ``dolfinx.fem.petsc`` class. Each subclass overrides this to call its own base class's ``solve()`` - directly ({py:meth}`dolfinx.fem.petsc.LinearProblem.solve` / - {py:meth}`NonlinearProblem.solve`) - rather than ``self.solve()``, which would recurse back into this - Problem's own overridden, tape-recording ``solve()``. + directly (``dolfinx.fem.petsc.LinearProblem.solve``/ + ``NonlinearProblem.solve``) rather than ``self.solve()``, which would + recurse back into this Problem's own overridden, tape-recording + ``solve()``. Returns: - The solution {py:class}`~dolfinx_adjoint.Function`, or one per - output block for a blocked problem. + The solution Function, or one per output block for a blocked problem. """ def _record_and_solve(self, annotate: bool) -> _Function | typing.Sequence[_Function]: """Shared ``solve()`` skeleton for both classes. - Records a tape block (via the subclass's - {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._make_block` hook) when - annotating, refreshes the forward solver's placeholder coefficients - from the user's own current values (a prior recompute -- see - {py:meth}`*ProblemBlock.prepare_recompute_component` - -- may have left them holding a checkpointed/candidate value instead), - solves (via the subclass's - {py:meth}`~dolfinx_adjoint.solvers._ProblemBase._dolfinx_solve` hook), - and records the block's outputs. - Args: annotate: Whether to record this solve as a block on the working tape. Returns: - The solution {py:class}`~dolfinx_adjoint.Function`, or one per output block for a blocked problem. + The solution Function, or one per output block for a blocked problem. """ annotate = pyadjoint.annotate_tape({"annotate": annotate}) block = self._make_block() if annotate else None @@ -713,6 +650,10 @@ def _record_and_solve(self, annotate: bool) -> _Function | typing.Sequence[_Func tape = pyadjoint.get_working_tape() tape.add_block(block) + # Refresh the forward solver's placeholders from the user's own current + # values: a prior recompute (see *ProblemBlock.prepare_recompute_component, + # blocks/solvers.py) 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() @@ -811,21 +752,19 @@ def __init__( 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 a form is blocked from the user-side, it can be made without + # a {py:class}`ufl.MixedFunctionSpace`. Therefore we modify the + # form to use a {py:class}`ufl.MixedFunctionSpace` and split the form into its + # components. 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] + self._u = find_or_create_then_overload(u, L) # type: ignore[arg-type] - # A caller-omitted prefix must still be unique per Problem (see - # _PROBLEM_PREFIX_COUNTER) so PETSc's process-global options database - # never lets two Problems' SNES/KSP options bleed into each other. + # Unique, synchronized prefix for every solver instance (as SNES requires sync in prefix + # across processes). if petsc_options_prefix is None: petsc_options_prefix = f"dxa_linear_problem_{next(_PROBLEM_PREFIX_COUNTER)}_" @@ -839,34 +778,29 @@ def __init__( 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`. + # We replace all input coefficients with local coefficients, + # so that we don't disturb the input data during re-computations + # adjoints, etc. The local coefficients are stored in self._value_placeholders. + # The unknown `u` is not replaced, as it shouldn't be part of a linear problem's coefficients. u_list = self._u if isinstance(self._u, list) else [self._u] coefficients = collect_coefficients(a) | collect_coefficients(L) + if set(u_list).issubset(coefficients): + raise ValueError("The unknown `u` should not be part of the coefficients of a linear problem.") if P is not None: coefficients |= collect_coefficients(P) - coefficients -= set(u_list) + if set(u_list).issubset(coefficients): + raise ValueError("The unknown `u` should not be part of the coefficients of a linear problem.") + self._value_placeholders: dict[dolfinx.fem.Function, dolfinx.fem.Function] = { c: dolfinx.fem.Function(c.function_space) for c in coefficients } - - # Initialize linear solver + a_R, L_R, P_R = recursive_replace((a, L, P), self._value_placeholders) # type: ignore[misc] super().__init__( - a=recursive_replace(a, self._value_placeholders), # type: ignore[arg-type] - L=recursive_replace(L, self._value_placeholders), # type: ignore[arg-type] + a=a_R, # type: ignore[arg-type] + L=L_R, # type: ignore[arg-type] bcs=bcs, u=self._u, # type: ignore[arg-type] - P=recursive_replace(P, self._value_placeholders), # type: ignore[arg-type] + P=P_R, # type: ignore[arg-type] kind=kind, # type: ignore[arg-type] petsc_options_prefix=petsc_options_prefix, petsc_options=petsc_options, @@ -1032,19 +966,18 @@ def __init__( self._adj_options = adjoint_petsc_options self._tlm_options = tlm_petsc_options - # Assign mixed-space `part` indices to the test functions in a blocked - # residual once, here, mirroring LinearProblem: needed so a blocked - # residual's per-block forms can be safely combined into one - # whole-system form (via sum_form) inside _get_or_build_residual_template. + # If a form is blocked from the user-side, it can be made without + # a {py:class}`ufl.MixedFunctionSpace`. Therefore we modify the + # form to use a {py:class}`ufl.MixedFunctionSpace` and split the form into its + # components. if not isinstance(F, ufl.Form): F = assign_mixed_parts(F) # type: ignore[arg-type] - self._u = resolve_u(u, F) # type: ignore[arg-type] + self._u = find_or_create_then_overload(u, F) # type: ignore[arg-type] self._bcs = [] if bcs is None else bcs - # A caller-omitted prefix must still be unique per Problem (see - # _PROBLEM_PREFIX_COUNTER) so PETSc's process-global options database - # never lets two Problems' SNES/KSP options bleed into each other. + # Unique, synchronized prefix for every solver instance (as SNES requires sync in prefix + # across processes). if petsc_options_prefix is None: petsc_options_prefix = f"dxa_nonlinear_problem_{next(_PROBLEM_PREFIX_COUNTER)}_" @@ -1052,9 +985,7 @@ def __init__( # might appear in a hand-supplied Jacobian but not in F itself (e.g. a # stabilization term); the Jacobian _get_or_build_dFdu_template uses # for adjoint/TLM/Hessian purposes is always derived symbolically from - # F, never from this. Named distinctly from dolfinx.fem.petsc.NonlinearProblem's - # own `_J` (its compiled Jacobian, set by super().__init__() below) to avoid - # colliding with it. + # F, never from this. self._user_J = J self._rhs = F self._petsc_options = petsc_options @@ -1064,27 +995,11 @@ def __init__( self._petsc_options_prefix = petsc_options_prefix self._kind = kind - # 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. + # We replace all input coefficients with local coefficients, + # so that we don't disturb the input data during re-computations + # adjoints, etc. The local coefficients are stored in self._value_placeholders. + # The unknown is not replaced in the residual, but when deriving the + # adjoint and TLM solutions, through `_residual_state_placeholder` u_list = self._u if isinstance(self._u, list) else [self._u] coefficients = collect_coefficients(F) - set(u_list) if J is not None: @@ -1094,10 +1009,11 @@ def __init__( } # Initialize nonlinear solver + F_R, J_R, P_R = recursive_replace((F, J, P), self._value_placeholders) # type: ignore[misc] super().__init__( - F=recursive_replace(F, self._value_placeholders), # type: ignore[arg-type] - J=recursive_replace(J, self._value_placeholders), # type: ignore[arg-type] - P=recursive_replace(P, self._value_placeholders), # type: ignore[arg-type] + F=F_R, # type: ignore[arg-type] + J=J_R, # type: ignore[arg-type] + P=P_R, # type: ignore[arg-type] bcs=self._bcs, u=self._u, # type: ignore[arg-type] kind=kind, # type: ignore[arg-type] From a0dd7553ea79ec4b3c088b53c82de2bb7752e574 Mon Sep 17 00:00:00 2001 From: Joergen Schartum Dokken Date: Tue, 1 Sep 2026 11:52:34 +0000 Subject: [PATCH 5/8] Fix docstrings --- src/dolfinx_adjoint/ufl_utils.py | 24 ++---------------------- 1 file changed, 2 insertions(+), 22 deletions(-) diff --git a/src/dolfinx_adjoint/ufl_utils.py b/src/dolfinx_adjoint/ufl_utils.py index 59a28d9..bcd6551 100644 --- a/src/dolfinx_adjoint/ufl_utils.py +++ b/src/dolfinx_adjoint/ufl_utils.py @@ -184,14 +184,9 @@ def sum_form(form: NestedSequence[ufl.Form | None]) -> ufl.Form | None: def compute_adjoint(form: ufl.Form) -> typing.Sequence[typing.Sequence[ufl.Form]] | ufl.Form: """Compute the adjoint of a (possibly blocked) bilinear form. - A module-level function, not a method: it needs no ``Block``/``Problem`` state, - just ``form`` itself, so both ``_ProblemBlockBase`` (``blocks/solvers.py``) and - ``_ProblemBase`` (``solvers.py``) can call it directly rather than one reaching - into a "private" method defined on the other. - Args: - form: A bilinear form :math:`a(u, v)`, either a single ``ufl.Form`` or a - blocked (nested list) system. + form: A bilinear form :math:`a(u, v)`. Blocked forms should be summed with + {py:func}`sum_form` before passing to this function. Returns: The transposed form :math:`a(v, u)`, decomposed back into blocks (via @@ -203,21 +198,6 @@ def compute_adjoint(form: ufl.Form) -> typing.Sequence[typing.Sequence[ufl.Form] def recursive_replace(form: ufl.Form | typing.Sequence | None, placeholders: dict) -> ufl.Form | typing.Sequence | None: """Recursively apply {py:func}`ufl.replace` 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 {py:class}`~pyadjoint.Block`) exists to avoid: a - self-referential ``_replace`` closure inside a - {py:class}`~dolfinx_adjoint.LinearProblem`/{py:class}`~dolfinx_adjoint.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. - Args: form: A single form, ``None``, or an arbitrarily nested sequence of forms/``None`` (e.g. a blocked system). From 394a54d1baf2706338ec9622920f193abd8577f2 Mon Sep 17 00:00:00 2001 From: Joergen Schartum Dokken Date: Tue, 1 Sep 2026 11:55:19 +0000 Subject: [PATCH 6/8] Replace typing.Optional --- src/dolfinx_adjoint/blocks/assembly.py | 18 +++++++++--------- .../blocks/function_assigner.py | 2 +- src/dolfinx_adjoint/types/function.py | 10 ++++------ 3 files changed, 14 insertions(+), 16 deletions(-) diff --git a/src/dolfinx_adjoint/blocks/assembly.py b/src/dolfinx_adjoint/blocks/assembly.py index 57ccb71..f982c1e 100644 --- a/src/dolfinx_adjoint/blocks/assembly.py +++ b/src/dolfinx_adjoint/blocks/assembly.py @@ -11,7 +11,7 @@ def assemble_compiled_form( - form: dolfinx.fem.Form, tensor: typing.Optional[typing.Union[dolfinx.la.Vector, _SpecialVector | float]] = None + form: dolfinx.fem.Form, tensor: typing.Union[dolfinx.la.Vector, _SpecialVector | float] | None = None ) -> typing.Union[dolfinx.la.Vector, _SpecialVector, float]: """Assemble a compiled form and optionally apply Dirichlet boundary condition. @@ -54,10 +54,10 @@ class AssembleBlock(Block): def __init__( self, form: ufl.Form, - ad_block_tag: typing.Optional[str] = None, - jit_options: typing.Optional[dict] = None, - form_compiler_options: typing.Optional[dict] = None, - entity_maps: typing.Optional[typing.Sequence[dolfinx.mesh.EntityMap]] = None, + ad_block_tag: str | None = None, + jit_options: dict | None = None, + form_compiler_options: dict | None = None, + entity_maps: typing.Sequence[dolfinx.mesh.EntityMap] | None = None, ): super(AssembleBlock, self).__init__(ad_block_tag=ad_block_tag) @@ -87,10 +87,10 @@ def compute_action_adjoint( self, adj_input: typing.Union[float, dolfinx.la.Vector], arity_form: int, - form: typing.Optional[ufl.Form] = None, - c_rep: typing.Optional[typing.Union[ufl.Coefficient, ufl.Constant]] = None, - space: typing.Optional[dolfinx.fem.FunctionSpace] = None, - dform: typing.Optional[dolfinx.fem.Form] = None, + form: ufl.Form | None = None, + c_rep: typing.Union[ufl.Coefficient, ufl.Constant] | None = None, + space: dolfinx.fem.FunctionSpace | None = None, + dform: dolfinx.fem.Form | None = None, ): """This computes the action of the adjoint of the derivative of `form` wrt `c_rep` on `adj_input`. diff --git a/src/dolfinx_adjoint/blocks/function_assigner.py b/src/dolfinx_adjoint/blocks/function_assigner.py index 95583f5..f0ea816 100644 --- a/src/dolfinx_adjoint/blocks/function_assigner.py +++ b/src/dolfinx_adjoint/blocks/function_assigner.py @@ -49,7 +49,7 @@ def __init__( self, other: np.inexact | int | float | _Function | ufl.core.expr.Expr, func: _Function, - ad_block_tag: typing.Optional[str] = None, + ad_block_tag: str | None = None, ): super().__init__(ad_block_tag=ad_block_tag) diff --git a/src/dolfinx_adjoint/types/function.py b/src/dolfinx_adjoint/types/function.py index 3f92f4a..eae65a7 100644 --- a/src/dolfinx_adjoint/types/function.py +++ b/src/dolfinx_adjoint/types/function.py @@ -51,8 +51,8 @@ class Function(dolfinx.fem.Function, FloatingType): def __init__( self, V: dolfinx.fem.FunctionSpace, - x: typing.Optional[dolfinx.la.Vector] = None, - name: typing.Optional[str] = None, + x: dolfinx.la.Vector | None = None, + name: str | None = None, dtype: npt.DTypeLike = dolfinx.default_scalar_type, **kwargs, ): @@ -105,7 +105,7 @@ def _ad_create_checkpoint(self): def _ad_restore_at_checkpoint(self, checkpoint): return checkpoint - def _ad_dot(self, other: typing.Self, options: typing.Optional[dict] = None): + def _ad_dot(self, other: typing.Self, options: dict | None = None): """Compute the inner product of the current function with ``other`` in the Riesz representation. Args: @@ -165,9 +165,7 @@ def _ad_add(self, other: typing.Self) -> typing.Self: return r @no_annotations - def _ad_convert_riesz( - self, value: dolfinx.la.Vector, riesz_map: typing.Optional[dict] = None - ) -> dolfinx.fem.Function: + def _ad_convert_riesz(self, value: dolfinx.la.Vector, riesz_map: dict | None = None) -> dolfinx.fem.Function: """Convert a vector to a Riesz representation of the function.""" options = {} if riesz_map is None else riesz_map riesz_representation = options.get("riesz_representation", "l2") From b28e462420732f729b25d63270e5014a167cc897 Mon Sep 17 00:00:00 2001 From: Joergen Schartum Dokken Date: Tue, 1 Sep 2026 12:03:44 +0000 Subject: [PATCH 7/8] Add more coverage to nonlinear problem and simplify documentation --- tests/test_blocked_problem.py | 112 +++++++++++++++++++++++++++++++++- tests/test_solver_reuse.py | 110 +++++++++++++++++++++++++++++++++ tests/test_tlm_update.py | 33 +++++----- 3 files changed, 238 insertions(+), 17 deletions(-) diff --git a/tests/test_blocked_problem.py b/tests/test_blocked_problem.py index 940cc2f..0e97aa9 100644 --- a/tests/test_blocked_problem.py +++ b/tests/test_blocked_problem.py @@ -8,7 +8,14 @@ import ufl from dolfinx_adjoint import Function, assemble_scalar -from dolfinx_adjoint.solvers import LinearProblem +from dolfinx_adjoint.solvers import LinearProblem, NonlinearProblem + +direct_solve = { + "ksp_type": "preonly", + "pc_type": "lu", + "ksp_error_if_not_converged": True, + "pc_factor_mat_solver_type": "mumps", +} @pytest.fixture(scope="module") @@ -120,3 +127,106 @@ def L1(mesh, q): dHddu = hessian._ad_dot(f) min_rate = pyadjoint.taylor_test(Jh, z, f, dJdm=dJdm, Hm=dHddu) assert np.isclose(min_rate, 3.0, rtol=0.1, atol=0.1), f"Expected convergence rate close to 3.0, got {min_rate}" + + +@pytest.mark.parametrize("use_mixed_space", [True, False]) +def test_nonlinear_solver(use_mixed_space: bool, mesh_2D): + """As ``test_solver``, but for a blocked ``NonlinearProblem``: a Navier-Stokes-like + velocity/pressure system with a viscosity control, genuinely nonlinear in the state + via the convective term, exercising the same forward/adjoint/TLM/Hessian paths as + ``test_solver`` for the nonlinear (rather than linear) blocked residual. + """ + pyadjoint.get_working_tape().clear_tape() + mesh = mesh_2D + el_u = basix.ufl.element("P", mesh.basix_cell(), 2, shape=(mesh.geometry.dim,)) + el_p = basix.ufl.element("P", mesh.basix_cell(), 1) + V = dolfinx.fem.functionspace(mesh, el_u) + Q = dolfinx.fem.functionspace(mesh, el_p) + Z = dolfinx.fem.functionspace(mesh, ("DG", 0)) + dx = ufl.Measure("dx", domain=mesh) + + mu = Function(Z, name="viscosity") + mu.interpolate(lambda x: 1.0 + 0.5 * np.sin(np.pi * x[0])) + + uh, ph = Function(V, name="velocity"), Function(Q, name="pressure") + + # A moderate, non-conservative body force: large enough that the Taylor remainders + # stay clear of round-off, small enough that the Newton solve below converges + # reliably (the convective term is quadratic in the state). + x = ufl.SpatialCoordinate(mesh) + f = 10.0 * ufl.as_vector((ufl.sin(ufl.pi * x[1]), ufl.cos(ufl.pi * x[0]))) + + if use_mixed_space: + W = ufl.MixedFunctionSpace(*[V, Q]) + v, q = ufl.TestFunctions(W) + F = ufl.extract_blocks( + ufl.inner(mu * ufl.grad(uh), ufl.grad(v)) * dx + + ufl.inner(ufl.dot(ufl.grad(uh), uh), v) * dx + + ufl.inner(ph, ufl.div(v)) * dx + - ufl.inner(f, v) * dx + + ufl.inner(q, ufl.div(uh)) * dx + ) + else: + v, q = ufl.TestFunction(V), ufl.TestFunction(Q) + F = [ + ufl.inner(mu * ufl.grad(uh), ufl.grad(v)) * dx + + ufl.inner(ufl.dot(ufl.grad(uh), uh), v) * dx + + ufl.inner(ph, ufl.div(v)) * dx + - ufl.inner(f, v) * dx, + ufl.inner(q, ufl.div(uh)) * 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.zeros((mesh.geometry.dim,), dtype=dolfinx.default_scalar_type)) + bc = dolfinx.fem.dirichletbc(bc_val, boundary_dofs, V) + + forward_options = { + "snes_type": "newtonls", + "snes_error_if_not_converged": True, + "snes_atol": 1e-9, + "snes_rtol": 1e-9, + "snes_stol": 1e-12, + "ksp_type": "preonly", + "pc_type": "lu", + "pc_factor_mat_solver_type": "mumps", + } + problem = NonlinearProblem( + F, + u=[uh, ph], + bcs=[bc], + # Distinct from the default prefix: this fixture's tight snes_atol/rtol/stol + # must not leak into -- or collide with -- another NonlinearProblem elsewhere + # in the suite that happens to use the default prefix. + petsc_options_prefix=f"dxa_blocked_nonlinear_test_{use_mixed_space}_", + petsc_options=forward_options, + adjoint_petsc_options=direct_solve, + tlm_petsc_options=direct_solve, + ) + problem.solve() + + # Quartic in the state and with no constant offset, for the same round-off-avoidance + # reason as test_solver's objective. + J = assemble_scalar(ufl.inner(uh, uh) ** 2 * dx) + + control = pyadjoint.Control(mu) + Jh = pyadjoint.ReducedFunctional(J, control) + d = Function(Z) + d.interpolate(lambda x: 1.0 + 0.3 * np.cos(np.pi * x[1])) + e = Function(Z) + e.interpolate(lambda x: 0.2 * np.sin(3 * x[0])) + + min_rate = pyadjoint.taylor_test(Jh, d, e, dJdm=0) + assert np.isclose(min_rate, 1.0, rtol=1e-1, atol=1e-1), f"Expected convergence rate close to 1.0, got {min_rate}" + + Jh.derivative() + min_rate = pyadjoint.taylor_test(Jh, d, e) + assert np.isclose(min_rate, 2.0, rtol=1e-1, atol=1e-1), f"Expected convergence rate close to 2.0, got {min_rate}" + + Jh(d) + dJdm = Jh.derivative()._ad_dot(e) + hessian = Jh.hessian(e) + dHddu = hessian._ad_dot(e) + min_rate = pyadjoint.taylor_test(Jh, d, e, dJdm=dJdm, Hm=dHddu) + assert np.isclose(min_rate, 3.0, rtol=0.1, atol=0.1), f"Expected convergence rate close to 3.0, got {min_rate}" diff --git a/tests/test_solver_reuse.py b/tests/test_solver_reuse.py index b81738f..558e1d1 100644 --- a/tests/test_solver_reuse.py +++ b/tests/test_solver_reuse.py @@ -13,6 +13,7 @@ from mpi4py import MPI +import basix.ufl import dolfinx import numpy as np import pyadjoint @@ -804,3 +805,112 @@ def test_default_petsc_options_prefix_is_unique_per_problem(): # An explicitly-passed prefix must be honored unchanged. explicit = LinearProblem(a, L, petsc_options_prefix="my_custom_prefix_") assert explicit._petsc_options_prefix == "my_custom_prefix_" + + +def test_nonlinear_blocked_problem_templates_compiled_once(): + """As ``test_nonlinear_adjoint_lhs_compiled_once``/``test_nonlinear_tlm_rhs_templates_compiled_once``, + but for a *blocked* ``NonlinearProblem`` (a Navier-Stokes-like velocity/pressure system): + the adjoint solver, TLM solver, per-dependency TLM right-hand-side templates, and the + blocked Hessian templates (``NonlinearProblem._get_or_build_hessian_templates``'s + ``isinstance(self._u, list)`` branch, which otherwise has no caching-focused regression + coverage at all) must all be compiled exactly once and never rebuilt across repeated + ``derivative()``/``hessian()`` calls at different control values. + """ + pyadjoint.get_working_tape().clear_tape() + + mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 6, 6) + el_u = basix.ufl.element("P", mesh.basix_cell(), 2, shape=(mesh.geometry.dim,)) + el_p = basix.ufl.element("P", mesh.basix_cell(), 1) + V = dolfinx.fem.functionspace(mesh, el_u) + Q = dolfinx.fem.functionspace(mesh, el_p) + Z = dolfinx.fem.functionspace(mesh, ("DG", 0)) + dx = ufl.Measure("dx", domain=mesh) + + mu = Function(Z, name="viscosity") + mu.interpolate(lambda x: 1.0 + 0.5 * np.sin(np.pi * x[0])) + + uh, ph = Function(V, name="velocity"), Function(Q, name="pressure") + v, q = ufl.TestFunction(V), ufl.TestFunction(Q) + + x = ufl.SpatialCoordinate(mesh) + f = 10.0 * ufl.as_vector((ufl.sin(ufl.pi * x[1]), ufl.cos(ufl.pi * x[0]))) + F0 = ( + ufl.inner(mu * ufl.grad(uh), ufl.grad(v)) * dx + + ufl.inner(ufl.dot(ufl.grad(uh), uh), v) * dx + + ufl.inner(ph, ufl.div(v)) * dx + - ufl.inner(f, v) * dx + ) + F1 = ufl.inner(q, ufl.div(uh)) * 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) + zero = dolfinx.fem.Constant(mesh, np.zeros(mesh.geometry.dim, dtype=dolfinx.default_scalar_type)) + bc = dolfinx.fem.dirichletbc(zero, dofs, V) + + forward_options = { + "snes_type": "newtonls", + "snes_error_if_not_converged": True, + "snes_atol": 1e-9, + "snes_rtol": 1e-9, + "snes_stol": 1e-12, + "ksp_type": "preonly", + "pc_type": "lu", + "pc_factor_mat_solver_type": "mumps", + } + adjoint_options = { + "ksp_type": "preonly", + "pc_type": "lu", + "ksp_error_if_not_converged": True, + "pc_factor_mat_solver_type": "mumps", + } + problem = NonlinearProblem( + [F0, F1], + u=[uh, ph], + bcs=[bc], + petsc_options_prefix="dxa_blocked_nonlinear_reuse_test_", + petsc_options=forward_options, + adjoint_petsc_options=adjoint_options, + tlm_petsc_options=adjoint_options, + ) + problem.solve() + + J = assemble_scalar(ufl.inner(uh, uh) ** 2 * dx) + control = pyadjoint.Control(mu) + Jh = pyadjoint.ReducedFunctional(J, control) + + dm = Function(Z) + dm.interpolate(lambda x: 0.2 * np.sin(3 * x[0])) + Jh.hessian(dm) + + adjoint_solver = problem._get_or_build_adjoint_solver() + compiled_adjoint_lhs = adjoint_solver._a + assert compiled_adjoint_lhs is not None + + tlm_solver = problem._get_or_build_tlm_solver() + compiled_tlm_lhs = tlm_solver._a + assert compiled_tlm_lhs is not None + + tlm_templates, _, _ = problem._get_or_build_tlm_rhs_templates() + compiled_tlm_ids = {c: id(form) for c, form in tlm_templates.items()} + assert compiled_tlm_ids, "no TLM RHS templates were built" + + hessian_templates = problem._get_or_build_hessian_templates() + assert isinstance(hessian_templates.soa_self, list), "expected a per-row list for a blocked problem" + + mu2 = Function(Z) + mu2.interpolate(lambda x: 2.0 + np.cos(np.pi * x[0])) + Jh(mu2) + Jh.derivative() + Jh.hessian(dm) + + assert adjoint_solver._a is compiled_adjoint_lhs, "adjoint LHS was rebuilt after evaluating at a new point" + assert tlm_solver._a is compiled_tlm_lhs, "TLM LHS was rebuilt after evaluating at a new point" + + tlm_templates_after, _, _ = problem._get_or_build_tlm_rhs_templates() + for c, form in tlm_templates_after.items(): + assert id(form) == compiled_tlm_ids[c], f"TLM RHS template for {c.name} was rebuilt" + + assert problem._get_or_build_hessian_templates() is hessian_templates, ( + "blocked Hessian templates were rebuilt after evaluating at a new point" + ) diff --git a/tests/test_tlm_update.py b/tests/test_tlm_update.py index e9b6ed3..f4b1646 100644 --- a/tests/test_tlm_update.py +++ b/tests/test_tlm_update.py @@ -1,11 +1,15 @@ -"""Regression tests for the tangent-linear solver cached on ``LinearProblemBlock``. - -``construct_tlm_solver`` compiles ``_compute_residual_derivative()`` once and caches the -resulting solver on the block. That form is built from ``block_variable.saved_output``, -which is a fresh object after every control update, so the cached operator can stay -pinned to the evaluation point it was first built at. When the control appears in the -bilinear form -- so that dF/du genuinely depends on it -- every Hessian computed after -the first one is then silently wrong. +"""Regression tests for cached adjoint/TLM/Hessian templates staying correct across evaluations. + +``_ProblemBase`` (``solvers.py``) builds ``dF/du``, the TLM right-hand side, and the Hessian +templates once per ``LinearProblem``/``NonlinearProblem`` and reuses them -- refreshed at each +new evaluation point -- across every block that Problem records, rather than recompiling per +block/solve. These tests check that reuse never leaves an operator evaluated at a stale point: +when the control appears in a way that makes ``dF/du`` genuinely depend on it, a cached-but- +unrefreshed operator would make every Hessian computed after the first one silently wrong. +Covers scalar and blocked (multi-field) problems, for both ``LinearProblem`` and +``NonlinearProblem``, plus a warm-start pollution check (``test_hessian_mpi_breakdown`` -- +not actually MPI-specific, despite the name: it checks that solving at one control value in +between two evaluations at another value doesn't perturb the second evaluation's result). """ from mpi4py import MPI @@ -36,9 +40,8 @@ def mesh_2D(): def _viscous_stokes(mesh): """Blocked Stokes-like problem whose control ``mu`` sits inside ``a[0][0]``. - The control has to enter the bilinear form for this test to have any teeth: if it - only enters ``L``, dF/du is independent of the control and a stale tangent-linear - operator is indistinguishable from a fresh one. + The control is part of the bilinear form, so that dF/du!=0, and we get + a genuinely control-dependent tangent-linear operator. """ el_u = basix.ufl.element("P", mesh.basix_cell(), 2, shape=(mesh.geometry.dim,)) el_p = basix.ufl.element("P", mesh.basix_cell(), 1) @@ -52,11 +55,9 @@ def _viscous_stokes(mesh): u, p = ufl.TrialFunction(V), ufl.TrialFunction(Q) v, q = ufl.TestFunction(V), ufl.TestFunction(Q) - # A rotational (non-conservative) body force. A constant force would be balanced - # exactly by the pressure in a closed incompressible box, leaving u == 0 and making - # the functional independent of the control. The 1e3 only sets the scale of the - # state, which keeps the Taylor remainders clear of pyadjoint's absolute - # machine-precision warning threshold. + # A rotational (non-conservative) body force. Avoid using constant force to avoid + # velocity being 0 always, independent of control. Scaled to ensure decent-sized Taylor + # remainders. x = ufl.SpatialCoordinate(mesh) f = 1e3 * ufl.as_vector((ufl.sin(ufl.pi * x[1]), ufl.cos(ufl.pi * x[0]))) From 7167121f5259afdd9c0ebb0ee7a2c37c2ee70310 Mon Sep 17 00:00:00 2001 From: Joergen Schartum Dokken Date: Tue, 1 Sep 2026 12:25:00 +0000 Subject: [PATCH 8/8] Minor improvements on documentation and addition of 3rd order taylor test --- README.md | 4 ++-- demos/poisson_mother.py | 4 ++-- tests/test_blocked_problem.py | 14 ++++++++++++++ 3 files changed, 18 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index b5e08ce..0e51f47 100644 --- a/README.md +++ b/README.md @@ -41,13 +41,13 @@ conda install -c conda-forge dolfinx-adjoint To install the latest development version directly from the repository, use: ```bash -python3 -m pip install git+[https://github.com/scientificcomputing/dolfinx-adjoint.git](https://github.com/scientificcomputing/dolfinx-adjoint.git) +python3 -m pip install git+https://github.com/scientificcomputing/dolfinx-adjoint.git ``` If you plan to actively modify the code, clone the repository and install the optional dependencies for testing, development, and documentation generation: ```bash -git clone [https://github.com/scientificcomputing/dolfinx-adjoint.git](https://github.com/scientificcomputing/dolfinx-adjoint.git) +git clone https://github.com/scientificcomputing/dolfinx-adjoint.git cd dolfinx-adjoint python3 -m pip install -e ".[all]" ``` diff --git a/demos/poisson_mother.py b/demos/poisson_mother.py index 1e87c50..60da429 100644 --- a/demos/poisson_mother.py +++ b/demos/poisson_mother.py @@ -156,7 +156,7 @@ def refinement_region(x, tol=1e-14): # the `dolfinx.fem.petsc.LinearProblem` class. # # ```{note} -# When creating the :py:func:`dolfinx_adjoint.LinearProblem`, we can specify the solver options that +# When creating the :py:class:`dolfinx_adjoint.LinearProblem`, we can specify the solver options that # are passed on to the underlying PETSc Krylov subspace solver. # This is also the place to pass in solver options for the first and second order adjoint equations # and the tangent linear model (TLM) equation. @@ -182,7 +182,7 @@ def refinement_region(x, tol=1e-14): # ```{note} # Note that we can pass in solver options for the adjoint equation via the keyword argument -# `adjoint_petsc_options`. As we a solving a linear, symmetric problem, i,e, a self-adjoint problem, +# `adjoint_petsc_options`. As we are solving a linear, symmetric problem, i.e. a self-adjoint problem, # we use the same options for both the forward and adjoint problems. # ``` diff --git a/tests/test_blocked_problem.py b/tests/test_blocked_problem.py index 0e97aa9..b88740b 100644 --- a/tests/test_blocked_problem.py +++ b/tests/test_blocked_problem.py @@ -230,3 +230,17 @@ def test_nonlinear_solver(use_mixed_space: bool, mesh_2D): dHddu = hessian._ad_dot(e) min_rate = pyadjoint.taylor_test(Jh, d, e, dJdm=dJdm, Hm=dHddu) assert np.isclose(min_rate, 3.0, rtol=0.1, atol=0.1), f"Expected convergence rate close to 3.0, got {min_rate}" + + # A second, independent evaluation point/direction: a cached-but-unrefreshed + # adjoint/TLM/Hessian operator (see tests/test_tlm_update.py) could pass the + # check above yet still be silently wrong here. + mu2 = Function(Z) + mu2.interpolate(lambda x: 2.0 + np.sin(x[1])) + h2 = Function(Z) + h2.interpolate(lambda x: 0.5 * np.cos(4 * x[0])) + Jh(mu2) + dJdm = Jh.derivative()._ad_dot(h2) + hessian = Jh.hessian(h2) + dHddu = hessian._ad_dot(h2) + min_rate = pyadjoint.taylor_test(Jh, mu2, h2, dJdm=dJdm, Hm=dHddu) + assert np.isclose(min_rate, 3.0, rtol=0.1, atol=0.1), f"Expected convergence rate close to 3.0, got {min_rate}"