From d2ae8a230c6f506d8fa4e59c239bec6a4268408b Mon Sep 17 00:00:00 2001 From: Joergen Schartum Dokken Date: Tue, 25 Aug 2026 14:06:16 +0000 Subject: [PATCH 01/26] Various fixes to nonlinear solve. After long dicussion with Gemini (August 2026), I disregarded all advice it gave, and figured out that one of the two issues were: - Since we use SNES and setCTX in the setFunction and setJacobian operations, we can't use the classical python-esque replacement of compiled forms. This invalidated the recompute for non-linear problems. Another issue covered by Claude (in another PR spawned by @finsberg) is that `u` was added as its own dependency if the problem was truely nonlinear. This is now fixed. Added tests to check this. --- src/dolfinx_adjoint/blocks/solvers.py | 138 ++++++++++++-------------- tests/test_nonlinear_problem.py | 93 +++++++++++++++++ 2 files changed, 155 insertions(+), 76 deletions(-) create mode 100644 tests/test_nonlinear_problem.py diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index 1ff15be..2e75261 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -106,11 +106,15 @@ def __init__( # NOTE: Add mesh and constants as dependencies later on try: - for c in self._lhs.coefficients(): # type: ignore - self.add_dependency(c, no_duplicates=True) - for c in self._rhs.coefficients(): # type: ignore - self.add_dependency(c, no_duplicates=True) + u_list = self._u if isinstance(self._u, list) else [self._u] + if self._lhs is not None: + for c in self._lhs.coefficients(): # type: ignore + if c not in u_list: # Exclude the unknown + self.add_dependency(c, no_duplicates=True) + for c in self._rhs.coefficients(): # type: ignore + if c not in u_list: # Exclude the unknown + self.add_dependency(c, no_duplicates=True) except AttributeError: raise NotImplementedError("Blocked systems not implemented yet.") self._compiled_lhs = dolfinx.fem.form( @@ -502,30 +506,37 @@ def evaluate_adj_component( 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}`.""" + """Evaluate the adjoint component, i.e. :math:`\\frac{\\partial F}{\\partial m}`.""" 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)}") + # Compute the sensitivity of the residual with respect to the parameter dFdm = -ufl.derivative(residual, c_rep, dc) + if dFdm.empty(): + # Generate a dummy form to safely extract the correct Vector wrapper type + dFdm = dolfinx.fem.form(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): @@ -759,7 +770,10 @@ def __init__( self._u: dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function] if isinstance(u, dolfinx.fem.Function): self._u = pyadjoint.create_overloaded_object(u) - self._rhs = F + replace_dict = {u: self._u} + self._rhs = ufl.replace(F, replace_dict) + self._lhs = ufl.replace(J, replace_dict) if J is not None else None + self._preconditioner = ufl.replace(P, replace_dict) if P is not None else None else: self._u = [pyadjoint.create_overloaded_object(ui) for ui in u] assert isinstance(F, typing.Iterable) @@ -768,12 +782,18 @@ def __init__( # NOTE: Add mesh and constants as dependencies later on try: - for c in self._lhs.coefficients(): # type: ignore - self.add_dependency(c, no_duplicates=True) - for c in self._rhs.coefficients(): # type: ignore - self.add_dependency(c, no_duplicates=True) + u_list = self._u if isinstance(self._u, list) else [self._u] + if self._lhs is not None: + for c in self._lhs.coefficients(): + if c not in u_list: # Exclude unknown + self.add_dependency(c, no_duplicates=True) + if self._rhs is not None: + for c in self._rhs.coefficients(): + if c not in u_list: # Exclude unknown + self.add_dependency(c, no_duplicates=True) except AttributeError: raise NotImplementedError("Blocked systems not implemented yet.") + self._compiled_lhs = dolfinx.fem.form( self._lhs, # type: ignore jit_options=jit_options, @@ -868,71 +888,37 @@ def _replace_coefficients_in_form(self, form: ufl.Form) -> ufl.Form: def prepare_recompute_component(self, inputs, relevant_outputs): """Prepare for recomputing the block with different control inputs.""" - # Create initial guess for the KSP solver - # Form independnet compilation would make it possible to use the same KSP for all re-evaluations. - if isinstance(self._u, Function): - initial_guess = dolfinx.fem.Function(self._u.function_space, name=self._u.name + "_initial_guess") - else: - initial_guess = [dolfinx.fem.Function(u.function_space, name=u.name + "_initial_guess") for u in self._u] + # As opposed to the linear problem, we need to update the coefficients in place, + # as the nonlinear problem snes.setContext doesn't reflect in place updates on the solver. + for block_variable in self.get_dependencies(): + coeff = block_variable.output + if isinstance(coeff, dolfinx.fem.Function): + coeff.x.array[:] = block_variable.saved_output.x.array[:] + coeff.x.scatter_forward() - # Replace values in the DirichletBC if it is dependent on a control - # NOTE: Currently assume that BCS are control independent. - bcs = self._bcs - # for block_variable in self.get_dependencies(): - # c = block_variable.output - # c_rep = block_variable.saved_output + # Warm-start original unknown objects in place + u_list = self._forward_solver._u if isinstance(self._forward_solver._u, list) else [self._forward_solver._u] + for idx, out_bv in relevant_outputs: + u_list[idx].x.array[:] = out_bv.saved_output.x.array[:] + u_list[idx].x.scatter_forward() - # if isinstance(c, dolfinx.fem.DirichletBC): - # bcs.append(c_rep) - - # Replace form coefficients with checkpointed values. - # Loop through the dependencies of the lhs and rhs, check if they are in the respective form - lhs = self._replace_coefficients_in_form(self._lhs) - rhs = self._replace_coefficients_in_form(self._rhs) - preconditioner = ( - self._replace_coefficients_in_form(self._preconditioner) if self._preconditioner is not None else None - ) - compiled_lhs = dolfinx.fem.form( - lhs, - jit_options=self._jit_options, - form_compiler_options=self._form_compiler_options, - entity_maps=self._entity_maps, - ) - compiled_rhs = dolfinx.fem.form( - rhs, - jit_options=self._jit_options, - form_compiler_options=self._form_compiler_options, - entity_maps=self._entity_maps, - ) - compiled_preconditioner = ( - dolfinx.fem.form( - preconditioner, - jit_options=self._jit_options, - form_compiler_options=self._form_compiler_options, - entity_maps=self._entity_maps, - ) - if preconditioner is not None - else None - ) - - # Replace the compiled forms with those with new coefficients. - self._forward_solver._a = compiled_lhs - self._forward_solver._L = compiled_rhs - self._forward_solver._P = compiled_preconditioner - self._forward_solver.bcs = bcs - self._forward_solver._u = initial_guess + return None def recompute_component( self, inputs: typing.Iterable[Function], block_variable, idx: int, prepared: None ) -> typing.Union[dolfinx.fem.Function, typing.Iterable[dolfinx.fem.Function]]: """Recompute the block with the prepared linear problem.""" - solution = self._forward_solver.solve() - return solution + with pyadjoint.tape.stop_annotating(): + self._forward_solver.solve() + + if isinstance(self._forward_solver._u, list): + return self._forward_solver._u[idx] + else: + return self._forward_solver._u 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): @@ -980,21 +966,17 @@ def _compute_residual(self) -> typing.Union[ufl.Form, list[ufl.Form]]: """ # NOTE: Should probably be possible to compile this form once. replacement_functions = self.get_outputs() - F_form: typing.Union[ufl.Form, list[ufl.Form]] = [] - assert isinstance(self._rhs, ufl.Form) + assert isinstance(self._rhs, (ufl.Form, typing.Sequence)) 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, Function): + if isinstance(self._u, dolfinx.fem.Function): F_form = ufl.replace(self._rhs, replacement_map) else: - assert isinstance(F_form, list) - assert isinstance(self._rhs, typing.Sequence) - assert len(F_form) == len(self._rhs) - for j, rhs_j in enumerate(self._rhs): - F_form[j] = ufl.replace(rhs_j, replacement_map) + 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]]]: @@ -1044,14 +1026,13 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar bcs = [] for bc in self._bcs: bcs.append(bc) + dFdm = ufl.ZeroBaseForm((ufl.TestFunction(V),)) for block_variable in self.get_dependencies(): tlm_value = block_variable.tlm_value - # c = block_variable.output c_rep = block_variable.saved_output if tlm_value is None: continue - dFdm += ufl.derivative(-F, c_rep, tlm_value) if isinstance(dFdm, float): @@ -1134,6 +1115,11 @@ def evaluate_adj_component( 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( diff --git a/tests/test_nonlinear_problem.py b/tests/test_nonlinear_problem.py new file mode 100644 index 0000000..70ec9a3 --- /dev/null +++ b/tests/test_nonlinear_problem.py @@ -0,0 +1,93 @@ +from mpi4py import MPI +import dolfinx +import numpy as np +import pyadjoint +import ufl + +from dolfinx_adjoint import Function, assemble_scalar +from dolfinx_adjoint.solvers import NonlinearProblem + + +def test_sequential_nonlinear_problems(): + """ + Test two cascaded non-linear PDEs. + PDE 1: -div(u1 * grad(u1)) = f + PDE 2: -div(u2 * grad(u2)) = u1 + """ + pyadjoint.get_working_tape().clear_tape() + mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 8, 7) + + V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) + + f = Function(V, name="control") + # Keep the control positive to avoid degeneracy in the diffusion tensor + f.interpolate(lambda x: 2.0 + np.sin(x[0])) + + u1 = Function(V, name="state_1") + u1.interpolate(lambda x: np.ones_like(x[0])) # Non-zero initial guess + v1 = ufl.TestFunction(V) + + F1 = (1 + u1**2) * ufl.inner(ufl.grad(u1), ufl.grad(v1)) * ufl.dx(domain=mesh) - f * v1 * ufl.dx(domain=mesh) + + # Setup PDE 2 + u2 = Function(V, name="state_2") + u2.interpolate(lambda x: np.ones_like(x[0])) # Non-zero initial guess + v2 = ufl.TestFunction(V) + F2 = (2 + u2**2) * ufl.inner(ufl.grad(u2), ufl.grad(v2)) * ufl.dx(domain=mesh) - u1 * v2 * ufl.dx(domain=mesh) + + # 4. Boundary Conditions (u = 1.0 on boundary) + 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) + + # Use SNES options for the nonlinear solver + options = { + "snes_monitor": None, + "snes_error_if_not_converged": True, + "snes_type": "newtonls", + "ksp_type": "preonly", + "pc_type": "lu", + "pc_factor_mat_solver_type": "mumps", + } + + # 5. Solve the Cascade + problem1 = NonlinearProblem(F1, u=u1, bcs=[bc], petsc_options=options, adjoint_petsc_options=options) + problem1.solve() + + problem2 = NonlinearProblem(F2, u=u2, bcs=[bc], petsc_options=options, adjoint_petsc_options=options) + problem2.solve() + + # 6. Objective (using the cubed error to ensure a 3.0 Hessian rate) + d = pyadjoint.AdjFloat(0.2) + error = (u2 - d) ** 3 * ufl.dx(domain=mesh) + J = assemble_scalar(error) + + # 7. Taylor Tests + control = pyadjoint.Control(f) + Jh = pyadjoint.ReducedFunctional(J, control) + + ctrl_eval = Function(V) + ctrl_eval.interpolate(lambda x: 4.0 + np.sin(x[0])) + + pert = Function(V) + pert.interpolate(lambda x: 15.1 * np.cos(x[1])) + + Jh(ctrl_eval) + min_rate_grad = pyadjoint.taylor_test(Jh, ctrl_eval, pert, dJdm=0) + print("\n--- 1st-order Taylor test ---") + Jh(ctrl_eval) + min_rate_grad = pyadjoint.taylor_test(Jh, ctrl_eval, pert) + assert np.isclose(min_rate_grad, 2.0, rtol=1e-2, atol=1e-2), f"Expected 2.0, got {min_rate_grad}" + + print("\n--- 2nd-order Taylor test ---") + Jh(ctrl_eval) + dJdm = Jh.derivative()._ad_dot(pert) + dHddu = Jh.hessian(pert)._ad_dot(pert) + min_rate_hess = pyadjoint.taylor_test(Jh, ctrl_eval, pert, dJdm=dJdm, Hm=dHddu) + assert np.isclose(min_rate_hess, 3.0, rtol=1e-2, atol=1e-2), f"Expected 3.0, got {min_rate_hess}" + + +if __name__ == "__main__": + test_sequential_nonlinear_problems() From 90d099434e95fa88f4c322aef7e2e74ae5dc4dfd Mon Sep 17 00:00:00 2001 From: Joergen Schartum Dokken Date: Tue, 25 Aug 2026 14:47:49 +0000 Subject: [PATCH 02/26] Start on blocked linear problem --- src/dolfinx_adjoint/blocks/assembly.py | 1 - src/dolfinx_adjoint/blocks/solvers.py | 73 +++++++++++++++----- tests/test_blocked_problem.py | 95 ++++++++++++++++++++++++++ 3 files changed, 150 insertions(+), 19 deletions(-) create mode 100644 tests/test_blocked_problem.py diff --git a/src/dolfinx_adjoint/blocks/assembly.py b/src/dolfinx_adjoint/blocks/assembly.py index dc50824..57ccb71 100644 --- a/src/dolfinx_adjoint/blocks/assembly.py +++ b/src/dolfinx_adjoint/blocks/assembly.py @@ -188,7 +188,6 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies): c_rep = block_variable.saved_output if coeff in self.form.coefficients(): replaced_coeffs[coeff] = c_rep - form = ufl.replace(self.form, replaced_coeffs) return form diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index 2e75261..a74f78b 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -105,17 +105,25 @@ def __init__( self._u = [pyadjoint.create_overloaded_object(ui) for ui in u] # NOTE: Add mesh and constants as dependencies later on - try: - u_list = self._u if isinstance(self._u, list) else [self._u] - if self._lhs is not None: - for c in self._lhs.coefficients(): # type: ignore - if c not in u_list: # Exclude the unknown - self.add_dependency(c, no_duplicates=True) - + if isinstance(self._u, dolfinx.fem.Function): + for c in self._lhs.coefficients(): # type: ignore + if c != self._u: # Exclude the unknown + self.add_dependency(c, no_duplicates=True) for c in self._rhs.coefficients(): # type: ignore - if c not in u_list: # Exclude the unknown + if c != self._u: # Exclude the unknown self.add_dependency(c, no_duplicates=True) - except AttributeError: + elif isinstance(self._u, typing.Iterable): + for Ai in self._lhs: # type: ignore + for Aij in Ai: + if Aij is not None: + for c in Aij.coefficients(): + if c not in self._u: + self.add_dependency(c, no_duplicates=True) + for i, part in enumerate(self._rhs): # type: ignore + for c in part.coefficients(): + if c not in self._u: + self.add_dependency(c, no_duplicates=True) + else: raise NotImplementedError("Blocked systems not implemented yet.") self._compiled_lhs = dolfinx.fem.form( self._lhs, @@ -193,23 +201,42 @@ def _recover_bcs(self): bcs.append(c_rep) return bcs - def _create_replace_map(self, form: ufl.Form) -> dict[Function, Function]: + def _create_replace_map(self, form: ufl.Form | typing.Iterable[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 coeff in form.coefficients(): - replace_map[coeff] = block_variable.saved_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 _replace_coefficients_in_form(self, form: ufl.Form) -> ufl.Form: + def _replace_coefficients_in_form( + self, form: ufl.Form | typing.Iterable[ufl.Form] + ) -> ufl.Form | typing.Iterable[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) + if isinstance(form, ufl.Form): + return ufl.replace(form, replace_map) + elif isinstance(form, typing.Iterable): + replaced_forms = [] + for f in form: + if f is None: + replaced_forms.append(None) + elif isinstance(f, typing.Iterable): + replaced_forms.append(self._replace_coefficients_in_form(f)) + else: + replaced_forms.append(ufl.replace(f, replace_map)) + return replaced_forms def prepare_recompute_component(self, inputs, relevant_outputs): """Prepare for recomputing the block with different control inputs.""" @@ -257,13 +284,19 @@ def prepare_recompute_component(self, inputs, relevant_outputs): self._forward_solver._P = compiled_preconditioner self._forward_solver.bcs = self._bcs self._forward_solver._u = initial_guess + with pyadjoint.stop_annotating(): + solution = self._forward_solver.solve() + return solution def recompute_component( self, inputs: typing.Iterable[Function], block_variable, idx: int, prepared: None ) -> typing.Union[dolfinx.fem.Function, typing.Iterable[dolfinx.fem.Function]]: """Recompute the block with the prepared linear problem.""" - solution = self._forward_solver.solve() - return solution + if isinstance(prepared, dolfinx.fem.Function): + assert idx == 0 + return prepared + else: + return prepared[idx] def _should_compute_boundary_adjoint( self, relevant_dependencies: typing.List[tuple[int, pyadjoint.block_variable.BlockVariable]] @@ -303,8 +336,12 @@ def _compute_adjoint( 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)) + if form_ij is None: + tmp_form[i].append(None) + adj_form[i].append(None) + else: + 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 diff --git a/tests/test_blocked_problem.py b/tests/test_blocked_problem.py new file mode 100644 index 0000000..e3102f7 --- /dev/null +++ b/tests/test_blocked_problem.py @@ -0,0 +1,95 @@ +import typing + +from mpi4py import MPI +import basix.ufl +import dolfinx +import numpy as np +import pyadjoint +import pytest +import ufl + +from dolfinx_adjoint import Function, assemble_scalar +from dolfinx_adjoint.solvers import LinearProblem + + +@pytest.fixture(scope="module") +def mesh_2D(): + return dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 8, 7) + + +@pytest.fixture(scope="module") +def mesh_3D(): + return dolfinx.mesh.create_unit_cube(MPI.COMM_WORLD, 11, 13, 12, cell_type=dolfinx.mesh.CellType.hexahedron) + + +@pytest.mark.parametrize("constant", [np.float64(0.2), float(-0.13), int(3)]) +@pytest.mark.parametrize("mesh_var_name", ["mesh_2D"]) +def test_solver(mesh_var_name: str, request, constant: typing.Union[float, int, np.floating]): + pyadjoint.get_working_tape().clear_tape() + mesh = request.getfixturevalue(mesh_var_name) + 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) + W = ufl.MixedFunctionSpace(*[V, Q]) + u, p = ufl.TrialFunctions(W) + v, q = ufl.TestFunctions(W) + dx = ufl.Measure("dx", domain=mesh) + a = ufl.inner(ufl.grad(u), ufl.grad(v)) * dx + ufl.inner(p, ufl.div(v)) * dx + ufl.inner(q, ufl.div(u)) * dx + + f = Function(V, name="control") + f.interpolate(lambda x: (np.sin(x[0]), x[1])) + L = ufl.inner(f, v) * dx + L += dolfinx.fem.Constant(mesh, 0.0) * q * 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) + + options = { + "ksp_monitor": None, + "ksp_type": "preonly", + "pc_type": "lu", + "ksp_error_if_not_converged": True, + "pc_factor_mat_solver_type": "mumps", + } + uh, ph = (Function(V, name="state"), Function(Q, name="pressure")) + problem = LinearProblem( + ufl.extract_blocks(a), + ufl.extract_blocks(L), + u=[uh, ph], + bcs=[bc], + petsc_options=options, + adjoint_petsc_options=options, + tlm_petsc_options=options, + ) + problem.solve() + + d = pyadjoint.AdjFloat(constant) + x = ufl.SpatialCoordinate(mesh) + c = ufl.as_vector((d * ufl.sin(x[0]), d * ufl.cos(x[1]))) + error = ufl.inner(uh - c, uh - c) * ufl.inner(uh - c, uh - c) * ufl.dx + J = assemble_scalar(error) + + control = pyadjoint.Control(f) + Jh = pyadjoint.ReducedFunctional(J, control) + d = Function(V) + d.interpolate(lambda x: (10 * x[0], x[1])) + + e = Function(V) + e.interpolate(lambda x: (2000 * np.sin(x[0]), -1000 * x[1])) + min_rate = pyadjoint.taylor_test(Jh, d, e, dJdm=0) + assert np.isclose(min_rate, 1.0, rtol=1e-2, atol=1e-2), 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-2, atol=1e-2), 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=5e-3, atol=5e-3), f"Expected convergence rate close to 3.0, got {min_rate}" From 4f823d3c6ec0f59ba2bf95a83b9e9c20d9c2bea1 Mon Sep 17 00:00:00 2001 From: Joergen Schartum Dokken Date: Wed, 26 Aug 2026 08:10:46 +0000 Subject: [PATCH 03/26] Various fixes for blocked linear problem --- src/dolfinx_adjoint/blocks/solvers.py | 89 ++++++++++++++++++++++----- 1 file changed, 72 insertions(+), 17 deletions(-) diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index a74f78b..4e0c199 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -336,7 +336,7 @@ def _compute_adjoint( tmp_form.append([]) adj_form.append([]) for j, form_ij in enumerate(f_i): - if form_ij is None: + if form_ij is None or form_ij.empty(): tmp_form[i].append(None) adj_form[i].append(None) else: @@ -364,15 +364,15 @@ def _compute_residual(self) -> typing.Union[ufl.Form, list[ufl.Form]]: assert len(self._u) == len(replacement_functions), ( f"Expected {len(self._u)} output functions, got {len(replacement_functions)}" ) + r_funcs = [r.saved_output for r in replacement_functions] for i in range(len(self._u)): assert isinstance(F_form, list) res_i = ufl.ZeroBaseForm((self._u[i],)) for j in range(len(self._u)): - res_i += ufl.action(self._lhs[i][j], replacement_functions[j].saved_output) # type: ignore[index] + if self._lhs[i][j] is not None: + res_i += ufl.action(self._lhs[i][j], r_funcs) # type: ignore[index] res_i -= self._rhs[i] # type: ignore[index] F_form.append(res_i) - # NOTE: Will fail for blocked systems atm - assert isinstance(F_form, ufl.Form) replacement_map = self._create_replace_map(F_form) if isinstance(self._u, Function): F_form = ufl.replace(F_form, replacement_map) @@ -391,18 +391,32 @@ def _compute_residual_derivative(self) -> typing.Union[ufl.Form, list[list[ufl.F assert isinstance(F_form, ufl.Form) dFdu = ufl.derivative(F_form, outputs[0], ufl.TrialFunction(outputs[0].function_space)) else: + # Replacement trial function needs to be in mixed space if initial form is created with a mixed function space. + # This means re-using the trialfunctions from the lhs + trial_functions = [None for _ in range(len(outputs))] + for i in range(len(outputs)): + for j in range(len(outputs)): + if self._lhs[i][j] is not None: + trial_functions[j] = self._lhs[i][j].arguments()[1] 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))) + dFdu_ij = ufl.derivative(F_form[i], outputs[j], trial_functions[j]) + # Apply derivatives to avoid getting empty forms later on + dFdu_ij = ufl.algorithms.apply_derivatives.apply_derivatives( + ufl.algorithms.expand_derivatives(dFdu_ij) + ) + dFdu_ij = None if dFdu_ij.empty() else dFdu_ij + dFdu[-1].append(dFdu_ij) return dFdu def prepare_evaluate_tlm( self, inputs, tlm_inputs, relevant_outputs ) -> tuple[typing.Union[list[ufl.Form], ufl.Form], dolfinx.fem.Form]: F_form = self._compute_residual() + breakpoint() dFdu_compiled = dolfinx.fem.form( self._compute_residual_derivative(), jit_options=self._jit_options, @@ -428,21 +442,33 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar bcs = [] for bc in self._bcs: bcs.append(bc) - dFdm = ufl.ZeroBaseForm((ufl.TestFunction(V),)) + + # Extract the testfunction from input space to keep the part of a mixed function space + outputs = [output.saved_output for output in self.get_outputs()] + test_functions = [None for _ in range(len(outputs))] + for i in range(len(outputs)): + test_functions[i] = self._rhs[i].arguments()[0] + + dFdm = ufl.ZeroBaseForm((test_functions[idx],)) for block_variable in self.get_dependencies(): tlm_value = block_variable.tlm_value # c = block_variable.output c_rep = block_variable.saved_output + if tlm_value is None: continue - dFdm += ufl.derivative(-F, c_rep, tlm_value) + dFdm += ufl.derivative(-F[i], c_rep, tlm_value) if isinstance(dFdm, float): v = dFdu.arguments()[0] dFdm = ufl.ZeroBaseForm((v,)) + dFdm = ufl.algorithms.apply_derivatives.apply_derivatives(dFdm) dFdm = ufl.algorithms.expand_derivatives(dFdm) + if dFdm.empty(): + dFdm = ufl.ZeroBaseForm((test_functions[idx],)) + dFdm_compiled = dolfinx.fem.form( dFdm, jit_options=self._jit_options, @@ -516,10 +542,25 @@ def prepare_evaluate_adj( dFdu_adj = self._compute_adjoint(dFdu) # Extract dJ/du[v] from the adjoint inputs. - assert len(adj_inputs) == 1 - adj_rhs = adj_inputs[0] - dJdu = dolfinx.la.vector(adj_rhs.index_map, adj_rhs.block_size) - dJdu.array[:] = adj_rhs.array[:].copy() + if len(adj_inputs) == 1: + adj_rhs = adj_inputs[0] + dJdu = dolfinx.la.vector(adj_rhs.index_map, adj_rhs.block_size) + dJdu.array[:] = adj_rhs.array[:].copy() + else: + assert len(adj_inputs) == len(self.get_outputs()), ( + f"Expected {len(self.get_outputs())} adjoint inputs, got {len(adj_inputs)})" + ) + dJdu = [] + for adj_rhs, output in zip(adj_inputs, self.get_outputs(), strict=True): + if adj_rhs is None: + dJdu_i = dolfinx.la.vector( + output.output.index_map, output.output.function_space.dofmap.index_map_bs + ) + dJdu_i.array[:] = 0.0 + else: + dJdu_i = dolfinx.la.vector(adj_rhs.index_map, adj_rhs.block_size) + dJdu_i.array[:] = adj_rhs.array[:].copy() + dJdu.append(dJdu_i) # Solve adjoint problem compiled_dFdu = dolfinx.fem.form( @@ -529,10 +570,12 @@ def prepare_evaluate_adj( entity_maps=self._entity_maps, ) self._adjoint_solver._a = compiled_dFdu - self._adjoint_solver._b = dJdu.petsc_vec + if len(adj_inputs) == 1: + self._adjoint_solver._b = dJdu.petsc_vec + else: + dolfinx.la.petsc.assign([dJdu_i.petsc_vec.array_r for dJdu_i in dJdu], self._adjoint_solver._b) self._adjoint_solver._u = self._adjoint_solutions # type: ignore[assignment] self._adjoint_solver.solve() - return F_form def evaluate_adj_component( @@ -548,14 +591,25 @@ def evaluate_adj_component( residual = prepared c = block_variable.output c_rep = block_variable.saved_output - if isinstance(c, dolfinx.fem.Function): - dc = ufl.TrialFunction(c.function_space) + # Similar construction of trial function re-using the mixed function space trial function + # Replacement trial function needs to be in mixed space if initial form is created with a mixed function space. + # This means re-using the trialfunctions from the lhs + outputs = [output.saved_output for output in self.get_outputs()] + trial_functions = [None for _ in range(len(outputs))] + for i in range(len(outputs)): + for j in range(len(outputs)): + if self._lhs[i][j] is not None: + trial_functions[j] = self._lhs[i][j].arguments()[1] + dc = trial_functions[idx] else: raise NotImplementedError(f"Unsupported control {type(c)}") # Compute the sensitivity of the residual with respect to the parameter - dFdm = -ufl.derivative(residual, c_rep, dc) + if isinstance(residual, list): + dFdm = -ufl.derivative(residual[idx], c_rep, dc) + else: + dFdm = -ufl.derivative(residual, c_rep, dc) if dFdm.empty(): # Generate a dummy form to safely extract the correct Vector wrapper type dFdm = dolfinx.fem.form(ufl.ZeroBaseForm((dc,))) @@ -588,11 +642,12 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ assert len(outputs) == 1, "Hessian computation only implemented for single output blocks." assert len(tlm_output) == 1, "Hessian computation only implemented for single TLM output blocks." d2Fdu2 = ufl.algorithms.expand_derivatives(ufl.derivative(dFdu_form, outputs[0].saved_output, tlm_output[0])) - + breakpoint() # bdy = self._should_compute_boundary_adjoint(relevant_dependencies) assert len(hessian_inputs) == 1, "Hessian computation only implemented for single hessian input blocks." # Assemble right hand side of second order adjoint equation + # Note this term should always be zero for linear problems, but we include it for completeness. b_form = d2Fdu2 if d2Fdu2.empty() else ufl.action(ufl.adjoint(d2Fdu2), self._adjoint_solutions) for bo in self.get_dependencies(): c = bo.output From 271fc57b85d31185ae7c3f7dfb292c7247a9d9d5 Mon Sep 17 00:00:00 2001 From: Joergen Schartum Dokken Date: Wed, 26 Aug 2026 10:25:58 +0000 Subject: [PATCH 04/26] Fix adjoint and tlm for blocked problems. --- src/dolfinx_adjoint/blocks/interpolation.py | 2 +- src/dolfinx_adjoint/blocks/solvers.py | 313 +++++++++++--------- src/dolfinx_adjoint/types/function.py | 2 +- tests/test_blocked_problem.py | 14 +- tests/test_nonlinear_problem.py | 1 + 5 files changed, 191 insertions(+), 141 deletions(-) diff --git a/src/dolfinx_adjoint/blocks/interpolation.py b/src/dolfinx_adjoint/blocks/interpolation.py index e7ce1ce..06ca6ab 100644 --- a/src/dolfinx_adjoint/blocks/interpolation.py +++ b/src/dolfinx_adjoint/blocks/interpolation.py @@ -13,8 +13,8 @@ from ufl.algorithms.analysis import traverse_unique_terminals from ..compat import get_interpolation_points +from ..types.function import Function, _create_function from ..utils import unroll_dofmap -from ..types.function import _create_function, Function if typing.TYPE_CHECKING: from petsc4py import PETSc diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index 4e0c199..63ff789 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -172,10 +172,12 @@ def __init__( 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] self._adjoint_solver = LinearAdjointProblem( self._compute_adjoint(self._lhs), # type: ignore[arg-type] @@ -190,6 +192,7 @@ def __init__( kind=kind, # type: ignore[arg-type] entity_maps=self._entity_maps, ) # type: ignore[misc] + self._tlm_solver = None def _recover_bcs(self): bcs = [] @@ -201,6 +204,23 @@ def _recover_bcs(self): bcs.append(c_rep) return bcs + def construct_tlm_solver(self): + dFdu_form = self._compute_residual_derivative() + tlm_solver = LinearAdjointProblem( + dFdu_form, # type: ignore[arg-type] + self._rhs, # type: ignore[arg-type] + bcs=self._bcs, + u=self._tlm_solutions, # type: ignore[arg-type] + P=self._preconditioner, # type: ignore[arg-type] + form_compiler_options=self._form_compiler_options, + jit_options=self._jit_options, + petsc_options=self._tlm_petsc_options, + petsc_options_prefix=self._petsc_options_prefix, + kind=self._kind, # type: ignore[arg-type] + entity_maps=self._entity_maps, + ) # type: ignore[misc] + return tlm_solver + def _create_replace_map(self, form: ufl.Form | typing.Iterable[ufl.Form] | None) -> dict[Function, Function]: """Replace dependencies with latest checkpoint.""" replace_map = {} @@ -391,7 +411,8 @@ def _compute_residual_derivative(self) -> typing.Union[ufl.Form, list[list[ufl.F assert isinstance(F_form, ufl.Form) dFdu = ufl.derivative(F_form, outputs[0], ufl.TrialFunction(outputs[0].function_space)) else: - # Replacement trial function needs to be in mixed space if initial form is created with a mixed function space. + # Replacement trial function needs to be in mixed space if initial form is created + # with a mixed function space. # This means re-using the trialfunctions from the lhs trial_functions = [None for _ in range(len(outputs))] for i in range(len(outputs)): @@ -415,59 +436,43 @@ def _compute_residual_derivative(self) -> typing.Union[ufl.Form, list[list[ufl.F def prepare_evaluate_tlm( self, inputs, tlm_inputs, relevant_outputs ) -> tuple[typing.Union[list[ufl.Form], ufl.Form], dolfinx.fem.Form]: - F_form = self._compute_residual() - breakpoint() - dFdu_compiled = dolfinx.fem.form( - self._compute_residual_derivative(), - jit_options=self._jit_options, - form_compiler_options=self._form_compiler_options, - entity_maps=self._entity_maps, - ) - return F_form, dFdu_compiled # type: ignore[return-value] - - def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepared=None) -> dolfinx.fem.Function: - """Solve the TLM equation for the block variable. - - .. math:: - - \frac{\\partial F}{\\partial u} \frac{\\partial u}{\\partial m} = \frac{\\partial F}{\\partial m} - - """ - # FIXME: Think about blocks later - F, dFdu = prepared - V = self.get_outputs()[idx].output.function_space - - # FIXME: DirichletBC not block variable yet. Required later on. Currently all bcs should be homogenized - bcs = [] - for bc in self._bcs: - bcs.append(bc) - - # Extract the testfunction from input space to keep the part of a mixed function space - outputs = [output.saved_output for output in self.get_outputs()] - test_functions = [None for _ in range(len(outputs))] - for i in range(len(outputs)): - test_functions[i] = self._rhs[i].arguments()[0] + F_form = self._compute_residual() + if self._tlm_solver is None: + self._tlm_solver = self.construct_tlm_solver() + # Build RHS (dFdm) for the monolithic system + if isinstance(self._u, list): + test_funcs = [self._rhs[i].arguments()[0] for i in range(len(self._u))] + dFdm = [ufl.ZeroBaseForm((test,)) for test in test_funcs] + else: + test_funcs = [self._rhs.arguments()[0]] + dFdm = ufl.ZeroBaseForm((test_funcs[0],)) - dFdm = ufl.ZeroBaseForm((test_functions[idx],)) for block_variable in self.get_dependencies(): tlm_value = block_variable.tlm_value - # c = block_variable.output c_rep = block_variable.saved_output - if tlm_value is None: continue - dFdm += ufl.derivative(-F[i], c_rep, tlm_value) - - if isinstance(dFdm, float): - v = dFdu.arguments()[0] - dFdm = ufl.ZeroBaseForm((v,)) + # Accumulate sensitivities across all block components + if isinstance(self._u, list): + for i in range(len(self._u)): + term = ufl.derivative(-F_form[i], c_rep, tlm_value) + dFdm[i] += term + else: + term = ufl.derivative(-F_form, c_rep, tlm_value) + dFdm += term - dFdm = ufl.algorithms.apply_derivatives.apply_derivatives(dFdm) - dFdm = ufl.algorithms.expand_derivatives(dFdm) - if dFdm.empty(): - dFdm = ufl.ZeroBaseForm((test_functions[idx],)) + # Safely wrap zero forms to prevent compilation crashes + if isinstance(self._u, list): + for i in range(len(self._u)): + dFdm[i] = ufl.algorithms.expand_derivatives(dFdm[i]) + if dFdm[i] == 0 or dFdm[i].empty(): + dFdm[i] = ufl.ZeroBaseForm((test_funcs[i],)) + else: + dFdm = ufl.algorithms.expand_derivatives(dFdm) + if dFdm == 0 or dFdm.empty(): + dFdm = ufl.ZeroBaseForm((test_funcs[0],)) dFdm_compiled = dolfinx.fem.form( dFdm, @@ -475,58 +480,41 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar form_compiler_options=self._form_compiler_options, entity_maps=self._entity_maps, ) - dudm = dolfinx.fem.Function(V, name="du_dm_tlm_linearblock") - # Create and assemble TLM matrix - if not hasattr(self, "_A_tlm"): - self._A_tlm = dolfinx.fem.petsc.create_matrix(dFdu) - - self._A_tlm.zeroEntries() - dolfinx.fem.petsc.assemble_matrix(self._A_tlm, dFdu, bcs=bcs) # type: ignore[misc,arg-type] - self._A_tlm.assemble() - - # Create TLM KSP and attach matrix - if not hasattr(self, "_ksp_tlm"): - self._ksp_tlm = PETSc.KSP().create(self._A_tlm.getComm()) - self._ksp_tlm.setOperators(self._A_tlm) - - # Set TLM solver options - if self._tlm_petsc_options is not None: - prefix = self._petsc_options_prefix + "tlm_" - self._ksp_tlm.setOptionsPrefix(prefix) - opts = PETSc.Options() - opts.prefixPush(prefix) - for k, v in self._tlm_petsc_options.items(): - opts.setValue(k, v) - self._ksp_tlm.setFromOptions() - opts.prefixPop() - - # For some strange reason delValue doesn't respect prefixes - for k, v in self._tlm_petsc_options.items(): - opts.delValue(f"{prefix}{k}") - # Setup preconditioner - self._ksp_tlm.setUp() - - b_tlm = dolfinx.fem.create_vector(dolfinx.fem.extract_function_spaces(dFdm_compiled)) - b_tlm.array[:] = 0.0 - dolfinx.fem.petsc.assemble_vector(b_tlm.petsc_vec, dFdm_compiled) - if bcs is not None: - # This system should never be "blocked" - dolfinx.fem.petsc.apply_lifting(b_tlm.petsc_vec, [dFdu], bcs=[bcs], alpha=0) - dolfinx.la.petsc._ghost_update(b_tlm.petsc_vec, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) # type: ignore[arg-type] - for bc in bcs: - bc.set(b_tlm.array, alpha=0) - else: - dolfinx.la.petsc._ghost_update(b_tlm, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) # type: ignore[arg-type] + # 3. Assemble RHS Vector utilizing the internal block-allocated vector + b_petsc = self._tlm_solver._b + with b_petsc.localForm() as b_loc: + b_loc.set(0.0) - # Use the cached solver to skip reallocation and factorization! - self._ksp_tlm.solve(b_tlm.petsc_vec, dudm.x.petsc_vec) - dudm.x.scatter_forward() + dolfinx.fem.petsc.assemble_vector(b_petsc, dFdm_compiled) + dolfinx.la.petsc._ghost_update(b_petsc, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) - # Explicitly free the temporary RHS vector memory - b_tlm.petsc_vec.destroy() + # 4. Apply Homogeneous Boundary Conditions safely directly to the block vector + if self._bcs: + try: + for bc in self._bcs: + bc.set(b_petsc.array_w, alpha=0.0) + except RuntimeError: + # FEniCSx throws RuntimeError for flat .set() on a blocked array. + # We use the bcs_by_block utility to handle the nested extraction. + from dolfinx.fem.bcs import bcs_by_block - return dudm + V_ext = dolfinx.fem.extract_function_spaces(dFdm_compiled) + bcs_lift = bcs_by_block(V_ext, self._bcs) + dolfinx.fem.petsc.set_bc(b_petsc, bcs_lift, alpha=0.0) + + # 5. Solve the full monolithic TLM system + self._tlm_solver.solve() + + 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: + return self._tlm_solutions def prepare_evaluate_adj( self, @@ -592,16 +580,9 @@ def evaluate_adj_component( c = block_variable.output c_rep = block_variable.saved_output if isinstance(c, dolfinx.fem.Function): - # Similar construction of trial function re-using the mixed function space trial function - # Replacement trial function needs to be in mixed space if initial form is created with a mixed function space. - # This means re-using the trialfunctions from the lhs - outputs = [output.saved_output for output in self.get_outputs()] - trial_functions = [None for _ in range(len(outputs))] - for i in range(len(outputs)): - for j in range(len(outputs)): - if self._lhs[i][j] is not None: - trial_functions[j] = self._lhs[i][j].arguments()[1] - dc = trial_functions[idx] + # Need some clever construction of the TrialFunction to get a part of the mixed space + part = idx if isinstance(self._u, list) else None + dc = ufl.TrialFunction(c_rep.function_space, part=part) else: raise NotImplementedError(f"Unsupported control {type(c)}") @@ -627,7 +608,6 @@ def evaluate_adj_component( 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): @@ -639,16 +619,25 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ # Using the equation Form we derive dF/du, d^2F/du^2 * du/dm * direction. dFdu_form = self._compute_residual_derivative() - assert len(outputs) == 1, "Hessian computation only implemented for single output blocks." - assert len(tlm_output) == 1, "Hessian computation only implemented for single TLM output blocks." - d2Fdu2 = ufl.algorithms.expand_derivatives(ufl.derivative(dFdu_form, outputs[0].saved_output, tlm_output[0])) - breakpoint() + + # For linear forms d2Fdu2 is zero, but we include it for completeness. + if isinstance(dFdu_form, list): + summed = sum(sum(dFdu_ij for dFdu_ij in dFdu_i if dFdu_ij is not None) for dFdu_i in dFdu_form) + unknowns = [output.saved_output for output in self.get_outputs()] + d2Fdu2 = ufl.algorithms.expand_derivatives(ufl.derivative(summed, unknowns, tlm_output)) + else: + d2Fdu2 = ufl.algorithms.expand_derivatives( + ufl.derivative(dFdu_form, outputs[0].saved_output, tlm_output[0]) + ) + # bdy = self._should_compute_boundary_adjoint(relevant_dependencies) - assert len(hessian_inputs) == 1, "Hessian computation only implemented for single hessian input blocks." # Assemble right hand side of second order adjoint equation # Note this term should always be zero for linear problems, but we include it for completeness. + if not d2Fdu2.empty(): + raise RuntimeError(f"This term {d2Fdu2:s} should be zero for linear problems.") b_form = d2Fdu2 if d2Fdu2.empty() else ufl.action(ufl.adjoint(d2Fdu2), self._adjoint_solutions) + b_form = len(outputs) * [b_form] for bo in self.get_dependencies(): c = bo.output c_rep = bo.saved_output @@ -658,28 +647,82 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ if isinstance(c, (dolfinx.mesh.Mesh, dolfinx.fem.DirichletBC)): raise NotImplementedError(f"Hessian computation for {type(c)} control not implemented yet.") else: - dFdu_adj = ufl.action(ufl.adjoint(dFdu_form), self._adjoint_solutions) - b_form += ufl.derivative(dFdu_adj, c_rep, tlm_input) + dFdu_adj = self._compute_adjoint(dFdu_form) + if isinstance(dFdu_form, list): + dFdu_adj_applied = [ + ufl.ZeroBaseForm((self._adjoint_solutions[i],)) for i in range(len(self._adjoint_solutions)) + ] + + for i, dFdu_i in enumerate(dFdu_adj): + for j, dFdu_ij in enumerate(dFdu_i): + if dFdu_ij is not None and not dFdu_ij.empty(): + test_part, trial_part = dFdu_ij.arguments() + dFdu_ij_rep = ufl.replace( + dFdu_ij, + { + test_part: ufl.TestFunction(test_part.ufl_function_space()), + trial_part: ufl.TrialFunction(trial_part.ufl_function_space()), + }, + ) + dFdu_adj_applied[i] += ufl.action(dFdu_ij_rep, self._adjoint_solutions[j]) + for i in range(len(b_form)): + b_form[i] += ufl.derivative(dFdu_adj_applied[i], c_rep, tlm_input) + else: + dFdu_adj_applied = ufl.action(dFdu_adj, self._adjoint_solutions) + b_form[0] += ufl.derivative(dFdu_adj, c_rep, tlm_input) - b = dolfinx.la.vector(hessian_inputs[0].index_map, hessian_inputs[0].block_size) - b.array[:] = 0.0 - if not ufl.algorithms.apply_derivatives.apply_derivatives(b_form).empty(): - compiled_soa_rhs = dolfinx.fem.form( - b_form, - jit_options=self._jit_options, - form_compiler_options=self._form_compiler_options, - entity_maps=self._entity_maps, - ) - dolfinx.fem.petsc.assemble_vector(b.petsc_vec, compiled_soa_rhs) - b.scatter_reverse(dolfinx.la.InsertMode.add) - b.array[:] *= -1 + if len(outputs) == 1: + b = dolfinx.la.vector(hessian_inputs[0].index_map, hessian_inputs[0].block_size) + b.array[:] = 0.0 + b_form = ufl.algorithms.apply_derivatives.apply_derivatives(b_form[0]) + if not b_form.empty(): + compiled_soa_rhs = dolfinx.fem.form( + b_form, + jit_options=self._jit_options, + form_compiler_options=self._form_compiler_options, + entity_maps=self._entity_maps, + ) + dolfinx.fem.petsc.assemble_vector(b.petsc_vec, compiled_soa_rhs) + b.scatter_reverse(dolfinx.la.InsertMode.add) + b.array[:] *= -1 + + b.array[:] += hessian_inputs[0].array + b.scatter_forward() + self._adjoint_solver._b = b.petsc_vec - b.array[:] += hessian_inputs[0].array - b.scatter_forward() + else: + bs = [] + 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 + b_form[i] = ufl.algorithms.apply_derivatives.apply_derivatives(b_form[i]) + if not b_form[i].empty(): + compiled_soa_rhs = dolfinx.fem.form( + b_form[i], + jit_options=self._jit_options, + form_compiler_options=self._form_compiler_options, + entity_maps=self._entity_maps, + ) + dolfinx.fem.petsc.assemble_vector(bi.petsc_vec, compiled_soa_rhs) + bi.scatter_reverse(dolfinx.la.InsertMode.add) + bi.array[:] *= -1 + if hess_input is not None: + bi.array[:] += hess_input.array + bi.scatter_forward() + b = self._adjoint_solver.b + dolfinx.la.petsc.assign([bi.petsc_vec.array_r for bi in bs], b) # Compile SOA LHS dFdu_adj = dolfinx.fem.form( - ufl.adjoint(dFdu_form), + self._compute_adjoint(dFdu_form), jit_options=self._jit_options, form_compiler_options=self._form_compiler_options, entity_maps=self._entity_maps, @@ -687,7 +730,6 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ # Solve adjoint problem self._adjoint_solver._a = dFdu_adj - self._adjoint_solver._b = b.petsc_vec self._adjoint_solver._u = self._second_adjoint_solutions self._adjoint_solver.solve() return self._compute_residual(), self._adjoint_solutions, self._second_adjoint_solutions @@ -707,8 +749,7 @@ def evaluate_hessian_component( F_form, adj_sol, adj_sol2 = prepared outputs = self.get_outputs() - assert len(outputs) == 1, "Hessian computation only implemented for single output blocks." - tlm_output = outputs[0].tlm_value + tlm_output = outputs[idx].tlm_value c_rep = block_variable.saved_output @@ -731,8 +772,12 @@ def evaluate_hessian_component( W = c.function_space dc = ufl.TestFunction(W) - form_adj = ufl.action(F_form, adj_sol) - form_adj2 = ufl.action(F_form, adj_sol2) + if isinstance(F_form, list): + form_adj = sum(ufl.action(F_form[i], adj_sol) for i in range(len(F_form))) + form_adj2 = sum(ufl.action(F_form[i], adj_sol2) for i in range(len(F_form))) + else: + form_adj = ufl.action(F_form, adj_sol) + form_adj2 = ufl.action(F_form, adj_sol2) if isinstance(c, dolfinx.mesh.Mesh): raise NotImplementedError("Hessian computation for Mesh control not implemented yet.") # dFdm_adj = ufl.derivative(form_adj, X, dc) @@ -743,7 +788,7 @@ def evaluate_hessian_component( dFdm_adj2 = ufl.derivative(form_adj2, c_rep, dc) # TODO: Old comment claims this might break on split. Confirm if true or not. - d2Fdudm = ufl.algorithms.expand_derivatives(ufl.derivative(dFdm_adj, outputs[0].saved_output, tlm_output)) + d2Fdudm = ufl.algorithms.expand_derivatives(ufl.derivative(dFdm_adj, outputs[idx].saved_output, tlm_output)) d2Fdm2 = 0 # We need to add terms from every other dependency @@ -778,7 +823,7 @@ def evaluate_hessian_component( form_compiler_options=self._form_compiler_options, entity_maps=self._entity_maps, ) - hessian_output = _create_vector(compiled_hessian, hessian_form.arguments()[0].ufl_function_space()) + hessian_output = _create_vector(compiled_hessian, hessian_form.arguments()[idx].ufl_function_space()) hessian_output.array[:] = 0.0 assemble_compiled_form(compiled_hessian, hessian_output) hessian_output.array[:] *= -1.0 @@ -926,10 +971,12 @@ def __init__( if isinstance(self._u, dolfinx.fem.Function): self._adjoint_solutions = self._u.copy() # type: ignore[assignment] self._second_adjoint_solutions = self._u.copy() # type: ignore[assignment] + self._tlm_solutions = self._u.copy() # type: ignore[assignment] 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] if isinstance(F, ufl.Form): dFdu_adj = ufl.adjoint(ufl.derivative(F, u)) diff --git a/src/dolfinx_adjoint/types/function.py b/src/dolfinx_adjoint/types/function.py index f10814d..931d306 100644 --- a/src/dolfinx_adjoint/types/function.py +++ b/src/dolfinx_adjoint/types/function.py @@ -14,9 +14,9 @@ ) from pyadjoint.tape import no_annotations +from ..blocks._vector import _SpecialVector, _vector from ..blocks.assembly import assemble_compiled_form from ..utils import function_from_vector, gather -from ..blocks._vector import _vector, _SpecialVector def _create_function( diff --git a/tests/test_blocked_problem.py b/tests/test_blocked_problem.py index e3102f7..6ee2c37 100644 --- a/tests/test_blocked_problem.py +++ b/tests/test_blocked_problem.py @@ -1,6 +1,7 @@ import typing from mpi4py import MPI + import basix.ufl import dolfinx import numpy as np @@ -37,7 +38,8 @@ def test_solver(mesh_var_name: str, request, constant: typing.Union[float, int, dx = ufl.Measure("dx", domain=mesh) a = ufl.inner(ufl.grad(u), ufl.grad(v)) * dx + ufl.inner(p, ufl.div(v)) * dx + ufl.inner(q, ufl.div(u)) * dx - f = Function(V, name="control") + Z = dolfinx.fem.functionspace(mesh, ("DG", 0, (mesh.geometry.dim,))) + f = Function(Z, name="control") f.interpolate(lambda x: (np.sin(x[0]), x[1])) L = ufl.inner(f, v) * dx L += dolfinx.fem.Constant(mesh, 0.0) * q * dx @@ -75,13 +77,13 @@ def test_solver(mesh_var_name: str, request, constant: typing.Union[float, int, control = pyadjoint.Control(f) Jh = pyadjoint.ReducedFunctional(J, control) - d = Function(V) + d = Function(Z) d.interpolate(lambda x: (10 * x[0], x[1])) - e = Function(V) - e.interpolate(lambda x: (2000 * np.sin(x[0]), -1000 * x[1])) + e = Function(Z) + e.interpolate(lambda x: (1e3 * np.sin(x[1]), 1e3 * x[0])) min_rate = pyadjoint.taylor_test(Jh, d, e, dJdm=0) - assert np.isclose(min_rate, 1.0, rtol=1e-2, atol=1e-2), f"Expected convergence rate close to 1.0, got {min_rate}" + 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) @@ -92,4 +94,4 @@ def test_solver(mesh_var_name: str, request, constant: typing.Union[float, int, 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=5e-3, atol=5e-3), f"Expected convergence rate close to 3.0, got {min_rate}" + 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_nonlinear_problem.py b/tests/test_nonlinear_problem.py index 70ec9a3..7d8e37d 100644 --- a/tests/test_nonlinear_problem.py +++ b/tests/test_nonlinear_problem.py @@ -1,4 +1,5 @@ from mpi4py import MPI + import dolfinx import numpy as np import pyadjoint From 87d9b590ad217ac6d702f09065472db0d946f33f Mon Sep 17 00:00:00 2001 From: Joergen Schartum Dokken Date: Wed, 26 Aug 2026 15:45:53 +0000 Subject: [PATCH 05/26] Simplified block forms --- src/dolfinx_adjoint/blocks/solvers.py | 178 ++++++++++++-------------- src/dolfinx_adjoint/compat.py | 92 +++++++++++++ tests/test_blocked_problem.py | 44 ++++--- 3 files changed, 199 insertions(+), 115 deletions(-) diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index 63ff789..a60062c 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -12,6 +12,15 @@ from ..petsc_utils import LinearAdjointProblem, solve_linear_problem from ..types import Function from .assembly import _create_vector, _SpecialVector, assemble_compiled_form +from ..compat import compute_form_adjoint + + +def sum_form(form: typing.Sequence[typing.Sequence[ufl.Form] | ufl.Form] | ufl.Form) -> ufl.Form: + """Sum a blocked form into a single form.""" + if isinstance(form, ufl.Form): + return form + if isinstance(form, typing.Iterable): + return sum(sum_form(fi) for fi in form if fi is not None) class LinearProblemBlock(pyadjoint.Block): @@ -86,6 +95,37 @@ def __init__( 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 once that is based on a mixed functionspace. + if not isinstance(a, ufl.Form): + # Get all arguments from the RHS and LHS forms + trial_functions = len(a) * [None] + test_functions = len(a) * [None] + for i, ai in enumerate(a): + for j, aij in enumerate(ai): + if aij is not None: + test_functions[i] = aij.arguments()[0] + trial_functions[j] = aij.arguments()[1] + assert all(tf is not None for tf in trial_functions), "Not all trial functions were found." + assert all(tf is not None for tf in test_functions), "Not all test functions were found." + trial_parts = [tf.part() for tf in trial_functions] + test_parts = [tf.part() for tf in test_functions] + if any(tp is None for tp in trial_parts) or any(tp is None for tp in test_parts): + replace_map = {} + for i, tf in enumerate(trial_functions): + new_tf = ufl.TrialFunction(tf.ufl_function_space(), i) + replace_map[tf] = new_tf + for i, tf in enumerate(test_functions): + new_tf = ufl.TestFunction(tf.ufl_function_space(), i) + replace_map[tf] = new_tf + for i, ai in enumerate(a): + for j, aij in enumerate(ai): + if aij is not None: + a[i][j] = ufl.replace(aij, replace_map) + for i, Li in enumerate(L): + if Li is not None: + L[i] = ufl.replace(Li, replace_map) + self._lhs = a self._rhs = L self._preconditioner = P @@ -350,22 +390,8 @@ def _compute_adjoint( 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): - if form_ij is None or form_ij.empty(): - tmp_form[i].append(None) - adj_form[i].append(None) - else: - 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 + sum_form = sum([fij for fi in form for fij in fi if fij is not None]) + return ufl.extract_blocks(compute_form_adjoint(sum_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 @@ -373,33 +399,21 @@ def _compute_residual(self) -> typing.Union[ufl.Form, list[ufl.Form]]: """ # NOTE: Should probably be possible to compile this form once. replacement_functions = self.get_outputs() - F_form: typing.Union[ufl.Form, list[ufl.Form]] = [] + r_funcs = [r.saved_output for r in replacement_functions] if isinstance(self._u, Function): assert len(replacement_functions) == 1, ( f"Expected a single output function, got {len(replacement_functions)}" ) - F_form = ufl.action(self._lhs, replacement_functions[0].saved_output) - self._rhs + F_form = ufl.action(self._lhs, r_funcs[0]) - self._rhs else: # Blocked formulation (assuming no mixed function-space) assert len(self._u) == len(replacement_functions), ( f"Expected {len(self._u)} output functions, got {len(replacement_functions)}" ) - r_funcs = [r.saved_output for r in replacement_functions] - for i in range(len(self._u)): - assert isinstance(F_form, list) - res_i = ufl.ZeroBaseForm((self._u[i],)) - for j in range(len(self._u)): - if self._lhs[i][j] is not None: - res_i += ufl.action(self._lhs[i][j], r_funcs) # type: ignore[index] - res_i -= self._rhs[i] # type: ignore[index] - F_form.append(res_i) + 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) - if isinstance(self._u, Function): - F_form = ufl.replace(F_form, replacement_map) - else: - assert isinstance(F_form, list) - for j in range(len(F_form)): - F_form[j] = ufl.replace(F_form[j], replacement_map) + F_form = ufl.replace(F_form, replacement_map) return F_form def _compute_residual_derivative(self) -> typing.Union[ufl.Form, list[list[ufl.Form]]]: @@ -414,24 +428,11 @@ def _compute_residual_derivative(self) -> typing.Union[ufl.Form, list[list[ufl.F # Replacement trial function needs to be in mixed space if initial form is created # with a mixed function space. # This means re-using the trialfunctions from the lhs - trial_functions = [None for _ in range(len(outputs))] - for i in range(len(outputs)): - for j in range(len(outputs)): - if self._lhs[i][j] is not None: - trial_functions[j] = self._lhs[i][j].arguments()[1] - assert isinstance(F_form, list) - dFdu = [] - for i in range(len(outputs)): - dFdu.append([]) - for j in range(len(outputs)): - dFdu_ij = ufl.derivative(F_form[i], outputs[j], trial_functions[j]) - # Apply derivatives to avoid getting empty forms later on - dFdu_ij = ufl.algorithms.apply_derivatives.apply_derivatives( - ufl.algorithms.expand_derivatives(dFdu_ij) - ) - dFdu_ij = None if dFdu_ij.empty() else dFdu_ij - dFdu[-1].append(dFdu_ij) - return dFdu + trial_functions = sorted( + filter(lambda a: a.number() == 1, sum_form(self._lhs).arguments()), key=lambda a: a.part() + ) + dFdu = ufl.derivative(F_form, outputs, trial_functions) + return ufl.extract_blocks(dFdu) def prepare_evaluate_tlm( self, inputs, tlm_inputs, relevant_outputs @@ -442,8 +443,10 @@ def prepare_evaluate_tlm( self._tlm_solver = self.construct_tlm_solver() # Build RHS (dFdm) for the monolithic system if isinstance(self._u, list): - test_funcs = [self._rhs[i].arguments()[0] for i in range(len(self._u))] - dFdm = [ufl.ZeroBaseForm((test,)) for test in test_funcs] + test_funcs = sorted( + filter(lambda x: x.number() == 0, sum_form(self._rhs).arguments()), key=lambda a: a.part() + ) + dFdm = sum([ufl.ZeroBaseForm((test,)) for test in test_funcs]) else: test_funcs = [self._rhs.arguments()[0]] dFdm = ufl.ZeroBaseForm((test_funcs[0],)) @@ -455,22 +458,21 @@ def prepare_evaluate_tlm( continue # Accumulate sensitivities across all block components - if isinstance(self._u, list): - for i in range(len(self._u)): - term = ufl.derivative(-F_form[i], c_rep, tlm_value) - dFdm[i] += term - else: - term = ufl.derivative(-F_form, c_rep, tlm_value) - dFdm += term + dFdm += ufl.derivative(-F_form, c_rep, tlm_value) # Safely wrap zero forms to prevent compilation crashes + dFdm = ufl.algorithms.expand_derivatives(dFdm) if isinstance(self._u, list): - for i in range(len(self._u)): - dFdm[i] = ufl.algorithms.expand_derivatives(dFdm[i]) - if dFdm[i] == 0 or dFdm[i].empty(): - dFdm[i] = ufl.ZeroBaseForm((test_funcs[i],)) + blocks = ufl.extract_blocks(dFdm) + if len(blocks) != len(self._u): + # Some zero blocks, manually pad with zero forms + _dFdm = [ufl.ZeroBaseForm((test,)) for test in test_funcs] + for block in blocks: + args = block.arguments() + assert len(args) == 1, "Expected a single test function in the block." + _dFdm[args[0].part()] = block + dFdm = _dFdm else: - dFdm = ufl.algorithms.expand_derivatives(dFdm) if dFdm == 0 or dFdm.empty(): dFdm = ufl.ZeroBaseForm((test_funcs[0],)) @@ -527,14 +529,14 @@ def prepare_evaluate_adj( # Compute (dF/du[v])* for the linear problem. F_form = self._compute_residual() dFdu = self._compute_residual_derivative() - dFdu_adj = self._compute_adjoint(dFdu) - + dFdu_adj = compute_form_adjoint(sum_form(dFdu)) # Extract dJ/du[v] from the adjoint inputs. if len(adj_inputs) == 1: adj_rhs = adj_inputs[0] dJdu = dolfinx.la.vector(adj_rhs.index_map, adj_rhs.block_size) dJdu.array[:] = adj_rhs.array[:].copy() else: + dFdu_adj = ufl.extract_blocks(dFdu_adj) assert len(adj_inputs) == len(self.get_outputs()), ( f"Expected {len(self.get_outputs())} adjoint inputs, got {len(adj_inputs)})" ) @@ -621,10 +623,10 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ dFdu_form = self._compute_residual_derivative() # For linear forms d2Fdu2 is zero, but we include it for completeness. - if isinstance(dFdu_form, list): - summed = sum(sum(dFdu_ij for dFdu_ij in dFdu_i if dFdu_ij is not None) for dFdu_i in dFdu_form) + if isinstance(dFdu_form, tuple): unknowns = [output.saved_output for output in self.get_outputs()] - d2Fdu2 = ufl.algorithms.expand_derivatives(ufl.derivative(summed, unknowns, tlm_output)) + summed_form = sum_form(dFdu_form) + d2Fdu2 = ufl.algorithms.expand_derivatives(ufl.derivative(summed_form, unknowns, tlm_output)) else: d2Fdu2 = ufl.algorithms.expand_derivatives( ufl.derivative(dFdu_form, outputs[0].saved_output, tlm_output[0]) @@ -648,25 +650,10 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ raise NotImplementedError(f"Hessian computation for {type(c)} control not implemented yet.") else: dFdu_adj = self._compute_adjoint(dFdu_form) - if isinstance(dFdu_form, list): - dFdu_adj_applied = [ - ufl.ZeroBaseForm((self._adjoint_solutions[i],)) for i in range(len(self._adjoint_solutions)) - ] - - for i, dFdu_i in enumerate(dFdu_adj): - for j, dFdu_ij in enumerate(dFdu_i): - if dFdu_ij is not None and not dFdu_ij.empty(): - test_part, trial_part = dFdu_ij.arguments() - dFdu_ij_rep = ufl.replace( - dFdu_ij, - { - test_part: ufl.TestFunction(test_part.ufl_function_space()), - trial_part: ufl.TrialFunction(trial_part.ufl_function_space()), - }, - ) - dFdu_adj_applied[i] += ufl.action(dFdu_ij_rep, self._adjoint_solutions[j]) - for i in range(len(b_form)): - b_form[i] += ufl.derivative(dFdu_adj_applied[i], c_rep, tlm_input) + if isinstance(dFdu_form, tuple): + summed_form = sum_form(dFdu_adj) + dFdu_adj_applied = ufl.action(summed_form, self._adjoint_solutions) + b_form = ufl.extract_blocks(ufl.derivative(dFdu_adj_applied, c_rep, tlm_input)) else: dFdu_adj_applied = ufl.action(dFdu_adj, self._adjoint_solutions) b_form[0] += ufl.derivative(dFdu_adj, c_rep, tlm_input) @@ -674,10 +661,10 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ if len(outputs) == 1: b = dolfinx.la.vector(hessian_inputs[0].index_map, hessian_inputs[0].block_size) b.array[:] = 0.0 - b_form = ufl.algorithms.apply_derivatives.apply_derivatives(b_form[0]) - if not b_form.empty(): + form_i = ufl.algorithms.apply_derivatives.apply_derivatives(b_form[0]) + if not form_i.empty(): compiled_soa_rhs = dolfinx.fem.form( - b_form, + form_i, jit_options=self._jit_options, form_compiler_options=self._form_compiler_options, entity_maps=self._entity_maps, @@ -703,10 +690,10 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ ) bs.append(bi) bi.array[:] = 0.0 - b_form[i] = ufl.algorithms.apply_derivatives.apply_derivatives(b_form[i]) - if not b_form[i].empty(): + form_i = ufl.algorithms.apply_derivatives.apply_derivatives(b_form[i]) + if not form_i.empty(): compiled_soa_rhs = dolfinx.fem.form( - b_form[i], + form_i, jit_options=self._jit_options, form_compiler_options=self._form_compiler_options, entity_maps=self._entity_maps, @@ -1097,7 +1084,6 @@ def _compute_adjoint( 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 diff --git a/src/dolfinx_adjoint/compat.py b/src/dolfinx_adjoint/compat.py index a040be9..faa5e75 100644 --- a/src/dolfinx_adjoint/compat.py +++ b/src/dolfinx_adjoint/compat.py @@ -1,4 +1,10 @@ import dolfinx +from ufl.algebra import Conj +from ufl.algorithms.formsplitter import extract_blocks +from ufl.algorithms.map_integrands import map_integrands +from ufl.algorithms.replace import replace +from ufl.argument import Argument + try: from ufl.algorithms.extract_linear_combination import extract_linear_combination @@ -208,3 +214,89 @@ def get_interpolation_points(V: dolfinx.fem.FunctionSpace): return V.element.interpolation_points() # type: ignore[operator] except TypeError: return V.element.interpolation_points + + +# Workaround until https://github.com/FEniCS/ufl/pull/508 is in all stable releases we support +def compute_form_adjoint( + form, + reordered_arguments: tuple[Argument, Argument] | tuple[tuple[Argument, Argument], ...] | None = None, +): + """Compute the adjoint of a bilinear form. + + This works simply by swapping the number of the two arguments, + but keeping their elements and places in the integrand expressions. + + Args: + form: A UFL bilinear form. + reordered_arguments: Optional explicit arguments to use for the adjoint form. + - For standard finite element spaces: A single tuple `(new_u, new_v)` + representing the replacement trial and test functions. + - For mixed function spaces: A sequence of tuples, with one `(new_u, new_v)` + pair for each *subspace*. For example, `((new_u0, new_v0), (new_u1, new_v1))`. + The test function mappings are extracted using the block row index `i`, + and the trial function mappings using the block column index `j`. + + Returns: + The adjoint of the bilinear form. + """ + if form.empty(): + return form + + arguments = form.arguments() + + # Check if mixed space + is_mixed = any(arg.part() is not None for arg in arguments) + + def validate_mapping(old_v: Argument, old_u: Argument, new_v: Argument, new_u: Argument, check_parts=False): + """Validate the mapping of old arguments to new arguments.""" + if new_u.number() >= new_v.number(): + raise ValueError("Ordering of new arguments is the same as the old arguments!") + if new_u.ufl_function_space() != old_u.ufl_function_space(): + raise ValueError("Element mismatch between new and old arguments (trial functions).") + if new_v.ufl_function_space() != old_v.ufl_function_space(): + raise ValueError("Element mismatch between new and old arguments (test functions).") + + if check_parts and (new_u.part() != old_v.part() or new_v.part() != old_u.part()): + raise ValueError("Ordering of new arguments is the same as the old arguments!") + + if not is_mixed: + if len(arguments) != 2: + raise ValueError("Expecting bilinear form.") + + v, u = arguments + if v.number() >= u.number(): + raise ValueError("Mistaken assumption in code!") + if reordered_arguments is None: + assert u.part() is None and v.part() is None + new_u = Argument(u.ufl_function_space(), number=v.number()) + new_v = Argument(v.ufl_function_space(), number=u.number()) + else: + assert isinstance(reordered_arguments, tuple) and len(reordered_arguments) == 2 + u_arg, v_arg = reordered_arguments[0], reordered_arguments[1] + assert isinstance(u_arg, Argument) and isinstance(v_arg, Argument) + new_u, new_v = u_arg, v_arg + + validate_mapping(v, u, new_v, new_u, check_parts=True) + + return map_integrands(Conj, replace(form, {v: new_v, u: new_u})) + else: + form_blocked = extract_blocks(form, arity=2) + # Apply mapping block-by-block and sum + form_adj = 0 + assert isinstance(form_blocked, tuple) + for i, row in enumerate(form_blocked): + assert isinstance(row, tuple) + for j, block in enumerate(row): + if block is not None: + v, u = block.arguments() + if reordered_arguments is not None: + new_v = reordered_arguments[i][1] + new_u = reordered_arguments[j][0] + else: + new_v = Argument(v.ufl_function_space(), number=u.number(), part=v.part()) + new_u = Argument(u.ufl_function_space(), number=v.number(), part=u.part()) + local_map = {v: new_v, u: new_u} + validate_mapping(v, u, new_v, new_u) + form_adj += map_integrands(Conj, replace(block, local_map)) + + return form_adj diff --git a/tests/test_blocked_problem.py b/tests/test_blocked_problem.py index 6ee2c37..fb6e12b 100644 --- a/tests/test_blocked_problem.py +++ b/tests/test_blocked_problem.py @@ -18,31 +18,38 @@ def mesh_2D(): return dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 8, 7) -@pytest.fixture(scope="module") -def mesh_3D(): - return dolfinx.mesh.create_unit_cube(MPI.COMM_WORLD, 11, 13, 12, cell_type=dolfinx.mesh.CellType.hexahedron) - - -@pytest.mark.parametrize("constant", [np.float64(0.2), float(-0.13), int(3)]) -@pytest.mark.parametrize("mesh_var_name", ["mesh_2D"]) -def test_solver(mesh_var_name: str, request, constant: typing.Union[float, int, np.floating]): +@pytest.mark.parametrize("use_mixed_space", [True, False]) +def test_solver(use_mixed_space: bool, mesh_2D): pyadjoint.get_working_tape().clear_tape() - mesh = request.getfixturevalue(mesh_var_name) + 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) - W = ufl.MixedFunctionSpace(*[V, Q]) - u, p = ufl.TrialFunctions(W) - v, q = ufl.TestFunctions(W) dx = ufl.Measure("dx", domain=mesh) - a = ufl.inner(ufl.grad(u), ufl.grad(v)) * dx + ufl.inner(p, ufl.div(v)) * dx + ufl.inner(q, ufl.div(u)) * dx + a00 = lambda u, v: ufl.inner(ufl.grad(u), ufl.grad(v)) * dx + a01 = lambda p, v: ufl.inner(p, ufl.div(v)) * dx + a10 = lambda q, u: ufl.inner(q, ufl.div(u)) * dx + L0 = lambda f, v: ufl.inner(f, v) * dx + L1 = lambda mesh, q: dolfinx.fem.Constant(mesh, 0.0) * q * dx Z = dolfinx.fem.functionspace(mesh, ("DG", 0, (mesh.geometry.dim,))) f = Function(Z, name="control") f.interpolate(lambda x: (np.sin(x[0]), x[1])) - L = ufl.inner(f, v) * dx - L += dolfinx.fem.Constant(mesh, 0.0) * q * dx + + if use_mixed_space: + W = ufl.MixedFunctionSpace(*[V, Q]) + u, p = ufl.TrialFunctions(W) + v, q = ufl.TestFunctions(W) + a = ufl.extract_blocks(a00(u, v) + a01(p, v) + a10(q, u)) + L = ufl.extract_blocks(L0(f, v) + L1(mesh, q)) + else: + u = ufl.TrialFunction(V) + p = ufl.TrialFunction(Q) + v = ufl.TestFunction(V) + q = ufl.TestFunction(Q) + a = [[a00(u, v), a01(p, v)], [a10(q, u), None]] + L = [L0(f, v), L1(mesh, q)] mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim) boundary_facets = dolfinx.mesh.exterior_facet_indices(mesh.topology) @@ -59,8 +66,8 @@ def test_solver(mesh_var_name: str, request, constant: typing.Union[float, int, } uh, ph = (Function(V, name="state"), Function(Q, name="pressure")) problem = LinearProblem( - ufl.extract_blocks(a), - ufl.extract_blocks(L), + a, + L, u=[uh, ph], bcs=[bc], petsc_options=options, @@ -69,9 +76,8 @@ def test_solver(mesh_var_name: str, request, constant: typing.Union[float, int, ) problem.solve() - d = pyadjoint.AdjFloat(constant) x = ufl.SpatialCoordinate(mesh) - c = ufl.as_vector((d * ufl.sin(x[0]), d * ufl.cos(x[1]))) + c = ufl.as_vector((2 * ufl.sin(x[0]), 3 * ufl.cos(x[1]))) error = ufl.inner(uh - c, uh - c) * ufl.inner(uh - c, uh - c) * ufl.dx J = assemble_scalar(error) From a974f1f150fe0ee7d1651b7148dc57eede00824a Mon Sep 17 00:00:00 2001 From: Joergen Schartum Dokken Date: Wed, 26 Aug 2026 15:57:03 +0000 Subject: [PATCH 06/26] Use direct solver to ensure that initial guess doesn't influence the result of the replay. --- src/dolfinx_adjoint/blocks/solvers.py | 9 +-------- tests/test_dirichlet_bc.py | 10 ++++++++-- 2 files changed, 9 insertions(+), 10 deletions(-) diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index a60062c..9385de9 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -301,13 +301,6 @@ def _replace_coefficients_in_form( def prepare_recompute_component(self, inputs, relevant_outputs): """Prepare for recomputing the block with different control inputs.""" - # Create initial guess for the KSP solver - # Form independnet compilation would make it possible to use the same KSP for all re-evaluations. - if isinstance(self._u, Function): - initial_guess = dolfinx.fem.Function(self._u.function_space, name=self._u.name + "_initial_guess") - else: - initial_guess = [dolfinx.fem.Function(u.function_space, name=u.name + "_initial_guess") for u in self._u] - # Replace form coefficients with checkpointed values. # Loop through the dependencies of the lhs and rhs, check if they are in the respective form lhs = self._replace_coefficients_in_form(self._lhs) @@ -343,7 +336,7 @@ def prepare_recompute_component(self, inputs, relevant_outputs): self._forward_solver._L = compiled_rhs self._forward_solver._P = compiled_preconditioner self._forward_solver.bcs = self._bcs - self._forward_solver._u = initial_guess + self._forward_solver._u = self._u with pyadjoint.stop_annotating(): solution = self._forward_solver.solve() return solution diff --git a/tests/test_dirichlet_bc.py b/tests/test_dirichlet_bc.py index 7660efc..0687a11 100644 --- a/tests/test_dirichlet_bc.py +++ b/tests/test_dirichlet_bc.py @@ -118,8 +118,14 @@ def test_time_dependent_bc_replay(): # Use native dolfinx here! PyAdjoint traces the bc_func inside it. bc = dirichletbc(bc_func, boundary_dofs) - - problem = LinearProblem(a, L, bcs=[bc], u=uh) + petsc_options = { + "ksp_monitor": None, + "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) J = 0.0 From 1bca24ca569a2bc95160fb4847abf53a9baa380f Mon Sep 17 00:00:00 2001 From: Joergen Schartum Dokken Date: Wed, 26 Aug 2026 16:46:58 +0000 Subject: [PATCH 07/26] Various dolfinx.la.vectors switched to petsc for safe management. --- src/dolfinx_adjoint/blocks/solvers.py | 89 ++++++++++++++------------- 1 file changed, 47 insertions(+), 42 deletions(-) diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index 9385de9..09fad18 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -526,25 +526,29 @@ def prepare_evaluate_adj( # Extract dJ/du[v] from the adjoint inputs. if len(adj_inputs) == 1: adj_rhs = adj_inputs[0] - dJdu = dolfinx.la.vector(adj_rhs.index_map, adj_rhs.block_size) - dJdu.array[:] = adj_rhs.array[:].copy() + dJdu = self._adjoint_solver._b + with dJdu.localForm() as dJdu_loc: + dJdu_loc.set(0.0) + dJdu.array_w[:] = adj_rhs.array_r[:].copy() else: dFdu_adj = ufl.extract_blocks(dFdu_adj) assert len(adj_inputs) == len(self.get_outputs()), ( f"Expected {len(self.get_outputs())} adjoint inputs, got {len(adj_inputs)})" ) - dJdu = [] - for adj_rhs, output in zip(adj_inputs, self.get_outputs(), strict=True): + dJdu = self._adjoint_solver._b + with dJdu.localForm() as dJdu_loc: + dJdu_loc.set(0.0) + import numpy as np + + arrs = [ + np.zeros(output.output.index_map.size_local * output.output.function_space.dofmap.index_map_bs) + for output in self.get_outputs() + ] + for i, adj_rhs in enumerate(adj_inputs): if adj_rhs is None: - dJdu_i = dolfinx.la.vector( - output.output.index_map, output.output.function_space.dofmap.index_map_bs - ) - dJdu_i.array[:] = 0.0 + arrs[i][:] = 0.0 else: - dJdu_i = dolfinx.la.vector(adj_rhs.index_map, adj_rhs.block_size) - dJdu_i.array[:] = adj_rhs.array[:].copy() - dJdu.append(dJdu_i) - + arrs[i][: len(arrs[i])] = adj_rhs.array[: len(arrs[i])] # Solve adjoint problem compiled_dFdu = dolfinx.fem.form( dFdu_adj, # type: ignore[arg-type] @@ -554,9 +558,9 @@ def prepare_evaluate_adj( ) self._adjoint_solver._a = compiled_dFdu if len(adj_inputs) == 1: - self._adjoint_solver._b = dJdu.petsc_vec + self._adjoint_solver._b = dJdu else: - dolfinx.la.petsc.assign([dJdu_i.petsc_vec.array_r for dJdu_i in dJdu], self._adjoint_solver._b) + dolfinx.la.petsc.assign(arrs, self._adjoint_solver._b) self._adjoint_solver._u = self._adjoint_solutions # type: ignore[assignment] self._adjoint_solver.solve() return F_form @@ -652,8 +656,9 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ b_form[0] += ufl.derivative(dFdu_adj, c_rep, tlm_input) if len(outputs) == 1: - b = dolfinx.la.vector(hessian_inputs[0].index_map, hessian_inputs[0].block_size) - b.array[:] = 0.0 + b = self._adjoint_solver._b + with b.localForm() as b_loc: + b_loc.set(0.0) form_i = ufl.algorithms.apply_derivatives.apply_derivatives(b_form[0]) if not form_i.empty(): compiled_soa_rhs = dolfinx.fem.form( @@ -662,27 +667,26 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ form_compiler_options=self._form_compiler_options, entity_maps=self._entity_maps, ) - dolfinx.fem.petsc.assemble_vector(b.petsc_vec, compiled_soa_rhs) - b.scatter_reverse(dolfinx.la.InsertMode.add) - b.array[:] *= -1 + dolfinx.fem.petsc.assemble_vector(b, compiled_soa_rhs) + b.ghostUpdate(PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) + b.array_w[:] *= -1 - b.array[:] += hessian_inputs[0].array - b.scatter_forward() - self._adjoint_solver._b = b.petsc_vec + b.array_w[:] += hessian_inputs[0].array + b.ghostUpdate(PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) + self._adjoint_solver._b = b else: bs = [] 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) + bi = dolfinx.la.petsc.create_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, + bi = dolfinx.la.petsc.create_vector( + [(out_i.function_space.dofmap.index_map, out_i.function_space.dofmap.index_map_bs)] ) bs.append(bi) - bi.array[:] = 0.0 + bi.array_w[:] = 0.0 form_i = ufl.algorithms.apply_derivatives.apply_derivatives(b_form[i]) if not form_i.empty(): compiled_soa_rhs = dolfinx.fem.form( @@ -691,14 +695,14 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ form_compiler_options=self._form_compiler_options, entity_maps=self._entity_maps, ) - dolfinx.fem.petsc.assemble_vector(bi.petsc_vec, compiled_soa_rhs) - bi.scatter_reverse(dolfinx.la.InsertMode.add) - bi.array[:] *= -1 + dolfinx.fem.petsc.assemble_vector(bi, compiled_soa_rhs) + bi.ghostUpdate(PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) + bi.array_w[:] *= -1 if hess_input is not None: - bi.array[:] += hess_input.array - bi.scatter_forward() - b = self._adjoint_solver.b - dolfinx.la.petsc.assign([bi.petsc_vec.array_r for bi in bs], b) + bi.array_w[:] += hess_input.array[: len(bi.array_w)] + bi.ghostUpdate(PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) + b = self._adjoint_solver._b + dolfinx.la.petsc.assign([bi.array_r for bi in bs], b) # Compile SOA LHS dFdu_adj = dolfinx.fem.form( @@ -1197,8 +1201,10 @@ def prepare_evaluate_adj( # Extract dJ/du[v] from the adjoint inputs. assert len(adj_inputs) == 1 adj_rhs = adj_inputs[0] - dJdu = dolfinx.la.vector(adj_rhs.index_map, adj_rhs.block_size) - dJdu.array[:] = adj_rhs.array[:].copy() + dJdu = self._adjoint_solver._b + with dJdu.localForm() as dJdu_loc: + dJdu_loc.set(0.0) + dJdu.array_w[:] = adj_rhs.array[:].copy() # Solve adjoint problem compiled_dFdu = dolfinx.fem.form( @@ -1208,7 +1214,6 @@ def prepare_evaluate_adj( entity_maps=self._entity_maps, ) self._adjoint_solver._a = compiled_dFdu - self._adjoint_solver._b = dJdu.petsc_vec self._adjoint_solver._u = self._adjoint_solutions # type: ignore[assignment] self._adjoint_solver.solve() return F_form @@ -1281,8 +1286,8 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ dFdu_adj = ufl.action(ufl.adjoint(dFdu_form), self._adjoint_solutions) b_form += ufl.derivative(dFdu_adj, c_rep, tlm_input) - b = dolfinx.la.vector(hessian_inputs[0].index_map, hessian_inputs[0].block_size) - b.array[:] = 0.0 + b = dolfinx.la.create_vector([(hessian_inputs[0].index_map, hessian_inputs[0].block_size)]) + b.array_w[:] = 0.0 if not ufl.algorithms.apply_derivatives.apply_derivatives(b_form).empty(): compiled_soa_rhs = dolfinx.fem.form( b_form, @@ -1292,9 +1297,9 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ ) dolfinx.fem.petsc.assemble_vector(b.petsc_vec, compiled_soa_rhs) b.scatter_reverse(dolfinx.la.InsertMode.add) - b.array[:] *= -1 + b.array_w[:] *= -1 - b.array[:] += hessian_inputs[0].array + b.array_w[:] += hessian_inputs[0].array # Compile SOA LHS dFdu_adj = dolfinx.fem.form( @@ -1305,7 +1310,7 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ ) self._adjoint_solver._a = dFdu_adj - self._adjoint_solver._b = b.petsc_vec + self._adjoint_solver._b = b self._adjoint_solver._u = self._second_adjoint_solutions self._adjoint_solver.solve() return self._compute_residual(), self._adjoint_solutions, self._second_adjoint_solutions From e508da2e501aad469f4a5906619a2141918916a1 Mon Sep 17 00:00:00 2001 From: Joergen Schartum Dokken Date: Thu, 27 Aug 2026 07:28:48 +0000 Subject: [PATCH 08/26] Use localform on petsc --- src/dolfinx_adjoint/blocks/solvers.py | 22 ++++++++++++++-------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index 09fad18..4708c73 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -7,6 +7,8 @@ import dolfinx.fem.petsc import pyadjoint import ufl +import numpy as np + from dolfinx.fem.function import Function as _Function from ..petsc_utils import LinearAdjointProblem, solve_linear_problem @@ -527,9 +529,8 @@ def prepare_evaluate_adj( if len(adj_inputs) == 1: adj_rhs = adj_inputs[0] dJdu = self._adjoint_solver._b - with dJdu.localForm() as dJdu_loc: - dJdu_loc.set(0.0) - dJdu.array_w[:] = adj_rhs.array_r[:].copy() + with dJdu.localForm() as dJdu_loc, adj_rhs.localForm() as adj_rhs_loc: + dJdu_loc.array[:] = adj_rhs_loc.array[:] else: dFdu_adj = ufl.extract_blocks(dFdu_adj) assert len(adj_inputs) == len(self.get_outputs()), ( @@ -538,7 +539,6 @@ def prepare_evaluate_adj( dJdu = self._adjoint_solver._b with dJdu.localForm() as dJdu_loc: dJdu_loc.set(0.0) - import numpy as np arrs = [ np.zeros(output.output.index_map.size_local * output.output.function_space.dofmap.index_map_bs) @@ -669,7 +669,8 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ ) dolfinx.fem.petsc.assemble_vector(b, compiled_soa_rhs) b.ghostUpdate(PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) - b.array_w[:] *= -1 + + b.scale(-1) b.array_w[:] += hessian_inputs[0].array b.ghostUpdate(PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) @@ -686,7 +687,8 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ [(out_i.function_space.dofmap.index_map, out_i.function_space.dofmap.index_map_bs)] ) bs.append(bi) - bi.array_w[:] = 0.0 + with bi.localForm() as bi_loc: + bi_loc.set(0.0) form_i = ufl.algorithms.apply_derivatives.apply_derivatives(b_form[i]) if not form_i.empty(): compiled_soa_rhs = dolfinx.fem.form( @@ -697,12 +699,16 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ ) dolfinx.fem.petsc.assemble_vector(bi, compiled_soa_rhs) bi.ghostUpdate(PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) - bi.array_w[:] *= -1 + bi.scale(-1) + if hess_input is not None: - bi.array_w[:] += hess_input.array[: len(bi.array_w)] + with bi.localForm() as bi_loc: + bi_loc.array[:] += hess_input.array bi.ghostUpdate(PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) b = self._adjoint_solver._b dolfinx.la.petsc.assign([bi.array_r for bi in bs], b) + for bi in bs: + bi.destroy() # Compile SOA LHS dFdu_adj = dolfinx.fem.form( From 654abce37db3c0a9f68cdc469ba2648468e95233 Mon Sep 17 00:00:00 2001 From: Joergen Schartum Dokken Date: Thu, 27 Aug 2026 09:42:24 +0000 Subject: [PATCH 09/26] Various fixes thanks to @finsberg and his Claude prompting. Summary of my own thoughts/changes: -As pyadjoint eagerly creates new functions, we have to recompile the forms (maybe we should use the data independent form creators, would align with redoxnics ideas). - Add convenience function to get arguments sorted by part. - Minor bug-fixes on using PETSc/non-PETSc vectors. - Correct usage of idx in evaluate_hessian_component. idx only relates to inputs. Output is in control space (always). Added check of this. - Instead of recreating vectors to attach to the tlm/adjoint solvers, instead use the vector that is created once and zero it out. Better for evyerone. --- src/dolfinx_adjoint/blocks/solvers.py | 89 ++++++++++--------- tests/test_blocked_problem.py | 11 +++ tests/test_tlm_update.py | 121 ++++++++++++++++++++++++++ 3 files changed, 182 insertions(+), 39 deletions(-) create mode 100644 tests/test_tlm_update.py diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index 4708c73..bb7ed1a 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -17,6 +17,11 @@ from ..compat import compute_form_adjoint +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 sum_form(form: typing.Sequence[typing.Sequence[ufl.Form] | ufl.Form] | ufl.Form) -> ufl.Form: """Sum a blocked form into a single form.""" if isinstance(form, ufl.Form): @@ -423,9 +428,7 @@ def _compute_residual_derivative(self) -> typing.Union[ufl.Form, list[list[ufl.F # Replacement trial function needs to be in mixed space if initial form is created # with a mixed function space. # This means re-using the trialfunctions from the lhs - trial_functions = sorted( - filter(lambda a: a.number() == 1, sum_form(self._lhs).arguments()), key=lambda a: a.part() - ) + trial_functions = get_sorted_arguments(sum_form(self._lhs).arguments(), 1) dFdu = ufl.derivative(F_form, outputs, trial_functions) return ufl.extract_blocks(dFdu) @@ -436,11 +439,17 @@ def prepare_evaluate_tlm( F_form = self._compute_residual() if self._tlm_solver is None: self._tlm_solver = self.construct_tlm_solver() + # Even if the solver is cached, we need to replace the form, as the output from pyadjoint + # is stored in a new function. + self._tlm_solver._a = dolfinx.fem.form( + self._compute_residual_derivative(), + jit_options=self._jit_options, + form_compiler_options=self._form_compiler_options, + entity_maps=self._entity_maps, + ) # Build RHS (dFdm) for the monolithic system if isinstance(self._u, list): - test_funcs = sorted( - filter(lambda x: x.number() == 0, sum_form(self._rhs).arguments()), key=lambda a: a.part() - ) + test_funcs = get_sorted_arguments(sum_form(self._rhs).arguments(), 0) dFdm = sum([ufl.ZeroBaseForm((test,)) for test in test_funcs]) else: test_funcs = [self._rhs.arguments()[0]] @@ -529,7 +538,7 @@ def prepare_evaluate_adj( if len(adj_inputs) == 1: adj_rhs = adj_inputs[0] dJdu = self._adjoint_solver._b - with dJdu.localForm() as dJdu_loc, adj_rhs.localForm() as adj_rhs_loc: + with dJdu.localForm() as dJdu_loc, adj_rhs.petsc_vec.localForm() as adj_rhs_loc: dJdu_loc.array[:] = adj_rhs_loc.array[:] else: dFdu_adj = ufl.extract_blocks(dFdu_adj) @@ -653,7 +662,7 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ b_form = ufl.extract_blocks(ufl.derivative(dFdu_adj_applied, c_rep, tlm_input)) else: dFdu_adj_applied = ufl.action(dFdu_adj, self._adjoint_solutions) - b_form[0] += ufl.derivative(dFdu_adj, c_rep, tlm_input) + b_form[0] += ufl.derivative(dFdu_adj_applied, c_rep, tlm_input) if len(outputs) == 1: b = self._adjoint_solver._b @@ -672,7 +681,8 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ b.scale(-1) - b.array_w[:] += hessian_inputs[0].array + with b.localForm() as b_loc: + b_loc.array[:] += hessian_inputs[0].array[:] b.ghostUpdate(PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) self._adjoint_solver._b = b @@ -680,15 +690,14 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ bs = [] for i, hess_input in enumerate(hessian_inputs): if hess_input is not None: - bi = dolfinx.la.petsc.create_vector([(hess_input.index_map, hess_input.block_size)]) + bi = dolfinx.la.vector(hess_input.index_map, hess_input.block_size) else: out_i = self.get_outputs()[i].saved_output - bi = dolfinx.la.petsc.create_vector( - [(out_i.function_space.dofmap.index_map, out_i.function_space.dofmap.index_map_bs)] + bi = dolfinx.la.vector( + out_i.function_space.dofmap.index_map, out_i.function_space.dofmap.index_map_bs ) bs.append(bi) - with bi.localForm() as bi_loc: - bi_loc.set(0.0) + 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( @@ -697,18 +706,18 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ form_compiler_options=self._form_compiler_options, entity_maps=self._entity_maps, ) - dolfinx.fem.petsc.assemble_vector(bi, compiled_soa_rhs) - bi.ghostUpdate(PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) - bi.scale(-1) + 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: - with bi.localForm() as bi_loc: - bi_loc.array[:] += hess_input.array - bi.ghostUpdate(PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) + bi.array[:] += hess_input.array + + bi.scatter_forward() b = self._adjoint_solver._b - dolfinx.la.petsc.assign([bi.array_r for bi in bs], b) - for bi in bs: - bi.destroy() + local_arrays = [bi.array[: bi.index_map.size_local * bi.block_size] for bi in bs] + dolfinx.la.petsc.assign(local_arrays, b) # Compile SOA LHS dFdu_adj = dolfinx.fem.form( @@ -739,7 +748,7 @@ def evaluate_hessian_component( F_form, adj_sol, adj_sol2 = prepared outputs = self.get_outputs() - tlm_output = outputs[idx].tlm_value + tlm_output = [output.tlm_value for output in outputs] c_rep = block_variable.saved_output @@ -778,7 +787,8 @@ def evaluate_hessian_component( dFdm_adj2 = ufl.derivative(form_adj2, c_rep, dc) # TODO: Old comment claims this might break on split. Confirm if true or not. - d2Fdudm = ufl.algorithms.expand_derivatives(ufl.derivative(dFdm_adj, outputs[idx].saved_output, tlm_output)) + sa = [output.saved_output for output in outputs] + d2Fdudm = ufl.algorithms.expand_derivatives(ufl.derivative(dFdm_adj, sa, tlm_output)) d2Fdm2 = 0 # We need to add terms from every other dependency @@ -813,7 +823,11 @@ def evaluate_hessian_component( form_compiler_options=self._form_compiler_options, entity_maps=self._entity_maps, ) - hessian_output = _create_vector(compiled_hessian, hessian_form.arguments()[idx].ufl_function_space()) + + 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 @@ -1145,7 +1159,6 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar \frac{\\partial F}{\\partial u} \frac{\\partial u}{\\partial m} = \frac{\\partial F}{\\partial m} """ - # FIXME: Think about blocks later F, dFdu = prepared V = self.get_outputs()[idx].output.function_space @@ -1208,9 +1221,8 @@ def prepare_evaluate_adj( assert len(adj_inputs) == 1 adj_rhs = adj_inputs[0] dJdu = self._adjoint_solver._b - with dJdu.localForm() as dJdu_loc: - dJdu_loc.set(0.0) - dJdu.array_w[:] = adj_rhs.array[:].copy() + with dJdu.localForm() as dJdu_loc, adj_rhs.petsc_vec.localForm() as adj_rhs_loc: + dJdu_loc.array[:] = adj_rhs_loc.array[:] # Solve adjoint problem compiled_dFdu = dolfinx.fem.form( @@ -1291,9 +1303,9 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ else: dFdu_adj = ufl.action(ufl.adjoint(dFdu_form), self._adjoint_solutions) b_form += ufl.derivative(dFdu_adj, c_rep, tlm_input) - - b = dolfinx.la.create_vector([(hessian_inputs[0].index_map, hessian_inputs[0].block_size)]) - b.array_w[:] = 0.0 + b = self._adjoint_solver._b + with b.localForm() as b_loc: + b_loc.set(0.0) if not ufl.algorithms.apply_derivatives.apply_derivatives(b_form).empty(): compiled_soa_rhs = dolfinx.fem.form( b_form, @@ -1301,11 +1313,11 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ form_compiler_options=self._form_compiler_options, entity_maps=self._entity_maps, ) - dolfinx.fem.petsc.assemble_vector(b.petsc_vec, compiled_soa_rhs) - b.scatter_reverse(dolfinx.la.InsertMode.add) - b.array_w[:] *= -1 - - b.array_w[:] += hessian_inputs[0].array + dolfinx.fem.petsc.assemble_vector(b, compiled_soa_rhs) + b.ghostUpdate(PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) # type: ignore [arg-type] + b.scale(-1) + with b.localForm() as b_loc, hessian_inputs[0].petsc_vec.localForm() as hess_loc: + b_loc.array[:] += hess_loc.array[:] # Compile SOA LHS dFdu_adj = dolfinx.fem.form( @@ -1316,7 +1328,6 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ ) self._adjoint_solver._a = dFdu_adj - self._adjoint_solver._b = b self._adjoint_solver._u = self._second_adjoint_solutions self._adjoint_solver.solve() return self._compute_residual(), self._adjoint_solutions, self._second_adjoint_solutions diff --git a/tests/test_blocked_problem.py b/tests/test_blocked_problem.py index fb6e12b..c4af583 100644 --- a/tests/test_blocked_problem.py +++ b/tests/test_blocked_problem.py @@ -101,3 +101,14 @@ def test_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}" + + z = Function(Z) + z.interpolate(lambda x: (np.sin(x[1]), -(x[0] ** 2))) + f = Function(Z) + f.interpolate(lambda x: (1e4 * x[0], 1e5 * np.sin(x[1]))) + Jh(z) + dJdm = Jh.derivative()._ad_dot(f) + hessian = Jh.hessian(f) + 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}" diff --git a/tests/test_tlm_update.py b/tests/test_tlm_update.py new file mode 100644 index 0000000..6d73808 --- /dev/null +++ b/tests/test_tlm_update.py @@ -0,0 +1,121 @@ +"""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. +""" + +from mpi4py import MPI + +import basix.ufl +import dolfinx +import numpy as np +import pyadjoint +import pytest +import ufl + +from dolfinx_adjoint import Function, assemble_scalar +from dolfinx_adjoint.solvers import LinearProblem + +direct_solve = { + "ksp_type": "preonly", + "pc_type": "lu", + "ksp_error_if_not_converged": True, + "pc_factor_mat_solver_type": "mumps", +} + + +@pytest.fixture(scope="module") +def mesh_2D(): + return dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 8, 7) + + +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. + """ + 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])) + + 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. + x = ufl.SpatialCoordinate(mesh) + f = 1e3 * ufl.as_vector((ufl.sin(ufl.pi * x[1]), ufl.cos(ufl.pi * x[0]))) + + a = [ + [ufl.inner(mu * ufl.grad(u), ufl.grad(v)) * dx, ufl.inner(p, ufl.div(v)) * dx], + [ufl.inner(q, ufl.div(u)) * dx, None], + ] + L = [ufl.inner(f, v) * dx, dolfinx.fem.Constant(mesh, 0.0) * q * 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) + + uh, ph = Function(V, name="velocity"), Function(Q, name="pressure") + problem = LinearProblem( + a, + L, + u=[uh, ph], + bcs=[bc], + petsc_options=direct_solve, + adjoint_petsc_options=direct_solve, + tlm_petsc_options=direct_solve, + ) + problem.solve() + + # Quartic in the state and with no constant offset, so that the Taylor remainders + # stay well above round-off: a functional dominated by a control-independent term + # bottoms out at machine precision before the third-order rate is visible. + J = assemble_scalar(ufl.inner(uh, uh) ** 2 * dx) + return pyadjoint.ReducedFunctional(J, pyadjoint.Control(mu)), Z + + +@pytest.mark.parametrize("warm_up_at_another_point", [False, True]) +def test_hessian_is_independent_of_previous_evaluation_points(warm_up_at_another_point, mesh_2D): + """The Hessian at ``m2`` must not depend on whether ``J`` was evaluated at ``m1`` first. + + Both parametrizations run the identical second-order Taylor test at ``m2``. The only + difference is a prior evaluate-and-differentiate sweep at a *different* control value, + which must not change the answer. + """ + pyadjoint.get_working_tape().clear_tape() + Jh, Z = _viscous_stokes(mesh_2D) + + m1, m2 = Function(Z), Function(Z) + m1.interpolate(lambda x: 1.0 + 0.5 * np.sin(np.pi * x[0])) + m2.interpolate(lambda x: 2.0 + 0.5 * np.cos(np.pi * x[1])) + h = Function(Z) + h.interpolate(lambda x: 0.3 + 0.2 * np.sin(3 * x[0])) + + if warm_up_at_another_point: + Jh(m1) + Jh.derivative() + Jh.hessian(h) + + Jh(m2) + dJdm = Jh.derivative()._ad_dot(h) + Hm = Jh.hessian(h)._ad_dot(h) + + min_rate = pyadjoint.taylor_test(Jh, m2, h, dJdm=dJdm, Hm=Hm) + assert np.isclose(min_rate, 3.0, rtol=0.1, atol=0.1), f"Expected convergence rate close to 3.0, got {min_rate}" From 9c5721e4094356f8d84ab32dcc3a5325c204077b Mon Sep 17 00:00:00 2001 From: Joergen Schartum Dokken Date: Thu, 27 Aug 2026 09:44:32 +0000 Subject: [PATCH 10/26] Ruff --- src/dolfinx_adjoint/blocks/solvers.py | 5 ++--- src/dolfinx_adjoint/compat.py | 1 - tests/test_blocked_problem.py | 1 - 3 files changed, 2 insertions(+), 5 deletions(-) diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index bb7ed1a..e838063 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -5,16 +5,15 @@ from petsc4py import PETSc import dolfinx.fem.petsc +import numpy as np import pyadjoint import ufl -import numpy as np - from dolfinx.fem.function import Function as _Function +from ..compat import compute_form_adjoint from ..petsc_utils import LinearAdjointProblem, solve_linear_problem from ..types import Function from .assembly import _create_vector, _SpecialVector, assemble_compiled_form -from ..compat import compute_form_adjoint def get_sorted_arguments(arguments: typing.Iterable[ufl.Argument], number: int) -> typing.Iterable[ufl.Argument]: diff --git a/src/dolfinx_adjoint/compat.py b/src/dolfinx_adjoint/compat.py index faa5e75..e8b7839 100644 --- a/src/dolfinx_adjoint/compat.py +++ b/src/dolfinx_adjoint/compat.py @@ -5,7 +5,6 @@ from ufl.algorithms.replace import replace from ufl.argument import Argument - try: from ufl.algorithms.extract_linear_combination import extract_linear_combination diff --git a/tests/test_blocked_problem.py b/tests/test_blocked_problem.py index c4af583..aeabfd3 100644 --- a/tests/test_blocked_problem.py +++ b/tests/test_blocked_problem.py @@ -1,4 +1,3 @@ -import typing from mpi4py import MPI From 4ec273313aec30716289a0844bb8ad28f8a1bc6c Mon Sep 17 00:00:00 2001 From: Joergen Schartum Dokken Date: Thu, 27 Aug 2026 09:55:09 +0000 Subject: [PATCH 11/26] Last fixes to make code nice. --- src/dolfinx_adjoint/blocks/solvers.py | 33 ++++++++++----------------- 1 file changed, 12 insertions(+), 21 deletions(-) diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index e838063..e7a046e 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -548,15 +548,15 @@ def prepare_evaluate_adj( with dJdu.localForm() as dJdu_loc: dJdu_loc.set(0.0) - arrs = [ - np.zeros(output.output.index_map.size_local * output.output.function_space.dofmap.index_map_bs) - for output in self.get_outputs() - ] - for i, adj_rhs in enumerate(adj_inputs): + arrs = [] + for adj_rhs, output in zip(adj_inputs, self.get_outputs()): + local_size = output.output.index_map.size_local * output.output.function_space.dofmap.index_map_bs if adj_rhs is None: - arrs[i][:] = 0.0 + arrs.append(np.zeros(local_size, dtype=dolfinx.default_scalar_type)) else: - arrs[i][: len(arrs[i])] = adj_rhs.array[: len(arrs[i])] + arrs.append(adj_rhs.array[:local_size]) + dolfinx.la.petsc.assign(arrs, dJdu) + # Solve adjoint problem compiled_dFdu = dolfinx.fem.form( dFdu_adj, # type: ignore[arg-type] @@ -565,10 +565,6 @@ def prepare_evaluate_adj( entity_maps=self._entity_maps, ) self._adjoint_solver._a = compiled_dFdu - if len(adj_inputs) == 1: - self._adjoint_solver._b = dJdu - else: - dolfinx.la.petsc.assign(arrs, self._adjoint_solver._b) self._adjoint_solver._u = self._adjoint_solutions # type: ignore[assignment] self._adjoint_solver.solve() return F_form @@ -594,10 +590,8 @@ def evaluate_adj_component( raise NotImplementedError(f"Unsupported control {type(c)}") # Compute the sensitivity of the residual with respect to the parameter - if isinstance(residual, list): - dFdm = -ufl.derivative(residual[idx], c_rep, dc) - else: - dFdm = -ufl.derivative(residual, c_rep, dc) + sum_res = sum_form(residual) + dFdm = -ufl.derivative(sum_res, c_rep, dc) if dFdm.empty(): # Generate a dummy form to safely extract the correct Vector wrapper type dFdm = dolfinx.fem.form(ufl.ZeroBaseForm((dc,))) @@ -770,12 +764,9 @@ def evaluate_hessian_component( W = c.function_space dc = ufl.TestFunction(W) - if isinstance(F_form, list): - form_adj = sum(ufl.action(F_form[i], adj_sol) for i in range(len(F_form))) - form_adj2 = sum(ufl.action(F_form[i], adj_sol2) for i in range(len(F_form))) - else: - form_adj = ufl.action(F_form, adj_sol) - form_adj2 = ufl.action(F_form, adj_sol2) + F_summed = sum_form(F_form) + form_adj = ufl.action(F_summed, adj_sol) + form_adj2 = ufl.action(F_summed, adj_sol2) if isinstance(c, dolfinx.mesh.Mesh): raise NotImplementedError("Hessian computation for Mesh control not implemented yet.") # dFdm_adj = ufl.derivative(form_adj, X, dc) From b5f35d4fffa4a0e852b7037bc2e3eda89ef8c035 Mon Sep 17 00:00:00 2001 From: Joergen Schartum Dokken Date: Thu, 27 Aug 2026 11:24:01 +0000 Subject: [PATCH 12/26] Further simplifications --- src/dolfinx_adjoint/blocks/solvers.py | 25 ++++++++----------------- 1 file changed, 8 insertions(+), 17 deletions(-) diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index e7a046e..3017c8f 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -622,14 +622,9 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ dFdu_form = self._compute_residual_derivative() # For linear forms d2Fdu2 is zero, but we include it for completeness. - if isinstance(dFdu_form, tuple): - 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)) - else: - d2Fdu2 = ufl.algorithms.expand_derivatives( - ufl.derivative(dFdu_form, outputs[0].saved_output, tlm_output[0]) - ) + unknowns = [output.saved_output for output in self.get_outputs()] + summed_form = sum_form(dFdu_form) + d2Fdu2 = ufl.algorithms.expand_derivatives(ufl.derivative(summed_form, unknowns, tlm_output)) # bdy = self._should_compute_boundary_adjoint(relevant_dependencies) @@ -638,7 +633,6 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ if not d2Fdu2.empty(): raise RuntimeError(f"This term {d2Fdu2:s} should be zero for linear problems.") b_form = d2Fdu2 if d2Fdu2.empty() else ufl.action(ufl.adjoint(d2Fdu2), self._adjoint_solutions) - b_form = len(outputs) * [b_form] for bo in self.get_dependencies(): c = bo.output c_rep = bo.saved_output @@ -649,19 +643,15 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ raise NotImplementedError(f"Hessian computation for {type(c)} control not implemented yet.") else: dFdu_adj = self._compute_adjoint(dFdu_form) - if isinstance(dFdu_form, tuple): - summed_form = sum_form(dFdu_adj) - dFdu_adj_applied = ufl.action(summed_form, self._adjoint_solutions) - b_form = ufl.extract_blocks(ufl.derivative(dFdu_adj_applied, c_rep, tlm_input)) - else: - dFdu_adj_applied = ufl.action(dFdu_adj, self._adjoint_solutions) - b_form[0] += ufl.derivative(dFdu_adj_applied, c_rep, tlm_input) + 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) if len(outputs) == 1: b = self._adjoint_solver._b with b.localForm() as b_loc: b_loc.set(0.0) - form_i = ufl.algorithms.apply_derivatives.apply_derivatives(b_form[0]) + form_i = ufl.algorithms.apply_derivatives.apply_derivatives(b_form) if not form_i.empty(): compiled_soa_rhs = dolfinx.fem.form( form_i, @@ -681,6 +671,7 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ else: 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) From 6bb07334285cc8570b381f850076ce6259ac9859 Mon Sep 17 00:00:00 2001 From: Joergen Schartum Dokken Date: Thu, 27 Aug 2026 11:37:48 +0000 Subject: [PATCH 13/26] Type-orama --- src/dolfinx_adjoint/blocks/solvers.py | 40 ++++++++++++++++++++------- src/dolfinx_adjoint/types/function.py | 2 +- tests/test_blocked_problem.py | 21 ++++++++++---- 3 files changed, 46 insertions(+), 17 deletions(-) diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index 3017c8f..c6acfd2 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -36,6 +36,7 @@ class LinearProblemBlock(pyadjoint.Block): """ _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 @@ -102,7 +103,8 @@ def __init__( 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 once that is based on a mixed functionspace. + # 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): # Get all arguments from the RHS and LHS forms trial_functions = len(a) * [None] @@ -114,20 +116,25 @@ def __init__( trial_functions[j] = aij.arguments()[1] assert all(tf is not None for tf in trial_functions), "Not all trial functions were found." assert all(tf is not None for tf in test_functions), "Not all test functions were found." - trial_parts = [tf.part() for tf in trial_functions] - test_parts = [tf.part() for tf in test_functions] + trial_parts = [tf.part() for tf in trial_functions] # type: ignore + test_parts = [tf.part() for tf in test_functions] # type: ignore + assert isinstance(a, typing.MutableSequence) if any(tp is None for tp in trial_parts) or any(tp is None for tp in test_parts): - replace_map = {} + replace_map: dict[ufl.Argument, ufl.Argument] = {} for i, tf in enumerate(trial_functions): - new_tf = ufl.TrialFunction(tf.ufl_function_space(), i) + assert tf is not None + new_tf = ufl.TrialFunction(tf.ufl_function_space(), part=i) replace_map[tf] = new_tf for i, tf in enumerate(test_functions): - new_tf = ufl.TestFunction(tf.ufl_function_space(), i) + assert tf is not None + new_tf = ufl.TestFunction(tf.ufl_function_space(), part=i) replace_map[tf] = new_tf for i, ai in enumerate(a): for j, aij in enumerate(ai): if aij is not None: - a[i][j] = ufl.replace(aij, replace_map) + assert isinstance(ai, typing.MutableSequence) + ai[j] = ufl.replace(aij, replace_map) + assert isinstance(L, typing.MutableSequence) for i, Li in enumerate(L): if Li is not None: L[i] = ufl.replace(Li, replace_map) @@ -294,7 +301,7 @@ def _replace_coefficients_in_form( if isinstance(form, ufl.Form): return ufl.replace(form, replace_map) elif isinstance(form, typing.Iterable): - replaced_forms = [] + replaced_forms: typing.MutableSequence[ufl.Form] = [] for f in form: if f is None: replaced_forms.append(None) @@ -303,6 +310,8 @@ def _replace_coefficients_in_form( else: replaced_forms.append(ufl.replace(f, replace_map)) return replaced_forms + else: + raise TypeError(f"Cannot replace coefficients in form of type {type(form)}") def prepare_recompute_component(self, inputs, relevant_outputs): """Prepare for recomputing the block with different control inputs.""" @@ -348,13 +357,18 @@ def prepare_recompute_component(self, inputs, relevant_outputs): return solution def recompute_component( - self, inputs: typing.Iterable[Function], block_variable, idx: int, prepared: None - ) -> typing.Union[dolfinx.fem.Function, typing.Iterable[dolfinx.fem.Function]]: + self, + inputs: typing.Iterable[Function], + block_variable, + idx: int, + prepared: dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function], + ) -> dolfinx.fem.Function: """Recompute the block with the prepared linear problem.""" if isinstance(prepared, dolfinx.fem.Function): assert idx == 0 return prepared else: + assert isinstance(prepared, typing.Iterable) return prepared[idx] def _should_compute_boundary_adjoint( @@ -440,6 +454,7 @@ def prepare_evaluate_tlm( self._tlm_solver = self.construct_tlm_solver() # Even if the solver is cached, we need to replace the form, as the output from pyadjoint # is stored in a new function. + assert isinstance(self._tlm_solver, LinearAdjointProblem) self._tlm_solver._a = dolfinx.fem.form( self._compute_residual_derivative(), jit_options=self._jit_options, @@ -519,6 +534,7 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar 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( @@ -823,6 +839,7 @@ class NonlinearProblemBlock(pyadjoint.Block): _adjoint_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] _rhs: ufl.Form | typing.Sequence[ufl.Form] @typing.overload @@ -906,10 +923,12 @@ def __init__( try: u_list = self._u if isinstance(self._u, list) else [self._u] if self._lhs is not None: + assert isinstance(self._lhs, ufl.Form) for c in self._lhs.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) @@ -1082,6 +1101,7 @@ def _compute_adjoint( 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 diff --git a/src/dolfinx_adjoint/types/function.py b/src/dolfinx_adjoint/types/function.py index 931d306..ef1c6a3 100644 --- a/src/dolfinx_adjoint/types/function.py +++ b/src/dolfinx_adjoint/types/function.py @@ -87,7 +87,7 @@ def _ad_init_object(cls, obj): return cls(obj.function_space, obj.x, obj.name) @property - def index_map(self) -> dolfinx.cpp.la.IndexMap: + def index_map(self) -> dolfinx.cpp.la.IndexMap: # type: ignore """Return the index map of the function's vector.""" return self.x.index_map diff --git a/tests/test_blocked_problem.py b/tests/test_blocked_problem.py index aeabfd3..940cc2f 100644 --- a/tests/test_blocked_problem.py +++ b/tests/test_blocked_problem.py @@ -1,4 +1,3 @@ - from mpi4py import MPI import basix.ufl @@ -26,11 +25,21 @@ def test_solver(use_mixed_space: bool, mesh_2D): V = dolfinx.fem.functionspace(mesh, el_u) Q = dolfinx.fem.functionspace(mesh, el_p) dx = ufl.Measure("dx", domain=mesh) - a00 = lambda u, v: ufl.inner(ufl.grad(u), ufl.grad(v)) * dx - a01 = lambda p, v: ufl.inner(p, ufl.div(v)) * dx - a10 = lambda q, u: ufl.inner(q, ufl.div(u)) * dx - L0 = lambda f, v: ufl.inner(f, v) * dx - L1 = lambda mesh, q: dolfinx.fem.Constant(mesh, 0.0) * q * dx + + def a00(u, v): + return ufl.inner(ufl.grad(u), ufl.grad(v)) * dx + + def a01(p, v): + return ufl.inner(p, ufl.div(v)) * dx + + def a10(q, u): + return ufl.inner(q, ufl.div(u)) * dx + + def L0(f, v): + return ufl.inner(f, v) * dx + + def L1(mesh, q): + return dolfinx.fem.Constant(mesh, 0.0) * q * dx Z = dolfinx.fem.functionspace(mesh, ("DG", 0, (mesh.geometry.dim,))) f = Function(Z, name="control") From a3a6eb645621d6601871069b81ba104a89f49dd2 Mon Sep 17 00:00:00 2001 From: Joergen Schartum Dokken Date: Thu, 27 Aug 2026 11:46:03 +0000 Subject: [PATCH 14/26] Add more typing --- src/dolfinx_adjoint/blocks/solvers.py | 16 ++++++++++------ src/dolfinx_adjoint/utils.py | 1 + 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index c6acfd2..a005d57 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -21,12 +21,12 @@ 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 sum_form(form: typing.Sequence[typing.Sequence[ufl.Form] | ufl.Form] | ufl.Form) -> ufl.Form: +def sum_form(form: typing.Iterable[typing.Iterable[ufl.Form] | ufl.Form] | ufl.Form) -> ufl.Form: """Sum a blocked form into a single form.""" if isinstance(form, ufl.Form): return form if isinstance(form, typing.Iterable): - return sum(sum_form(fi) for fi in form if fi is not None) + return sum(sum_form(fi) for fi in form if fi is not None) # type: ignore[return-value] class LinearProblemBlock(pyadjoint.Block): @@ -291,7 +291,7 @@ def _create_replace_map(self, form: ufl.Form | typing.Iterable[ufl.Form] | None) def _replace_coefficients_in_form( self, form: ufl.Form | typing.Iterable[ufl.Form] - ) -> ufl.Form | typing.Iterable[ufl.Form]: + ) -> ufl.Form | typing.Iterable[ufl.Form | None]: """Replace coefficients in the form with saved outputs. Args: @@ -301,12 +301,14 @@ def _replace_coefficients_in_form( if isinstance(form, ufl.Form): return ufl.replace(form, replace_map) elif isinstance(form, typing.Iterable): - replaced_forms: typing.MutableSequence[ufl.Form] = [] + replaced_forms: typing.MutableSequence[ufl.Form | None] = [] for f in form: if f is None: replaced_forms.append(None) elif isinstance(f, typing.Iterable): - replaced_forms.append(self._replace_coefficients_in_form(f)) + new_form = self._replace_coefficients_in_form(f) + assert isinstance(new_form, ufl.Form) + replaced_forms.append(new_form) else: replaced_forms.append(ufl.replace(f, replace_map)) return replaced_forms @@ -610,7 +612,7 @@ def evaluate_adj_component( dFdm = -ufl.derivative(sum_res, c_rep, dc) if dFdm.empty(): # Generate a dummy form to safely extract the correct Vector wrapper type - dFdm = dolfinx.fem.form(ufl.ZeroBaseForm((dc,))) + dFdm = dolfinx.fem.form(ufl.ZeroBaseForm((dc,))) # type: ignore[call-overload] dFdm_adj = ufl.adjoint(dFdm) sensitivity = ufl.action(dFdm_adj, self._adjoint_solutions) @@ -1110,6 +1112,7 @@ def _compute_residual(self) -> typing.Union[ufl.Form, list[ufl.Form]]: # 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] @@ -1119,6 +1122,7 @@ def _compute_residual(self) -> typing.Union[ufl.Form, list[ufl.Form]]: 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 diff --git a/src/dolfinx_adjoint/utils.py b/src/dolfinx_adjoint/utils.py index 5b5730e..6b0b01f 100644 --- a/src/dolfinx_adjoint/utils.py +++ b/src/dolfinx_adjoint/utils.py @@ -59,6 +59,7 @@ def assign_linear_combination(value: ufl.core.expr.Expr, function: dolfinx.fem.F function.x.array[:] = 0.0 floatifier = Floatify() for weight, func in pairs: + assert isinstance(func, dolfinx.fem.Function), "All operands in the linear combination must be Functions." if not func.function_space == function.function_space: raise ValueError("Function spaces of all functions in the linear combination must match for assignment.") function.x.array[:] += floatifier.process(weight) * func.x.array[:] From b4041ec4b3bbfb9f1bf26e803c129e0e0046ec06 Mon Sep 17 00:00:00 2001 From: Joergen Schartum Dokken Date: Thu, 27 Aug 2026 13:39:43 +0000 Subject: [PATCH 15/26] Further cleanup. Personal notes: - Using compute_adjoint action on problems that mixes ufl.MixedFunctionSpace arguments and non-mixed arguments is a bad idea (looking at you prepare_hessian). Therefore we use the trick of creating a Lagrangian and differentiate it. - Various nested type-hinting to make the recursive code mypy friendly. - Action on a form does not like lists, only place that has a bit of split handling at the momment (_compute_residual). --- src/dolfinx_adjoint/blocks/solvers.py | 138 ++++++++++++++------------ tests/test_linear_solver.py | 74 ++++++++++++++ 2 files changed, 151 insertions(+), 61 deletions(-) diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index a005d57..8216b4f 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -15,18 +15,47 @@ from ..types import Function from .assembly import _create_vector, _SpecialVector, assemble_compiled_form +type NestedMutableSequence[T] = T | typing.MutableSequence["NestedMutableSequence[T]"] +type NestedSequence[T] = T | typing.Sequence["NestedSequence[T]"] + + +def to_list(data): + if isinstance(data, (tuple, list)): + return [to_list(item) for item in data] + return data + 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 sum_form(form: typing.Iterable[typing.Iterable[ufl.Form] | ufl.Form] | ufl.Form) -> ufl.Form: +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 - if isinstance(form, typing.Iterable): - return sum(sum_form(fi) for fi in form if fi is not None) # type: ignore[return-value] + + 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 LinearProblemBlock(pyadjoint.Block): @@ -118,7 +147,7 @@ def __init__( assert all(tf is not None for tf in test_functions), "Not all test functions were found." trial_parts = [tf.part() for tf in trial_functions] # type: ignore test_parts = [tf.part() for tf in test_functions] # type: ignore - assert isinstance(a, typing.MutableSequence) + a = to_list(a) if any(tp is None for tp in trial_parts) or any(tp is None for tp in test_parts): replace_map: dict[ufl.Argument, ufl.Argument] = {} for i, tf in enumerate(trial_functions): @@ -129,6 +158,7 @@ def __init__( assert tf is not None new_tf = ufl.TestFunction(tf.ufl_function_space(), part=i) replace_map[tf] = new_tf + assert isinstance(a, typing.MutableSequence) for i, ai in enumerate(a): for j, aij in enumerate(ai): if aij is not None: @@ -274,7 +304,7 @@ def construct_tlm_solver(self): ) # type: ignore[misc] return tlm_solver - def _create_replace_map(self, form: ufl.Form | typing.Iterable[ufl.Form] | None) -> dict[Function, Function]: + 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(): @@ -290,8 +320,8 @@ def _create_replace_map(self, form: ufl.Form | typing.Iterable[ufl.Form] | None) return replace_map def _replace_coefficients_in_form( - self, form: ufl.Form | typing.Iterable[ufl.Form] - ) -> ufl.Form | typing.Iterable[ufl.Form | None]: + self, form: ufl.Form | NestedMutableSequence[ufl.Form] + ) -> ufl.Form | NestedMutableSequence[ufl.Form | None]: """Replace coefficients in the form with saved outputs. Args: @@ -300,14 +330,13 @@ def _replace_coefficients_in_form( replace_map = self._create_replace_map(form) if isinstance(form, ufl.Form): return ufl.replace(form, replace_map) - elif isinstance(form, typing.Iterable): - replaced_forms: typing.MutableSequence[ufl.Form | None] = [] + elif isinstance(form, typing.Sequence): + replaced_forms: typing.MutableSequence[NestedMutableSequence[ufl.Form | None] | ufl.Form | None] = [] for f in form: if f is None: replaced_forms.append(None) - elif isinstance(f, typing.Iterable): + elif isinstance(f, typing.Sequence): new_form = self._replace_coefficients_in_form(f) - assert isinstance(new_form, ufl.Form) replaced_forms.append(new_form) else: replaced_forms.append(ufl.replace(f, replace_map)) @@ -408,25 +437,19 @@ def _compute_adjoint( sum_form = sum([fij for fi in form for fij in fi if fij is not None]) return ufl.extract_blocks(compute_form_adjoint(sum_form)) - def _compute_residual(self) -> typing.Union[ufl.Form, list[ufl.Form]]: + def _compute_residual(self) -> 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. """ # 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 isinstance(self._u, Function): - assert len(replacement_functions) == 1, ( - f"Expected a single output function, got {len(replacement_functions)}" - ) - F_form = ufl.action(self._lhs, r_funcs[0]) - self._rhs - else: - # Blocked formulation (assuming no mixed function-space) - assert len(self._u) == len(replacement_functions), ( - f"Expected {len(self._u)} output functions, got {len(replacement_functions)}" - ) - summed_form = sum_form(self._lhs) - F_form = ufl.action(summed_form, r_funcs) - sum_form(self._rhs) + 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 @@ -436,15 +459,13 @@ def _compute_residual_derivative(self) -> typing.Union[ufl.Form, list[list[ufl.F 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: - # Replacement trial function needs to be in mixed space if initial form is created - # with a mixed function space. - # This means re-using the trialfunctions from the lhs - trial_functions = get_sorted_arguments(sum_form(self._lhs).arguments(), 1) - dFdu = ufl.derivative(F_form, outputs, trial_functions) + assert isinstance(F_form, ufl.Form), "Residual form must be a single UFL form." + 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, outputs, strict=True) + ] + dFdu = ufl.derivative(F_form, outputs, trial_functions) return ufl.extract_blocks(dFdu) def prepare_evaluate_tlm( @@ -544,13 +565,14 @@ def prepare_evaluate_adj( 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]]: + ) -> ufl.Form: """Prepare the block for evaluating the adjoint.""" # Compute (dF/du[v])* for the linear problem. F_form = self._compute_residual() dFdu = self._compute_residual_derivative() - dFdu_adj = compute_form_adjoint(sum_form(dFdu)) + summed = sum_form(dFdu) + dFdu_adj = compute_form_adjoint(summed) # Extract dJ/du[v] from the adjoint inputs. if len(adj_inputs) == 1: adj_rhs = adj_inputs[0] @@ -593,8 +615,8 @@ def evaluate_adj_component( 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]]: + prepared: ufl.Form, + ) -> _SpecialVector: """Evaluate the adjoint component, i.e. :math:`\\frac{\\partial F}{\\partial m}`.""" residual = prepared @@ -772,24 +794,24 @@ def evaluate_hessian_component( assert isinstance(c, dolfinx.fem.Function) W = c.function_space - dc = ufl.TestFunction(W) + # 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_summed = sum_form(F_form) - form_adj = ufl.action(F_summed, adj_sol) - form_adj2 = ufl.action(F_summed, adj_sol2) - if isinstance(c, dolfinx.mesh.Mesh): - raise NotImplementedError("Hessian computation for Mesh control not implemented yet.") - # dFdm_adj = ufl.derivative(form_adj, X, dc) - # dFdm_adj2 = ufl.derivative(form_adj2, X, dc) - else: - # Assume Function - dFdm_adj = ufl.derivative(form_adj, c_rep, dc) - dFdm_adj2 = ufl.derivative(form_adj2, c_rep, dc) + 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) + dL1dm = ufl.derivative(L1, c_rep, dc) + dL2dm = ufl.derivative(L2, c_rep, dc) - # TODO: Old comment claims this might break on split. Confirm if true or not. sa = [output.saved_output for output in outputs] - d2Fdudm = ufl.algorithms.expand_derivatives(ufl.derivative(dFdm_adj, sa, tlm_output)) + d2Fdudm = ufl.algorithms.expand_derivatives(ufl.derivative(dL1dm, sa, tlm_output)) - d2Fdm2 = 0 + 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: @@ -802,19 +824,13 @@ def evaluate_hessian_component( if tlm_input is None: continue - # If problem is non-linear we need to skip the output variable as a control, as we can't differentiate with - # respect to the initial guess - # if c2 == self._u and not self.linear: - # continue - - # TODO: If tlm_input is a Sum, this crashes in some instances? if isinstance(c2_rep, dolfinx.mesh.Mesh): X = ufl.SpatialCoordinate(c2_rep) - d2Fdm2 += ufl.algorithms.expand_derivatives(ufl.derivative(dFdm_adj, X, tlm_input)) + d2Fdm2 += ufl.algorithms.expand_derivatives(ufl.derivative(dL1dm, X, tlm_input)) else: - d2Fdm2 += ufl.algorithms.expand_derivatives(ufl.derivative(dFdm_adj, c2_rep, tlm_input)) + d2Fdm2 += ufl.algorithms.expand_derivatives(ufl.derivative(dL1dm, c2_rep, tlm_input)) - hessian_form = ufl.algorithms.expand_derivatives(d2Fdm2 + dFdm_adj2 + d2Fdudm) + hessian_form = ufl.algorithms.expand_derivatives(d2Fdm2 + dL2dm + d2Fdudm) compiled_hessian = dolfinx.fem.form( hessian_form, diff --git a/tests/test_linear_solver.py b/tests/test_linear_solver.py index 40d54cf..09cb830 100644 --- a/tests/test_linear_solver.py +++ b/tests/test_linear_solver.py @@ -88,3 +88,77 @@ def test_solver(mesh_var_name: str, request, constant: typing.Union[float, int, 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=5e-3, atol=5e-3), f"Expected convergence rate close to 3.0, got {min_rate}" + + +def test_linear_mixed_derivative_hessian(mesh_2D): + """Test LinearProblem with a control in the bilinear form (d2F/dudm != 0).""" + pyadjoint.get_working_tape().clear_tape() + mesh = mesh_2D + + V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) + uh = Function(V, name="state") + v = ufl.TestFunction(V) + u_trial = ufl.TrialFunction(V) + + # Control variable (conductivity) + m = Function(V, name="control") + m.interpolate(lambda x: 1.0 + x[0] ** 2 + x[1] ** 2) + + # Constant Source term + f = dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(1.0)) + + # Linear problem where the bilinear form explicitly depends on the control 'm' + a = m * ufl.inner(ufl.grad(u_trial), ufl.grad(v)) * ufl.dx + L = f * v * ufl.dx + + mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim) + boundary_facets = dolfinx.mesh.exterior_facet_indices(mesh.topology) + boundary_dofs = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim - 1, boundary_facets) + bc = dolfinx.fem.dirichletbc(dolfinx.default_scalar_type(0.0), boundary_dofs, V) + + # Forward Solve + petsc_options = { + "ksp_monitor": None, + "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, + adjoint_petsc_options=petsc_options, + tlm_petsc_options=petsc_options, + ) + problem.solve() + + # Define Objective + d = Function(V) + d.interpolate(lambda x: np.sin(np.pi * x[0])) + J = assemble_scalar(0.5 * ufl.inner(uh - d, uh - d) * ufl.dx) + + control = pyadjoint.Control(m) + Jh = pyadjoint.ReducedFunctional(J, control) + + # Perturbation + dm = Function(V) + dm.interpolate(lambda x: 2 * np.sin(x[0] * np.pi) * np.cos(x[1] * np.pi)) + + # Perturbation test + min_rate = pyadjoint.taylor_test(Jh, m, dm, 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(m) + + # Gradient Taylor Test + min_rate_grad = pyadjoint.taylor_test(Jh, m, dm) + assert np.isclose(min_rate_grad, 2.0, rtol=1e-1, atol=1e-1), f"Grad rate failed: {min_rate_grad}" + + # Hessian Taylor Test + Jh(m) + dJ = Jh.derivative()._ad_dot(dm) + H = Jh.hessian(dm)._ad_dot(dm) + min_rate_hess = pyadjoint.taylor_test(Jh, m, dm, dJdm=dJ, Hm=H) + assert np.isclose(min_rate_hess, 3.0, rtol=1e-1, atol=1e-1), f"Hessian rate failed: {min_rate_hess}" From 50fefb3528ee878f2dc33a2f86f0ce0b149c4050 Mon Sep 17 00:00:00 2001 From: Joergen Schartum Dokken Date: Thu, 27 Aug 2026 14:09:13 +0000 Subject: [PATCH 16/26] Move dFdu_adj outside of dependency loop in linear and nonlinear solver. Simplify the code for adjoint in linear solver. --- src/dolfinx_adjoint/blocks/solvers.py | 29 ++++++--------------------- 1 file changed, 6 insertions(+), 23 deletions(-) diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index 8216b4f..b033886 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -263,7 +263,7 @@ def __init__( self._tlm_solutions = [u.copy() for u in self._u] self._adjoint_solver = LinearAdjointProblem( - self._compute_adjoint(self._lhs), # type: ignore[arg-type] + self._compute_adjoint(sum_form(self._lhs)), # type: ignore[arg-type] self._rhs, # type: ignore[arg-type] bcs=self._bcs, u=self._adjoint_solutions, # type: ignore[arg-type] @@ -414,28 +414,11 @@ def _should_compute_boundary_adjoint( 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: ... - - @classmethod - def _compute_adjoint( - cls, form: typing.Union[ufl.Form, typing.Sequence[typing.Sequence[ufl.Form]]] - ) -> typing.Union[ufl.Form, typing.Sequence[typing.Sequence[ufl.Form]]]: + 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. """ - if isinstance(form, ufl.Form): - return ufl.adjoint(form) - else: - assert isinstance(form, typing.Iterable) - sum_form = sum([fij for fi in form for fij in fi if fij is not None]) - return ufl.extract_blocks(compute_form_adjoint(sum_form)) + return ufl.extract_blocks(compute_form_adjoint(form)) def _compute_residual(self) -> ufl.Form: """Convert the formulation :math:`a(u, v)=L(v)` into a residual :math:`F(u_b, v) = 0` where @@ -673,6 +656,7 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ if not d2Fdu2.empty(): raise RuntimeError(f"This term {d2Fdu2:s} should be zero for linear problems.") b_form = d2Fdu2 if d2Fdu2.empty() else ufl.action(ufl.adjoint(d2Fdu2), self._adjoint_solutions) + dFdu_adj = self._compute_adjoint(sum_form(dFdu_form)) for bo in self.get_dependencies(): c = bo.output c_rep = bo.saved_output @@ -682,7 +666,6 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ if isinstance(c, (dolfinx.mesh.Mesh, dolfinx.fem.DirichletBC)): raise NotImplementedError(f"Hessian computation for {type(c)} control not implemented yet.") else: - dFdu_adj = self._compute_adjoint(dFdu_form) 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) @@ -745,7 +728,7 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ # Compile SOA LHS dFdu_adj = dolfinx.fem.form( - self._compute_adjoint(dFdu_form), + self._compute_adjoint(sum_form(dFdu_form)), jit_options=self._jit_options, form_compiler_options=self._form_compiler_options, entity_maps=self._entity_maps, @@ -1313,6 +1296,7 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ # Assemble right hand side of second order adjoint equation b_form = d2Fdu2 if d2Fdu2.empty() else ufl.action(ufl.adjoint(d2Fdu2), self._adjoint_solutions) + dFdu_adj = ufl.action(ufl.adjoint(dFdu_form), self._adjoint_solutions) for bo in self.get_dependencies(): c = bo.output c_rep = bo.saved_output @@ -1322,7 +1306,6 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ if isinstance(c, (dolfinx.mesh.Mesh, dolfinx.fem.DirichletBC)): raise NotImplementedError(f"Hessian computation for {type(c)} control not implemented yet.") else: - dFdu_adj = ufl.action(ufl.adjoint(dFdu_form), self._adjoint_solutions) b_form += ufl.derivative(dFdu_adj, c_rep, tlm_input) b = self._adjoint_solver._b with b.localForm() as b_loc: From 874631cab8e640f9bd3df630e3756f017319ec85 Mon Sep 17 00:00:00 2001 From: Joergen Schartum Dokken Date: Thu, 27 Aug 2026 15:02:00 +0000 Subject: [PATCH 17/26] Simplify assigning mixed parts to data. Start debugging the time distributed control. --- demos/time_distributed_control.py | 39 ++++++--- src/dolfinx_adjoint/blocks/solvers.py | 112 ++++++++++++++++++-------- 2 files changed, 109 insertions(+), 42 deletions(-) diff --git a/demos/time_distributed_control.py b/demos/time_distributed_control.py index 04ed874..abef605 100644 --- a/demos/time_distributed_control.py +++ b/demos/time_distributed_control.py @@ -23,7 +23,7 @@ d = 16 * x[0] * (x[0] - 1) * x[1] * (x[1] - 1) * ufl.sin(ufl.pi * t) dt = dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(0.1)) # type: ignore -T = 1 +T = 0.3 V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) # type: ignore[arg-type] ctrls = OrderedDict() @@ -38,9 +38,12 @@ def solve_heat(ctrls): v = ufl.TestFunction(V) f = dolfinx_adjoint.Function(V, name="source") - u_0 = dolfinx_adjoint.Function(V, name="solution") - F = ((u - u_0) / dt * v + nu * ufl.inner(ufl.grad(u), ufl.grad(v)) - f * v) * ufl.dx + u_prev = dolfinx_adjoint.Function(V, name="u_prev") + uh = dolfinx_adjoint.Function(V, name="solution") + dolfinx_adjoint.assign(0.0, uh) + + F = ((u - u_prev) / dt * v + nu * ufl.inner(ufl.grad(u), ufl.grad(v)) - f * v) * ufl.dx a, L = ufl.system(F) mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim) exterior_facets = dolfinx.mesh.exterior_facet_indices(mesh.topology) @@ -48,13 +51,13 @@ def solve_heat(ctrls): bc = dolfinx.fem.dirichletbc(0.0, exterior_dofs, V) - j = 0.5 * float(dt) * dolfinx_adjoint.assemble_scalar((u_0 - d) ** 2 * ufl.dx) + j = 0.5 * float(dt) * dolfinx_adjoint.assemble_scalar((uh - d) ** 2 * ufl.dx) t_val = float(dt) problem = dolfinx_adjoint.LinearProblem( a, L, - u=u_0, + u=uh, bcs=[bc], petsc_options={ "ksp_type": "preonly", @@ -81,6 +84,7 @@ def solve_heat(ctrls): dolfinx_adjoint.assign(ctrls[t_val], f) # Update data function + dolfinx_adjoint.assign(uh, u_prev) # Solve PDE problem.solve() @@ -90,12 +94,12 @@ def solve_heat(ctrls): weight = 0.5 else: weight = 1 - j += weight * float(dt) * dolfinx_adjoint.assemble_scalar((u_0 - d) ** 2 * ufl.dx) + j += weight * float(dt) * dolfinx_adjoint.assemble_scalar((uh - d) ** 2 * ufl.dx) # Update time t_val += float(dt) dolfinx_adjoint.assign(t_val, t) - return u_0, d, j + return uh, d, j u, d, j = solve_heat(ctrls) @@ -111,16 +115,33 @@ def solve_heat(ctrls): J = j + dolfinx_adjoint.assemble_scalar(regularisation) m = [pyadjoint.Control(c) for c in ctrls.values()] - rf = pyadjoint.ReducedFunctional(J, m) +# Check accuracy of gradient and Hessian using Taylor test +with pyadjoint.stop_annotating(): + h = [dolfinx_adjoint.Function(ci.function_space) for ci in m] + for hi in h: + hi.x.array[:] = np.random.random(hi.x.array.shape) + + # Prove the gradient is mathematically exact + min_val = pyadjoint.taylor_test(rf, list(ctrls.values()), h) + assert np.isclose(min_val, 2.0, rtol=1e-2, atol=1e-2), f"Expected convergence rate close to 2.0, got {min_val}" + + # Prove the Hessian is mathematically exact + rf(list(ctrls.values())) + dJdm = sum(drfi._ad_dot(hi) for drfi, hi in zip(rf.derivative(), h, strict=True)) + H = rf.hessian(h) + dHddu = sum(Hi._ad_dot(hi) for Hi, hi in zip(H, h, strict=True)) + pyadjoint.taylor_test(rf, list(ctrls.values()), h, dJdm=dJdm, Hm=dHddu) + rf(list(ctrls.values())) + + tape = pyadjoint.get_working_tape() tape.visualise_dot("test.dot") opt_ctrls = pyadjoint.minimize( rf, method="BFGS", - # method="Newton-CG", options={"maxiter": 100, "disp": True}, ) diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index b033886..b4d6753 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -19,6 +19,82 @@ type NestedSequence[T] = T | typing.Sequence["NestedSequence[T]"] +def assign_mixed_parts( + *form_structs: typing.Sequence[NestedSequence[ufl.Form]], +) -> typing.Sequence[NestedSequence[ufl.Form]]: + """ + Recursively assigns mixed-space `part` indices to UFL Test and Trial functions + within nested iterables of forms. + + When solving monolithic block systems in FEniCSx, the UFL arguments must be tagged + with a `.part()` 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. + """ + replace_map = {} + + def _build_map(obj: NestedSequence[ufl.Form], indices: tuple[int, ...]) -> None: + """ + Recursively discover forms, tracking the depth and index of the nesting. + `indices` will be `(row,)` for vectors and `(row, col)` for matrices. + """ + if isinstance(obj, ufl.Form): + for arg in obj.arguments(): + # Only map arguments that haven't been assigned a part yet + if arg.part() is None and arg not in replace_map: + num = arg.number() + + # Because num is 0 for TestFunctions and 1 for TrialFunctions, + # it maps perfectly to our nested dimension indices! + # If num < len(indices), we have traversed deep enough to assign it. + if num < len(indices): + replace_map[arg] = ufl.Argument(arg.ufl_function_space(), number=num, part=indices[num]) + + elif isinstance(obj, typing.Iterable): + for i, item in enumerate(obj): + if item is not None: + # Append current topological index to the path and recurse + _build_map(item, indices + (i,)) + + 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 + + # 1. Build a shared map across all inputs (e.g., ensuring RHS TestFunctions + # perfectly match LHS TestFunctions) + for struct in form_structs: + _build_map(struct, ()) + + # 2. If no replacements are needed, exit early to save computation + if not replace_map: + return form_structs if len(form_structs) > 1 else form_structs[0] + + # 3. Apply the replacements and unpack if necessary + replaced = tuple(_replace(struct) for struct in form_structs) + return replaced if len(replaced) > 1 else replaced[0] + + def to_list(data): if isinstance(data, (tuple, list)): return [to_list(item) for item in data] @@ -135,39 +211,8 @@ def __init__( # 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): - # Get all arguments from the RHS and LHS forms - trial_functions = len(a) * [None] - test_functions = len(a) * [None] - for i, ai in enumerate(a): - for j, aij in enumerate(ai): - if aij is not None: - test_functions[i] = aij.arguments()[0] - trial_functions[j] = aij.arguments()[1] - assert all(tf is not None for tf in trial_functions), "Not all trial functions were found." - assert all(tf is not None for tf in test_functions), "Not all test functions were found." - trial_parts = [tf.part() for tf in trial_functions] # type: ignore - test_parts = [tf.part() for tf in test_functions] # type: ignore - a = to_list(a) - if any(tp is None for tp in trial_parts) or any(tp is None for tp in test_parts): - replace_map: dict[ufl.Argument, ufl.Argument] = {} - for i, tf in enumerate(trial_functions): - assert tf is not None - new_tf = ufl.TrialFunction(tf.ufl_function_space(), part=i) - replace_map[tf] = new_tf - for i, tf in enumerate(test_functions): - assert tf is not None - new_tf = ufl.TestFunction(tf.ufl_function_space(), part=i) - replace_map[tf] = new_tf - assert isinstance(a, typing.MutableSequence) - for i, ai in enumerate(a): - for j, aij in enumerate(ai): - if aij is not None: - assert isinstance(ai, typing.MutableSequence) - ai[j] = ufl.replace(aij, replace_map) - assert isinstance(L, typing.MutableSequence) - for i, Li in enumerate(L): - if Li is not None: - L[i] = ufl.replace(Li, replace_map) + a, L = assign_mixed_parts(a, L) + P, _ = assign_mixed_parts(P, L) if P is not None else None, None self._lhs = a self._rhs = L @@ -199,6 +244,7 @@ def __init__( 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 not in self._u: self.add_dependency(c, no_duplicates=True) From 16289d3b6cdf8d82d820af37e5b74daaec4c194f Mon Sep 17 00:00:00 2001 From: Joergen Schartum Dokken Date: Thu, 27 Aug 2026 16:31:46 +0000 Subject: [PATCH 18/26] Nighmare trying to be resolved. --- demos/time_distributed_control.py | 7 ++- src/dolfinx_adjoint/blocks/solvers.py | 89 +++++++++++++++++---------- 2 files changed, 62 insertions(+), 34 deletions(-) diff --git a/demos/time_distributed_control.py b/demos/time_distributed_control.py index abef605..11a1dfd 100644 --- a/demos/time_distributed_control.py +++ b/demos/time_distributed_control.py @@ -23,7 +23,7 @@ d = 16 * x[0] * (x[0] - 1) * x[1] * (x[1] - 1) * ufl.sin(ufl.pi * t) dt = dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(0.1)) # type: ignore -T = 0.3 +T = 1 V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) # type: ignore[arg-type] ctrls = OrderedDict() @@ -41,7 +41,6 @@ def solve_heat(ctrls): u_prev = dolfinx_adjoint.Function(V, name="u_prev") uh = dolfinx_adjoint.Function(V, name="solution") - dolfinx_adjoint.assign(0.0, uh) F = ((u - u_prev) / dt * v + nu * ufl.inner(ufl.grad(u), ufl.grad(v)) - f * v) * ufl.dx a, L = ufl.system(F) @@ -123,6 +122,10 @@ def solve_heat(ctrls): for hi in h: hi.x.array[:] = np.random.random(hi.x.array.shape) + # Prove the perturbation is mathematically exact + min_val = pyadjoint.taylor_test(rf, list(ctrls.values()), h, dJdm=0) + assert np.isclose(min_val, 1.0, rtol=5e-2, atol=1e-5), f"Expected convergence rate close to 1.0, got {min_val}" + # Prove the gradient is mathematically exact min_val = pyadjoint.taylor_test(rf, list(ctrls.values()), h) assert np.isclose(min_val, 2.0, rtol=1e-2, atol=1e-2), f"Expected convergence rate close to 2.0, got {min_val}" diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index b4d6753..3325644 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -233,27 +233,36 @@ def __init__( 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): - for c in self._lhs.coefficients(): # type: ignore - if c != self._u: # Exclude the unknown - self.add_dependency(c, no_duplicates=True) - for c in self._rhs.coefficients(): # type: ignore - if c != self._u: # Exclude the unknown - self.add_dependency(c, no_duplicates=True) + 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 not in self._u: - self.add_dependency(c, no_duplicates=True) - for i, part in enumerate(self._rhs): # type: ignore + 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 not in self._u: - self.add_dependency(c, no_duplicates=True) + 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 NotImplementedError("Blocked systems not implemented yet.") + raise RuntimeError(f"Unknown type for unknown function u={type(self._u)}.") self._compiled_lhs = dolfinx.fem.form( self._lhs, jit_options=jit_options, @@ -363,6 +372,16 @@ def _create_replace_map(self, form: ufl.Form | NestedMutableSequence[ufl.Form] | 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 _replace_coefficients_in_form( @@ -466,7 +485,7 @@ def _compute_adjoint(cls, form: ufl.Form) -> typing.Sequence[typing.Sequence[ufl """ return ufl.extract_blocks(compute_form_adjoint(form)) - def _compute_residual(self) -> ufl.Form: + def _compute_residual(self) -> tuple[ufl.Form, dict[ufl.Coefficient, ufl.Coefficient]]: """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. """ @@ -481,27 +500,28 @@ def _compute_residual(self) -> ufl.Form: 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 + return F_form, replacement_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.""" - F_form = self._compute_residual() - outputs = [output.saved_output for output in self.get_outputs()] + F_form, replacement_map = self._compute_residual() assert isinstance(F_form, ufl.Form), "Residual form must be a single UFL form." + outputs = self.get_outputs() + r_funcs = [replacement_map[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, outputs, strict=True) + for arg, output in zip(test_functions, r_funcs, strict=True) ] - dFdu = ufl.derivative(F_form, outputs, trial_functions) + dFdu = ufl.derivative(F_form, r_funcs, trial_functions) return ufl.extract_blocks(dFdu) def prepare_evaluate_tlm( self, inputs, tlm_inputs, relevant_outputs ) -> tuple[typing.Union[list[ufl.Form], ufl.Form], dolfinx.fem.Form]: - F_form = self._compute_residual() + F_form, replacement_map = self._compute_residual() if self._tlm_solver is None: self._tlm_solver = self.construct_tlm_solver() # Even if the solver is cached, we need to replace the form, as the output from pyadjoint @@ -526,7 +546,7 @@ def prepare_evaluate_tlm( c_rep = block_variable.saved_output if tlm_value is None: continue - + assert c_rep in replacement_map.values() # Accumulate sensitivities across all block components dFdm += ufl.derivative(-F_form, c_rep, tlm_value) @@ -594,14 +614,9 @@ def prepare_evaluate_adj( inputs: typing.Sequence[Function], adj_inputs: typing.Sequence[dolfinx.la.Vector], relevant_dependencies: typing.List[tuple[int, pyadjoint.block_variable.BlockVariable]], - ) -> ufl.Form: + ) -> tuple[ufl.Form, dict[ufl.Coefficient, ufl.Coefficient]]: """Prepare the block for evaluating the adjoint.""" - # Compute (dF/du[v])* for the linear problem. - F_form = self._compute_residual() - dFdu = self._compute_residual_derivative() - summed = sum_form(dFdu) - dFdu_adj = compute_form_adjoint(summed) # Extract dJ/du[v] from the adjoint inputs. if len(adj_inputs) == 1: adj_rhs = adj_inputs[0] @@ -609,7 +624,6 @@ def prepare_evaluate_adj( with dJdu.localForm() as dJdu_loc, adj_rhs.petsc_vec.localForm() as adj_rhs_loc: dJdu_loc.array[:] = adj_rhs_loc.array[:] else: - dFdu_adj = ufl.extract_blocks(dFdu_adj) assert len(adj_inputs) == len(self.get_outputs()), ( f"Expected {len(self.get_outputs())} adjoint inputs, got {len(adj_inputs)})" ) @@ -625,7 +639,15 @@ def prepare_evaluate_adj( else: arrs.append(adj_rhs.array[:local_size]) dolfinx.la.petsc.assign(arrs, dJdu) + dJdu.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD) + # Compute (dF/du[v])* for the linear problem. + F_form, replacement_map = self._compute_residual() + dFdu = self._compute_residual_derivative() + summed = sum_form(dFdu) + dFdu_adj = compute_form_adjoint(summed) + dFdu_adj = ufl.algorithms.apply_derivatives.apply_derivatives(ufl.algorithms.expand_derivatives(dFdu_adj)) + assert dFdu_adj.empty() is False, "Adjoint of dF/du[v] is empty. Check if the problem is linear." # Solve adjoint problem compiled_dFdu = dolfinx.fem.form( dFdu_adj, # type: ignore[arg-type] @@ -636,7 +658,7 @@ def prepare_evaluate_adj( self._adjoint_solver._a = compiled_dFdu self._adjoint_solver._u = self._adjoint_solutions # type: ignore[assignment] self._adjoint_solver.solve() - return F_form + return F_form, replacement_map def evaluate_adj_component( self, @@ -644,11 +666,11 @@ def evaluate_adj_component( adj_inputs: typing.Iterable[dolfinx.la.Vector], block_variable: pyadjoint.block_variable.BlockVariable, idx: int, - prepared: ufl.Form, + prepared: tuple[ufl.Form, dict[ufl.Coefficient, ufl.Coefficient]], ) -> _SpecialVector: """Evaluate the adjoint component, i.e. :math:`\\frac{\\partial F}{\\partial m}`.""" - residual = prepared + residual, replacement_map = prepared c = block_variable.output c_rep = block_variable.saved_output if isinstance(c, dolfinx.fem.Function): @@ -660,6 +682,8 @@ def evaluate_adj_component( # Compute the sensitivity of the residual with respect to the parameter sum_res = sum_form(residual) + assert c in replacement_map.keys() + assert c_rep == replacement_map[c] dFdm = -ufl.derivative(sum_res, c_rep, dc) if dFdm.empty(): # Generate a dummy form to safely extract the correct Vector wrapper type @@ -674,7 +698,6 @@ def evaluate_adj_component( 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) @@ -691,6 +714,7 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ dFdu_form = self._compute_residual_derivative() # For linear forms d2Fdu2 is zero, but we include it for completeness. + unknowns = [output.saved_output for output in self.get_outputs()] summed_form = sum_form(dFdu_form) d2Fdu2 = ufl.algorithms.expand_derivatives(ufl.derivative(summed_form, unknowns, tlm_output)) @@ -798,7 +822,7 @@ def evaluate_hessian_component( ): c = block_variable.output - F_form, adj_sol, adj_sol2 = prepared + (F_form, replacement_map), adj_sol, adj_sol2 = prepared outputs = self.get_outputs() tlm_output = [output.tlm_value for output in outputs] @@ -834,6 +858,7 @@ def evaluate_hessian_component( # 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) From 696d1f2e9d0b6e53f8be8cacb52ecb29125131cd Mon Sep 17 00:00:00 2001 From: Joergen Schartum Dokken Date: Thu, 27 Aug 2026 16:46:31 +0000 Subject: [PATCH 19/26] Add failing test. and updated AI slop code that doesn't work. --- demos/time_distributed_control.py | 57 ++++++++++++++++------- src/dolfinx_adjoint/blocks/solvers.py | 65 +++++++++++++++------------ 2 files changed, 78 insertions(+), 44 deletions(-) diff --git a/demos/time_distributed_control.py b/demos/time_distributed_control.py index 11a1dfd..5a7300d 100644 --- a/demos/time_distributed_control.py +++ b/demos/time_distributed_control.py @@ -41,7 +41,7 @@ def solve_heat(ctrls): u_prev = dolfinx_adjoint.Function(V, name="u_prev") uh = dolfinx_adjoint.Function(V, name="solution") - + dolfinx_adjoint.assign(0.0, uh) F = ((u - u_prev) / dt * v + nu * ufl.inner(ufl.grad(u), ufl.grad(v)) - f * v) * ufl.dx a, L = ufl.system(F) mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim) @@ -118,26 +118,51 @@ def solve_heat(ctrls): # Check accuracy of gradient and Hessian using Taylor test with pyadjoint.stop_annotating(): - h = [dolfinx_adjoint.Function(ci.function_space) for ci in m] + # Insert this diagnostic section into your demo right after J and m are defined: + + # 1. Generate a non-zero base control point m_pert + m_pert = [dolfinx_adjoint.Function(V, name=f"pert_ctrl_{t_val}") for t_val in ctrls.keys()] + for c in m_pert: + c.x.array[:] = np.random.uniform(0.1, 1.0, size=c.x.array.shape) + + # 2. Define random directions h + h = [pyadjoint.Control(dolfinx_adjoint.Function(V)) for _ in m] for hi in h: - hi.x.array[:] = np.random.random(hi.x.array.shape) + hi.control.x.array[:] = np.random.uniform(-0.1, 0.1, size=hi.control.x.array.shape) + + print("\n=== 1. Taylor Test at NON-ZERO Control Point ===") + min_val_pert = pyadjoint.taylor_test(rf, m_pert, h) + print(f"Convergence rate at perturbed point: {min_val_pert:.4f}") + + print("\n=== 2. Direct Finite Difference Gradient Verification ===") + eps = 1e-6 + + # Compute Adjoint Directional Derivative at m_pert + rf(m_pert) + grad_adj = rf.derivative() + adj_dir_deriv = sum(g._ad_dot(hi) for g, hi in zip(grad_adj, h, strict=True)) + + # Forward Perturbation J(m + eps*h) + m_plus = [dolfinx_adjoint.Function(V) for _ in m_pert] + for mp, m_p, hi in zip(m_plus, m_pert, h, strict=True): + mp.x.array[:] = m_p.x.array[:] + eps * hi.control.x.array[:] + J_plus = float(rf(m_plus)) - # Prove the perturbation is mathematically exact - min_val = pyadjoint.taylor_test(rf, list(ctrls.values()), h, dJdm=0) - assert np.isclose(min_val, 1.0, rtol=5e-2, atol=1e-5), f"Expected convergence rate close to 1.0, got {min_val}" + # Backward Perturbation J(m - eps*h) + m_minus = [dolfinx_adjoint.Function(V) for _ in m_pert] + for mm, m_p, hi in zip(m_minus, m_pert, h, strict=True): + mm.x.array[:] = m_p.x.array[:] - eps * hi.control.x.array[:] + J_minus = float(rf(m_minus)) - # Prove the gradient is mathematically exact - min_val = pyadjoint.taylor_test(rf, list(ctrls.values()), h) - assert np.isclose(min_val, 2.0, rtol=1e-2, atol=1e-2), f"Expected convergence rate close to 2.0, got {min_val}" + # Central Finite Difference + fd_dir_deriv = (J_plus - J_minus) / (2 * eps) - # Prove the Hessian is mathematically exact - rf(list(ctrls.values())) - dJdm = sum(drfi._ad_dot(hi) for drfi, hi in zip(rf.derivative(), h, strict=True)) - H = rf.hessian(h) - dHddu = sum(Hi._ad_dot(hi) for Hi, hi in zip(H, h, strict=True)) - pyadjoint.taylor_test(rf, list(ctrls.values()), h, dJdm=dJdm, Hm=dHddu) - rf(list(ctrls.values())) + print(f"Adjoint Directional Derivative: {adj_dir_deriv:.10e}") + print(f"Finite Difference Directional Dev: {fd_dir_deriv:.10e}") + rel_diff = abs(adj_dir_deriv - fd_dir_deriv) / (abs(fd_dir_deriv) + 1e-15) + print(f"Relative Mismatch: {rel_diff:.4e}") + assert rel_diff < 1e-4, f"Adjoint gradient mismatches finite differences! Relative error: {rel_diff}" tape = pyadjoint.get_working_tape() tape.visualise_dot("test.dot") diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index 3325644..f354587 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -384,41 +384,47 @@ def _create_replace_map(self, form: ufl.Form | NestedMutableSequence[ufl.Form] | replace_map.update(self._create_replace_map(f)) return replace_map - def _replace_coefficients_in_form( - self, form: ufl.Form | NestedMutableSequence[ufl.Form] - ) -> ufl.Form | NestedMutableSequence[ufl.Form | None]: - """Replace coefficients in the form with saved outputs. + def _create_recompute_replace_map(self, inputs: typing.Sequence[typing.Any]) -> dict: + """Map original dependency coefficients to active recomputation inputs.""" + replace_map = {} + for block_variable, input_val in zip(self.get_dependencies(), inputs, strict=True): + coeff = block_variable.output + val = input_val.output if isinstance(input_val, pyadjoint.block_variable.BlockVariable) else input_val + replace_map[coeff] = val + return replace_map - Args: - form: The UFL form to replace coefficients in. - """ - replace_map = self._create_replace_map(form) + def _replace_form_coefficients_recompute( + self, form: ufl.Form | NestedMutableSequence[ufl.Form] | None, replace_map: dict + ) -> ufl.Form | NestedMutableSequence[ufl.Form | None] | None: + """Recursively replace form coefficients for scalar or blocked form structures.""" + if form is None: + return None if isinstance(form, ufl.Form): - return ufl.replace(form, replace_map) + coeffs_in_form = form.coefficients() + sub_map = {k: v for k, v in replace_map.items() if k in coeffs_in_form} + return ufl.replace(form, sub_map) if sub_map else form elif isinstance(form, typing.Sequence): - replaced_forms: typing.MutableSequence[NestedMutableSequence[ufl.Form | None] | ufl.Form | None] = [] - for f in form: - if f is None: - replaced_forms.append(None) - elif isinstance(f, typing.Sequence): - new_form = self._replace_coefficients_in_form(f) - replaced_forms.append(new_form) - else: - replaced_forms.append(ufl.replace(f, replace_map)) - return replaced_forms + return [self._replace_form_coefficients_recompute(f, replace_map) for f in form] else: raise TypeError(f"Cannot replace coefficients in form of type {type(form)}") - def prepare_recompute_component(self, inputs, relevant_outputs): + def prepare_recompute_component( + self, inputs: typing.Sequence[typing.Any], relevant_outputs: typing.Sequence[typing.Any] + ) -> dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function]: """Prepare for recomputing the block with different control inputs.""" - # Replace form coefficients with checkpointed values. - # Loop through the dependencies of the lhs and rhs, check if they are in the respective form - lhs = self._replace_coefficients_in_form(self._lhs) - rhs = self._replace_coefficients_in_form(self._rhs) + replace_map = self._create_recompute_replace_map(inputs) + + # 1. Substitute forms using active candidate inputs instead of static tape checkpoints + lhs = self._replace_form_coefficients_recompute(self._lhs, replace_map) + rhs = self._replace_form_coefficients_recompute(self._rhs, replace_map) preconditioner = ( - self._replace_coefficients_in_form(self._preconditioner) if self._preconditioner is not None else None + self._replace_form_coefficients_recompute(self._preconditioner, replace_map) + if self._preconditioner is not None + else None ) + + # 2. Recompile UFL forms with candidate inputs compiled_lhs = dolfinx.fem.form( lhs, jit_options=self._jit_options, @@ -442,24 +448,27 @@ def prepare_recompute_component(self, inputs, relevant_outputs): else None ) - # Replace the compiled forms with those with new coefficients. + # 3. Hot-swap solver forms self._forward_solver._a = compiled_lhs self._forward_solver._L = compiled_rhs self._forward_solver._P = compiled_preconditioner self._forward_solver.bcs = self._bcs self._forward_solver._u = self._u + + # 4. Solve forward state while halting annotation with pyadjoint.stop_annotating(): solution = self._forward_solver.solve() + return solution def recompute_component( self, inputs: typing.Iterable[Function], - block_variable, + block_variable: pyadjoint.block_variable.BlockVariable, idx: int, prepared: dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function], ) -> dolfinx.fem.Function: - """Recompute the block with the prepared linear problem.""" + """Return the recomputed solution corresponding to the requested output index.""" if isinstance(prepared, dolfinx.fem.Function): assert idx == 0 return prepared From 84bc256a09695d4cde652bebdc7efc6eb81eefd0 Mon Sep 17 00:00:00 2001 From: Joergen Schartum Dokken Date: Thu, 27 Aug 2026 17:14:33 +0000 Subject: [PATCH 20/26] Fix code --- src/dolfinx_adjoint/blocks/solvers.py | 26 +++++++++++++++++--------- 1 file changed, 17 insertions(+), 9 deletions(-) diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index f354587..2e3ebb0 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -457,9 +457,8 @@ def prepare_recompute_component( # 4. Solve forward state while halting annotation with pyadjoint.stop_annotating(): - solution = self._forward_solver.solve() - - return solution + self._forward_solver.solve() + return self._u def recompute_component( self, @@ -468,13 +467,14 @@ def recompute_component( idx: int, prepared: dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function], ) -> dolfinx.fem.Function: - """Return the recomputed solution corresponding to the requested output index.""" + """Recompute and return an isolated copy of the solution state.""" if isinstance(prepared, dolfinx.fem.Function): assert idx == 0 - return prepared + # 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] + return prepared[idx].copy() def _should_compute_boundary_adjoint( self, relevant_dependencies: typing.List[tuple[int, pyadjoint.block_variable.BlockVariable]] @@ -514,17 +514,25 @@ def _compute_residual(self) -> tuple[ufl.Form, dict[ufl.Coefficient, ufl.Coeffic 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, replacement_map = self._compute_residual() + 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." + outputs = self.get_outputs() - r_funcs = [replacement_map[r.saved_output] for r in 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) - return ufl.extract_blocks(dFdu) + + if isinstance(self._u, list): + return ufl.extract_blocks(dFdu) + return dFdu def prepare_evaluate_tlm( self, inputs, tlm_inputs, relevant_outputs From 724cdc11a3fadc5b272d9b4a918174c9f2181b2c Mon Sep 17 00:00:00 2001 From: Joergen Schartum Dokken Date: Thu, 27 Aug 2026 17:14:41 +0000 Subject: [PATCH 21/26] Add more info to demo --- demos/time_distributed_control.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/demos/time_distributed_control.py b/demos/time_distributed_control.py index 5a7300d..1ad778a 100644 --- a/demos/time_distributed_control.py +++ b/demos/time_distributed_control.py @@ -133,8 +133,18 @@ def solve_heat(ctrls): print("\n=== 1. Taylor Test at NON-ZERO Control Point ===") min_val_pert = pyadjoint.taylor_test(rf, m_pert, h) print(f"Convergence rate at perturbed point: {min_val_pert:.4f}") + rf(m_pert) + print("\n=== 2. Second order taylor test at NON-ZERO Control Point ===") + dJdm = sum(drfi._ad_dot(hi) for drfi, hi in zip(rf.derivative(), h, strict=True)) + + H = rf.hessian([hi.control for hi in h]) - print("\n=== 2. Direct Finite Difference Gradient Verification ===") + # 2. Iterate and sum the Hessian dot products piecewise + dHddu = sum(Hi._ad_dot(hi) for Hi, hi in zip(H, h, strict=True)) + min_val = pyadjoint.taylor_test(rf, m_pert, h, dJdm=dJdm, Hm=dHddu) + print(f"Convergence rate at perturbed point with Hessian: {min_val:.4f}") + + print("\n=== 3. Direct Finite Difference Gradient Verification ===") eps = 1e-6 # Compute Adjoint Directional Derivative at m_pert @@ -164,6 +174,9 @@ def solve_heat(ctrls): assert rel_diff < 1e-4, f"Adjoint gradient mismatches finite differences! Relative error: {rel_diff}" +# Reset to ensure that we are at the original control point for the optimization +rf(list(ctrls.values())) + tape = pyadjoint.get_working_tape() tape.visualise_dot("test.dot") From c4672afa5e6b44794428019cdfbbca667ca5f3fe Mon Sep 17 00:00:00 2001 From: Joergen Schartum Dokken Date: Thu, 27 Aug 2026 17:16:32 +0000 Subject: [PATCH 22/26] Final fix --- src/dolfinx_adjoint/blocks/solvers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index 2e3ebb0..5120b24 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -667,7 +667,7 @@ def prepare_evaluate_adj( assert dFdu_adj.empty() is False, "Adjoint of dF/du[v] is empty. Check if the problem is linear." # Solve adjoint problem compiled_dFdu = dolfinx.fem.form( - dFdu_adj, # type: ignore[arg-type] + ufl.extract_blocks(dFdu_adj), # type: ignore[arg-type] jit_options=self._jit_options, form_compiler_options=self._form_compiler_options, entity_maps=self._entity_maps, From 654950e11185a8008a2227ddd00bf29bf8997a1d Mon Sep 17 00:00:00 2001 From: Henrik Finsberg Date: Thu, 27 Aug 2026 17:51:58 +0000 Subject: [PATCH 23/26] Fix mypy issues --- src/dolfinx_adjoint/blocks/solvers.py | 43 ++++++++++++++++++--------- src/dolfinx_adjoint/solvers.py | 6 ++-- 2 files changed, 32 insertions(+), 17 deletions(-) diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index 5120b24..b24174c 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -19,9 +19,15 @@ 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: typing.Sequence[NestedSequence[ufl.Form]], -) -> typing.Sequence[NestedSequence[ufl.Form]]: + *form_structs: NestedSequence[ufl.Form], +) -> NestedSequence[ufl.Form] | tuple[NestedSequence[ufl.Form], ...]: """ Recursively assigns mixed-space `part` indices to UFL Test and Trial functions within nested iterables of forms. @@ -212,7 +218,8 @@ def __init__( # once that is based on a mixed functionspace. if not isinstance(a, ufl.Form): a, L = assign_mixed_parts(a, L) - P, _ = assign_mixed_parts(P, L) if P is not None else None, None + if P is not None: + P, _ = assign_mixed_parts(P, L) self._lhs = a self._rhs = L @@ -237,6 +244,8 @@ def __init__( # 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(): @@ -393,8 +402,14 @@ def _create_recompute_replace_map(self, inputs: typing.Sequence[typing.Any]) -> replace_map[coeff] = val return replace_map + @typing.overload + def _replace_form_coefficients_recompute( + self, form: ufl.Form | NestedSequence[ufl.Form], replace_map: dict + ) -> ufl.Form | NestedMutableSequence[ufl.Form | None]: ... + @typing.overload + def _replace_form_coefficients_recompute(self, form: None, replace_map: dict) -> None: ... def _replace_form_coefficients_recompute( - self, form: ufl.Form | NestedMutableSequence[ufl.Form] | None, replace_map: dict + self, form: ufl.Form | NestedSequence[ufl.Form] | None, replace_map: dict ) -> ufl.Form | NestedMutableSequence[ufl.Form | None] | None: """Recursively replace form coefficients for scalar or blocked form structures.""" if form is None: @@ -426,20 +441,20 @@ def prepare_recompute_component( # 2. Recompile UFL forms with candidate inputs compiled_lhs = dolfinx.fem.form( - lhs, + lhs, # type: ignore[arg-type] jit_options=self._jit_options, form_compiler_options=self._form_compiler_options, entity_maps=self._entity_maps, ) compiled_rhs = dolfinx.fem.form( - rhs, + rhs, # type: ignore[arg-type] jit_options=self._jit_options, form_compiler_options=self._form_compiler_options, entity_maps=self._entity_maps, ) compiled_preconditioner = ( dolfinx.fem.form( - preconditioner, + preconditioner, # type: ignore[arg-type] jit_options=self._jit_options, form_compiler_options=self._form_compiler_options, entity_maps=self._entity_maps, @@ -451,7 +466,7 @@ def prepare_recompute_component( # 3. Hot-swap solver forms self._forward_solver._a = compiled_lhs self._forward_solver._L = compiled_rhs - self._forward_solver._P = compiled_preconditioner + self._forward_solver._preconditioner = compiled_preconditioner self._forward_solver.bcs = self._bcs self._forward_solver._u = self._u @@ -494,7 +509,7 @@ 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[ufl.Coefficient, ufl.Coefficient]]: + 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. """ @@ -631,7 +646,7 @@ def prepare_evaluate_adj( inputs: typing.Sequence[Function], adj_inputs: typing.Sequence[dolfinx.la.Vector], relevant_dependencies: typing.List[tuple[int, pyadjoint.block_variable.BlockVariable]], - ) -> tuple[ufl.Form, dict[ufl.Coefficient, ufl.Coefficient]]: + ) -> tuple[ufl.Form, dict[Function, Function]]: """Prepare the block for evaluating the adjoint.""" # Extract dJ/du[v] from the adjoint inputs. @@ -656,7 +671,7 @@ def prepare_evaluate_adj( else: arrs.append(adj_rhs.array[:local_size]) dolfinx.la.petsc.assign(arrs, dJdu) - dJdu.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD) + dJdu.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD) # type: ignore[arg-type] # Compute (dF/du[v])* for the linear problem. F_form, replacement_map = self._compute_residual() @@ -683,14 +698,14 @@ def evaluate_adj_component( adj_inputs: typing.Iterable[dolfinx.la.Vector], block_variable: pyadjoint.block_variable.BlockVariable, idx: int, - prepared: tuple[ufl.Form, dict[ufl.Coefficient, ufl.Coefficient]], + prepared: tuple[ufl.Form, dict[Function, Function]], ) -> _SpecialVector: """Evaluate the adjoint component, i.e. :math:`\\frac{\\partial F}{\\partial m}`.""" residual, replacement_map = prepared c = block_variable.output c_rep = block_variable.saved_output - if isinstance(c, dolfinx.fem.Function): + if isinstance(c, Function): # Need some clever construction of the TrialFunction to get a part of the mixed space part = idx if isinstance(self._u, list) else None dc = ufl.TrialFunction(c_rep.function_space, part=part) @@ -704,7 +719,7 @@ def evaluate_adj_component( dFdm = -ufl.derivative(sum_res, c_rep, dc) if dFdm.empty(): # Generate a dummy form to safely extract the correct Vector wrapper type - dFdm = dolfinx.fem.form(ufl.ZeroBaseForm((dc,))) # type: ignore[call-overload] + dFdm = dolfinx.fem.form(ufl.ZeroBaseForm((dc,))) # type: ignore[arg-type] dFdm_adj = ufl.adjoint(dFdm) sensitivity = ufl.action(dFdm_adj, self._adjoint_solutions) diff --git a/src/dolfinx_adjoint/solvers.py b/src/dolfinx_adjoint/solvers.py index 68f5aaf..2f9386b 100644 --- a/src/dolfinx_adjoint/solvers.py +++ b/src/dolfinx_adjoint/solvers.py @@ -153,7 +153,7 @@ def solve(self, annotate: bool = True) -> typing.Union[dolfinx.fem.Function, typ self._lhs, # type: ignore[arg-type] self._rhs, # type: ignore[arg-type] bcs=self.bcs, - u=self.u, + u=self.u, # type: ignore[arg-type] P=self._preconditioner, # type: ignore[arg-type] kind=self._kind, # type: ignore[arg-type] petsc_options=self._petsc_options, @@ -292,10 +292,10 @@ def solve(self, annotate: bool = True) -> typing.Union[dolfinx.fem.Function, typ annotate = pyadjoint.annotate_tape({"annotate": annotate}) if annotate: block = NonlinearProblemBlock( - J=self._lhs, + J=self._lhs, # type: ignore[arg-type] F=self._rhs, # type: ignore[arg-type] bcs=self._bcs, - u=self.u, + u=self.u, # type: ignore[arg-type] P=self._preconditioner, # type: ignore[arg-type] kind=self._kind, # type: ignore[arg-type] petsc_options=self._petsc_options, From 1637188ad33a0bf02492f5a006a8185b29da04bd Mon Sep 17 00:00:00 2001 From: Joergen Schartum Dokken Date: Thu, 27 Aug 2026 20:14:34 +0000 Subject: [PATCH 24/26] Add another test. Increase perturbation to avoid floating issues. Also let adjoint_solver have its own internal state and then copy to correct array post solve. --- src/dolfinx_adjoint/blocks/solvers.py | 37 ++++++++++++++---- tests/test_tlm_update.py | 56 ++++++++++++++++++++++++++- 2 files changed, 85 insertions(+), 8 deletions(-) diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index b24174c..9f95ce2 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -330,7 +330,6 @@ def __init__( self._compute_adjoint(sum_form(self._lhs)), # type: ignore[arg-type] self._rhs, # type: ignore[arg-type] bcs=self._bcs, - u=self._adjoint_solutions, # type: ignore[arg-type] P=self._preconditioner, # type: ignore[arg-type] form_compiler_options=self._form_compiler_options, jit_options=self._jit_options, @@ -470,6 +469,12 @@ def prepare_recompute_component( self._forward_solver.bcs = self._bcs self._forward_solver._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 # 4. Solve forward state while halting annotation with pyadjoint.stop_annotating(): self._forward_solver.solve() @@ -567,10 +572,10 @@ def prepare_evaluate_tlm( ) # Build RHS (dFdm) for the monolithic system if isinstance(self._u, list): - test_funcs = get_sorted_arguments(sum_form(self._rhs).arguments(), 0) + test_funcs = get_sorted_arguments(F_form.arguments(), 0) dFdm = sum([ufl.ZeroBaseForm((test,)) for test in test_funcs]) else: - test_funcs = [self._rhs.arguments()[0]] + test_funcs = [F_form.arguments()[0]] dFdm = ufl.ZeroBaseForm((test_funcs[0],)) for block_variable in self.get_dependencies(): @@ -688,8 +693,12 @@ def prepare_evaluate_adj( entity_maps=self._entity_maps, ) self._adjoint_solver._a = compiled_dFdu - self._adjoint_solver._u = self._adjoint_solutions # type: ignore[assignment] self._adjoint_solver.solve() + if isinstance(self._adjoint_solutions, list): + for adj_sol, sol in zip(self._adjoint_solutions, self._adjoint_solver.u): + adj_sol.x.array[:] = sol.x.array[:] + else: + self._adjoint_solutions.x.array[:] = self._adjoint_solver.u.x.array[:] return F_form, replacement_map def evaluate_adj_component( @@ -827,6 +836,7 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ b = self._adjoint_solver._b local_arrays = [bi.array[: bi.index_map.size_local * bi.block_size] for bi in bs] dolfinx.la.petsc.assign(local_arrays, b) + b.ghostUpdate(PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) # Compile SOA LHS dFdu_adj = dolfinx.fem.form( @@ -838,8 +848,13 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ # Solve adjoint problem self._adjoint_solver._a = dFdu_adj - self._adjoint_solver._u = self._second_adjoint_solutions self._adjoint_solver.solve() + if isinstance(self._second_adjoint_solutions, list): + for adj_sol, sol in zip(self._second_adjoint_solutions, self._adjoint_solver.u): + adj_sol.x.array[:] = sol.x.array[:] + else: + self._second_adjoint_solutions.x.array[:] = self._adjoint_solver.u.x.array[:] + return self._compute_residual(), self._adjoint_solutions, self._second_adjoint_solutions def evaluate_hessian_component( @@ -1094,7 +1109,6 @@ def __init__( dFdu_adj, # type: ignore[arg-type] self._rhs, # type: ignore[arg-type] bcs=self._bcs, - u=self._adjoint_solutions, # type: ignore[arg-type] P=self._preconditioner, # type: ignore[arg-type] form_compiler_options=self._form_compiler_options, jit_options=self._jit_options, @@ -1339,8 +1353,12 @@ def prepare_evaluate_adj( entity_maps=self._entity_maps, ) self._adjoint_solver._a = compiled_dFdu - self._adjoint_solver._u = self._adjoint_solutions # type: ignore[assignment] self._adjoint_solver.solve() + if isinstance(self._adjoint_solutions, list): + for adj_sol, sol in zip(self._adjoint_solutions, self._adjoint_solver.u): + adj_sol.x.array[:] = sol.x.array[:] + else: + self._adjoint_solutions.x.array[:] = self._adjoint_solver.u.x.array[:] return F_form def evaluate_adj_component( @@ -1437,6 +1455,11 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ self._adjoint_solver._a = dFdu_adj self._adjoint_solver._u = self._second_adjoint_solutions self._adjoint_solver.solve() + if isinstance(self._second_adjoint_solutions, list): + for adj_sol, sol in zip(self._second_adjoint_solutions, self._adjoint_solver.u): + adj_sol.x.array[:] = sol.x.array[:] + else: + self._second_adjoint_solutions.x.array[:] = self._adjoint_solver.u.x.array[:] return self._compute_residual(), self._adjoint_solutions, self._second_adjoint_solutions def evaluate_hessian_component( diff --git a/tests/test_tlm_update.py b/tests/test_tlm_update.py index 6d73808..9577309 100644 --- a/tests/test_tlm_update.py +++ b/tests/test_tlm_update.py @@ -106,7 +106,7 @@ def test_hessian_is_independent_of_previous_evaluation_points(warm_up_at_another m1.interpolate(lambda x: 1.0 + 0.5 * np.sin(np.pi * x[0])) m2.interpolate(lambda x: 2.0 + 0.5 * np.cos(np.pi * x[1])) h = Function(Z) - h.interpolate(lambda x: 0.3 + 0.2 * np.sin(3 * x[0])) + h.interpolate(lambda x: 10.0 + 3.2 * np.sin(3 * x[0])) if warm_up_at_another_point: Jh(m1) @@ -119,3 +119,57 @@ def test_hessian_is_independent_of_previous_evaluation_points(warm_up_at_another min_rate = pyadjoint.taylor_test(Jh, m2, h, dJdm=dJdm, Hm=Hm) assert np.isclose(min_rate, 3.0, rtol=0.1, atol=0.1), f"Expected convergence rate close to 3.0, got {min_rate}" + + +def test_hessian_mpi_breakdown(mesh_2D): + pyadjoint.get_working_tape().clear_tape() + Jh, Z = _viscous_stokes(mesh_2D) + + m1, m2 = Function(Z), Function(Z) + m1.interpolate(lambda x: 1.0 + 0.5 * np.sin(np.pi * x[0])) + m2.interpolate(lambda x: 2.0 + 0.5 * np.cos(np.pi * x[1])) + h = Function(Z) + h.interpolate(lambda x: 3.0 + 2.0 * np.sin(3 * x[0])) + + # === COLD START === + J_cold = float(Jh(m2)) + dJ_cold = Jh.derivative() + H_cold = Jh.hessian(h) + + # Extract underlying FEniCSx arrays securely + dJ_cold_array = dJ_cold.x.array.copy() if hasattr(dJ_cold, "x") else dJ_cold.array.copy() + H_cold_array = H_cold.x.array.copy() if hasattr(H_cold, "x") else H_cold.array.copy() + + # === WARM START (Pollution check) === + Jh(m1) + Jh.derivative() + Jh.hessian(h) + + J_warm = float(Jh(m2)) + dJ_warm = Jh.derivative() + H_warm = Jh.hessian(h) + + dJ_warm_array = dJ_warm.x.array if hasattr(dJ_warm, "x") else dJ_warm.array + H_warm_array = H_warm.x.array if hasattr(H_warm, "x") else H_warm.array + + # === ASSERTIONS === + + comm = mesh_2D.comm + rank = comm.rank + + # 1. Forward State Check + assert np.isclose(J_cold, J_warm), f"Rank {rank}: Forward evaluation J(m2) differs!" + + # 2. Gradient Check + grad_diff = np.linalg.norm(dJ_cold_array - dJ_warm_array) + assert grad_diff < 1e-10, f"Rank {rank}: Gradient differs after warm up! Diff: {grad_diff}" + + # 3. Hessian Check + hess_diff = np.linalg.norm(H_cold_array - H_warm_array) + assert hess_diff < 1e-10, f"Rank {rank}: Hessian differs after warm up! Diff: {hess_diff}" + + # 4. If arrays match, run Taylor test + dJdm = dJ_warm._ad_dot(h) + Hm = H_warm._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 3.0, got {min_rate}" From 6b371c3991ad739326e1698302a6e02c97ec0af4 Mon Sep 17 00:00:00 2001 From: Henrik Finsberg Date: Fri, 28 Aug 2026 06:40:52 +0000 Subject: [PATCH 25/26] Fix mypy issues --- .pre-commit-config.yaml | 1 + src/dolfinx_adjoint/blocks/solvers.py | 6 ++++-- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 720f67c..5e6d69f 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -27,5 +27,6 @@ repos: rev: v2.3.0 hooks: - id: mypy + language: system files: ^src/|^tests/ args: ["--config-file", "pyproject.toml"] diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index 9f95ce2..35e91e7 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -464,7 +464,7 @@ def prepare_recompute_component( # 3. Hot-swap solver forms self._forward_solver._a = compiled_lhs - self._forward_solver._L = compiled_rhs + self._forward_solver._L = compiled_rhs # type: ignore[assignment] self._forward_solver._preconditioner = compiled_preconditioner self._forward_solver.bcs = self._bcs self._forward_solver._u = self._u @@ -698,6 +698,7 @@ def prepare_evaluate_adj( for adj_sol, sol in zip(self._adjoint_solutions, self._adjoint_solver.u): adj_sol.x.array[:] = sol.x.array[:] else: + assert isinstance(self._adjoint_solutions, dolfinx.fem.Function) self._adjoint_solutions.x.array[:] = self._adjoint_solver.u.x.array[:] return F_form, replacement_map @@ -728,7 +729,7 @@ def evaluate_adj_component( dFdm = -ufl.derivative(sum_res, c_rep, dc) if dFdm.empty(): # Generate a dummy form to safely extract the correct Vector wrapper type - dFdm = dolfinx.fem.form(ufl.ZeroBaseForm((dc,))) # type: ignore[arg-type] + dFdm = dolfinx.fem.form(ufl.ZeroBaseForm((dc,))) # type: ignore[call-overload] dFdm_adj = ufl.adjoint(dFdm) sensitivity = ufl.action(dFdm_adj, self._adjoint_solutions) @@ -1358,6 +1359,7 @@ def prepare_evaluate_adj( for adj_sol, sol in zip(self._adjoint_solutions, self._adjoint_solver.u): adj_sol.x.array[:] = sol.x.array[:] else: + assert isinstance(self._adjoint_solutions, dolfinx.fem.Function) self._adjoint_solutions.x.array[:] = self._adjoint_solver.u.x.array[:] return F_form From 555ed890b456f628f3cadbc3f527208cc1fa2145 Mon Sep 17 00:00:00 2001 From: Henrik Finsberg Date: Fri, 28 Aug 2026 07:20:15 +0000 Subject: [PATCH 26/26] Another mypy issue --- src/dolfinx_adjoint/blocks/solvers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index 35e91e7..81c544a 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -463,7 +463,7 @@ def prepare_recompute_component( ) # 3. Hot-swap solver forms - self._forward_solver._a = compiled_lhs + self._forward_solver._a = compiled_lhs # type: ignore[assignment] self._forward_solver._L = compiled_rhs # type: ignore[assignment] self._forward_solver._preconditioner = compiled_preconditioner self._forward_solver.bcs = self._bcs