From 73d0a45fa0959e4f45822090484bfc1b4bef6237 Mon Sep 17 00:00:00 2001 From: Joergen Schartum Dokken Date: Thu, 27 Aug 2026 10:58:19 +0000 Subject: [PATCH 1/4] Remove unused variables --- src/dolfinx_adjoint/blocks/interpolation.py | 2 +- src/dolfinx_adjoint/blocks/solvers.py | 36 ++++----------------- src/dolfinx_adjoint/types/function.py | 4 +-- tests/test_nonlinear_problem.py | 1 + 4 files changed, 10 insertions(+), 33 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 2e75261..26f3792 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -117,18 +117,7 @@ def __init__( self.add_dependency(c, no_duplicates=True) except AttributeError: raise NotImplementedError("Blocked systems not implemented yet.") - self._compiled_lhs = dolfinx.fem.form( - self._lhs, - jit_options=jit_options, - form_compiler_options=form_compiler_options, - entity_maps=entity_maps, - ) - self._compiled_rhs = dolfinx.fem.form( - self._rhs, - jit_options=jit_options, - form_compiler_options=form_compiler_options, - entity_maps=entity_maps, - ) + # Cache form parameters for later # NOTE: Should probably be in a struct self._jit_options = jit_options @@ -762,7 +751,6 @@ def __init__( self._adjoint_petsc_options = adjoint_petsc_options self._tlm_petsc_options = tlm_petsc_options super().__init__(ad_block_tag=ad_block_tag) - self._lhs = J self._preconditioner = P # Create overloaded functions @@ -772,7 +760,7 @@ def __init__( self._u = pyadjoint.create_overloaded_object(u) 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 + J = 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] @@ -783,29 +771,17 @@ def __init__( # 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(): + if J is not None: + for c in J.coefficients(): # type: ignore 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(): + for c in self._rhs.coefficients(): # type: ignore 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, - form_compiler_options=form_compiler_options, - entity_maps=entity_maps, - ) - self._compiled_rhs = dolfinx.fem.form( - self._rhs, - jit_options=jit_options, - form_compiler_options=form_compiler_options, - entity_maps=entity_maps, - ) # Cache form parameters for later # NOTE: Should probably be in a struct self._jit_options = jit_options @@ -816,7 +792,7 @@ def __init__( self._bcs = bcs if bcs is not None else [] # Solver for recomputing the linear problem self._forward_solver = dolfinx.fem.petsc.NonlinearProblem( - J=self._lhs, # type: ignore[arg-type] + J=J, # type: ignore[arg-type] F=self._rhs, # type: ignore[arg-type] bcs=self._bcs, u=self._u, # type: ignore[arg-type] diff --git a/src/dolfinx_adjoint/types/function.py b/src/dolfinx_adjoint/types/function.py index f10814d..3f92f4a 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( @@ -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 [name-defined] """Return the index map of the function's vector.""" return self.x.index_map 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 6432bcdf7f35c859b8e14bd951763204c1fdf56f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=B8rgen=20Schartum=20Dokken?= Date: Thu, 27 Aug 2026 18:53:23 +0200 Subject: [PATCH 2/4] Apply suggestions from code review Co-authored-by: Henrik Finsberg --- src/dolfinx_adjoint/blocks/solvers.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index 26f3792..d073eef 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -510,7 +510,7 @@ def evaluate_adj_component( 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 = ufl.ZeroBaseForm((dc,)) dFdm_adj = ufl.adjoint(dFdm) sensitivity = ufl.action(dFdm_adj, self._adjoint_solutions) @@ -888,9 +888,9 @@ def recompute_component( self._forward_solver.solve() if isinstance(self._forward_solver._u, list): - return self._forward_solver._u[idx] + return self._forward_solver._u[idx]._ad_copy() else: - return self._forward_solver._u + return self._forward_solver._u._ad_copy() def _should_compute_boundary_adjoint( self, relevant_dependencies: typing.List[tuple[int, pyadjoint.block_variable.BlockVariable]] From 65baff1cd7d9dfb09328355d4858be78df2918b8 Mon Sep 17 00:00:00 2001 From: Jorgen Schartum Dokken Date: Fri, 28 Aug 2026 09:43:49 +0200 Subject: [PATCH 3/4] Typing --- src/dolfinx_adjoint/blocks/solvers.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index 039837d..3901f96 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -1155,15 +1155,16 @@ def prepare_recompute_component(self, inputs, relevant_outputs): def recompute_component( self, inputs: typing.Iterable[Function], block_variable, idx: int, prepared: None - ) -> typing.Union[dolfinx.fem.Function, typing.Iterable[dolfinx.fem.Function]]: + ) -> Function: """Recompute the block with the prepared linear problem.""" with pyadjoint.tape.stop_annotating(): self._forward_solver.solve() - if isinstance(self._forward_solver._u, list): - return self._forward_solver._u[idx]._ad_copy() + output = self._forward_solver._u[idx] else: - return self._forward_solver._u._ad_copy() + output = self._forward_solver._u + assert isinstance(output, Function) + return output def _should_compute_boundary_adjoint( self, relevant_dependencies: typing.List[tuple[int, pyadjoint.block_variable.BlockVariable]] From 1819a553f9d270904a59bc23a66f5cc9943b524b Mon Sep 17 00:00:00 2001 From: Henrik Finsberg Date: Fri, 28 Aug 2026 07:52:47 +0000 Subject: [PATCH 4/4] Swap J with self._lhs and emove try-except - it makes it harder to debug --- src/dolfinx_adjoint/blocks/solvers.py | 25 +++++++++++-------------- 1 file changed, 11 insertions(+), 14 deletions(-) diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index 3901f96..549c8be 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -1039,20 +1039,17 @@ def __init__( self._rhs = [ufl.replace(Fi, replace_dict) for Fi in F] # 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: - 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) - except AttributeError: - raise NotImplementedError("Blocked systems not implemented yet.") + u_list = self._u if isinstance(self._u, list) else [self._u] + if J is not None: + assert isinstance(J, ufl.Form) + for c in J.coefficients(): + if c not in u_list: # Exclude unknown + self.add_dependency(c, no_duplicates=True) + if self._rhs is not None: + assert isinstance(self._rhs, ufl.Form) + for c in self._rhs.coefficients(): + if c not in u_list: # Exclude unknown + self.add_dependency(c, no_duplicates=True) # Cache form parameters for later # NOTE: Should probably be in a struct