Adding overload for DirichletBC

pyproject.toml

@@ -13,6 +13,7 @@ dependencies = [
       "fenics-dolfinx>=0.10.0",
       "pyadjoint-ad>=2025.10.0",
       "typing_extensions; python_version < '3.11'",
+      "packaging>=24.2",
 ]
 
 

src/dolfinx_adjoint/__init__.py

@@ -7,7 +7,7 @@
 
 from .assembly import assemble_scalar, error_norm
 from .solvers import LinearProblem, NonlinearProblem
-from .types import Constant, Function
+from .types import Constant, Function, dirichletbc
 from .types.function import assign
 
 meta = metadata("dolfinx_adjoint")
@@ -24,6 +24,7 @@
 __all__ = [
     "Constant",
     "Function",
+    "dirichletbc",
     "LinearProblem",
     "NonlinearProblem",
     "assemble_scalar",

src/dolfinx_adjoint/blocks/assembly.py

@@ -129,6 +129,7 @@ def compute_action_adjoint(
                 form_compiler_options=self._form_compiler_options,
                 entity_maps=self._entity_maps,
             )
+
             if space is None:
                 # If space is not supplied infer it from the form
                 assert len(dform.arguments()) == 1
@@ -143,6 +144,9 @@ def compute_action_adjoint(
             # assemble_compiled_form(compiled_adjoint, self._cached_vectors[id(space)])
             assemble_compiled_form(compiled_adjoint, vector)
             # return a vector scaled by the scalar `adj_input`
+            # Safegaurd against None seeds from PyAdjoint
+            if adj_input is None:
+                adj_input = 1.0
             vector.array[:] *= vector.x.array.dtype.type(adj_input)
             vector.scatter_forward()
 

src/dolfinx_adjoint/blocks/dirichletbc.py

@@ -0,0 +1,42 @@
+import dolfinx
+import numpy as np
+import numpy.typing as npt
+from pyadjoint.block import Block
+
+
+class DirichletBCBlock(Block):
+    """A block representing a DirichletBC in the adjoint framework.
+
+    Args:
+        value: The value of the Dirichlet BC.
+        dofs: An array of degree-of-freedom indices in `V` where the BC should be applied.
+        V: The function space associated with the Dirichlet BC.
+        ad_block_tag: An optional tag to identify this block in the adjoint framework.
+
+    """
+
+    def __init__(
+        self,
+        value: dolfinx.fem.Function | dolfinx.fem.Constant,
+        dofs: npt.NDArray[np.int32],
+        V: dolfinx.fem.FunctionSpace | None = None,
+        ad_block_tag: str | None = None,
+    ):
+        super().__init__(ad_block_tag=ad_block_tag)
+        self._dofs = dofs
+        self._V = V
+        self.add_dependency(value)
+
+    @property
+    def dofs(self):
+        return self._dofs
+
+    @property
+    def V(self):
+        return self._V
+
+    def prepare_recompute_component(self, inputs, relevant_outputs):
+        return inputs[0] if inputs else None
+
+    def recompute_component(self, inputs, block_variable, idx, prepared):
+        return block_variable.saved_output

src/dolfinx_adjoint/blocks/function_assigner.py

@@ -178,15 +178,18 @@ def prepare_recompute_component(self, inputs, relevant_outputs):
     def recompute_component(self, inputs, block_variable, idx, prepared):
         if self.expr is None:
             prepared = inputs[0]
-        output = dolfinx.fem.Function(
-            block_variable.output.function_space, name="f{block_variable.output.name}_AssignBlockRecompute"
-        )
+
+        # We should return the exact object instance to maintain C++ memory bindings
+        # (especially for DirichletBCs), updating it in-place.
+        output = block_variable.saved_output
+
         try:
             if output.function_space == prepared.function_space:
                 output.x.array[:] = prepared.x.array[:]
         except AttributeError:
             # Handling float value
             output.x.array[:] = prepared
+
         return output
 
     def __str__(self):

src/dolfinx_adjoint/blocks/solvers.py

@@ -64,6 +64,7 @@ def __init__(
                 self.add_dependency(c, no_duplicates=True)
             for c in self._rhs.coefficients():  # type: ignore
                 self.add_dependency(c, no_duplicates=True)
+
         except AttributeError:
             raise NotImplementedError("Blocked systems not implemented yet.")
         self._compiled_lhs = dolfinx.fem.form(
@@ -86,6 +87,13 @@ def __init__(
         self._petsc_options = petsc_options if petsc_options is not None else {}
         self._petsc_options_prefix = petsc_options_prefix
         self._bcs = bcs if bcs is not None else []
+
+        # Add dependencies from the boundary conditions
+        if self._bcs is not None:
+            for bc in self._bcs:
+                if hasattr(bc, "block_variable"):
+                    self.add_dependency(bc, no_duplicates=True)
+
         # Solver for recomputing the linear problem
         self._forward_solver = dolfinx.fem.petsc.LinearProblem(
             a=self._lhs,
@@ -162,16 +170,6 @@ def prepare_recompute_component(self, inputs, relevant_outputs):
         else:
             initial_guess = [dolfinx.fem.Function(u.function_space, name=u.name + "_initial_guess") for u in self._u]
 
-        # 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
-
-        #     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)
@@ -206,7 +204,7 @@ def prepare_recompute_component(self, inputs, relevant_outputs):
         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.bcs = self._bcs
         self._forward_solver._u = initial_guess
 
     def recompute_component(
@@ -354,8 +352,36 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar
             entity_maps=self._entity_maps,
         )
         dudm = dolfinx.fem.Function(V, name="du_dm_tlm_linearblock")
-        A_tlm = dolfinx.fem.petsc.assemble_matrix(dFdu, bcs=bcs)
-        A_tlm.assemble()
+        # 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))  # type: ignore[arg-type]
         b_tlm.array[:] = 0.0
         dolfinx.fem.petsc.assemble_vector(b_tlm.petsc_vec, dFdm_compiled)
@@ -368,7 +394,14 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar
                 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]
-        solve_linear_problem(A_tlm, dudm.x, b_tlm, petsc_options=self._tlm_petsc_options)
+
+        # 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()
+
+        # Explicitly free the temporary RHS vector memory
+        b_tlm.petsc_vec.destroy()
+
         return dudm
 
     def prepare_evaluate_adj(
@@ -434,6 +467,7 @@ def evaluate_adj_component(
             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
 
@@ -481,6 +515,7 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_
             b.array[:] *= -1
 
         b.array[:] += hessian_inputs[0].array
+        b.scatter_forward()
 
         # Compile SOA LHS
         dFdu_adj = dolfinx.fem.form(
@@ -584,6 +619,7 @@ def evaluate_hessian_component(
             entity_maps=self._entity_maps,
         )
         hessian_output = _create_vector(compiled_hessian, hessian_form.arguments()[0].ufl_function_space())
+        hessian_output.array[:] = 0.0
         assemble_compiled_form(compiled_hessian, hessian_output)
         hessian_output.array[:] *= -1.0
         return hessian_output
@@ -999,6 +1035,7 @@ def evaluate_adj_component(
             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
 
@@ -1055,7 +1092,6 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_
             entity_maps=self._entity_maps,
         )
 
-        # Solve adjoint problem
         self._adjoint_solver._a = dFdu_adj
         self._adjoint_solver._b = b.petsc_vec
         self._adjoint_solver._u = self._second_adjoint_solutions
@@ -1147,6 +1183,7 @@ def evaluate_hessian_component(
             entity_maps=self._entity_maps,
         )
         hessian_output = _create_vector(compiled_hessian, hessian_form.arguments()[0].ufl_function_space())
+        hessian_output.array[:] = 0.0
         assemble_compiled_form(compiled_hessian, hessian_output)
         hessian_output.array[:] *= -1.0
         return hessian_output

src/dolfinx_adjoint/types/__init__.py

@@ -1,3 +1,4 @@
-__all__ = ["Function", "Constant"]
+__all__ = ["Function", "Constant", "dirichletbc"]
 
+from .dirichletbc import dirichletbc
 from .function import Constant, Function

src/dolfinx_adjoint/types/dirichletbc.py

@@ -0,0 +1,97 @@
+from typing import Any
+
+import dolfinx
+import numpy as np
+import numpy.typing as npt
+import pyadjoint
+from packaging.version import Version
+from pyadjoint.overloaded_type import FloatingType
+
+from ..blocks.dirichletbc import DirichletBCBlock
+from .function import Function
+
+
+class DirichletBC(dolfinx.fem.DirichletBC, FloatingType):
+    """A class overloading :py:class:`dolfinx.fem.DirichletBC` to support
+    it being used as a control variable in the adjoint framework.
+
+    Args:
+        g: The value of the Dirichlet BC.
+        dofs: An array of degree-of-freedom indices in `V` where the BC should be applied.
+        **kwargs: Additional keyword arguments to pass to the
+            :py:func:`pyadjoint.overloaded_type.FloatingType` constructor.
+
+    """
+
+    def __init__(self, g: Function, dofs: npt.NDArray[np.int32], **kwargs):
+        dtype = g.dtype
+
+        cpp_bc: (
+            dolfinx.cpp.fem.DirichletBC_float32
+            | dolfinx.cpp.fem.DirichletBC_float64
+            | dolfinx.cpp.fem.DirichletBC_complex64
+            | dolfinx.cpp.fem.DirichletBC_complex128
+        )
+        if np.issubdtype(dtype, np.float32):
+            assert isinstance(g._cpp_object, dolfinx.cpp.fem.Function_float32)
+            cpp_bc = dolfinx.cpp.fem.DirichletBC_float32(g._cpp_object, dofs)
+        elif np.issubdtype(dtype, np.float64):
+            assert isinstance(g._cpp_object, dolfinx.cpp.fem.Function_float64)
+            cpp_bc = dolfinx.cpp.fem.DirichletBC_float64(g._cpp_object, dofs)
+        elif np.issubdtype(dtype, np.complex64):
+            assert isinstance(g._cpp_object, dolfinx.cpp.fem.Function_complex64)
+            cpp_bc = dolfinx.cpp.fem.DirichletBC_complex64(g._cpp_object, dofs)
+        elif np.issubdtype(dtype, np.complex128):
+            assert isinstance(g._cpp_object, dolfinx.cpp.fem.Function_complex128)
+            cpp_bc = dolfinx.cpp.fem.DirichletBC_complex128(g._cpp_object, dofs)
+        else:
+            raise NotImplementedError(f"Type {dtype} not supported.")
+
+        bc_kwargs: dict[str, Any] = {}
+        # If dolfinx-version is 0.12 we need to pass the following
+        # due to https://github.com/FEniCS/dolfinx/pull/4342/
+        if Version(dolfinx.__version__).minor >= 11:
+            bc_kwargs["V"] = g.function_space
+            bc_kwargs["g"] = g
+
+        super().__init__(cpp_bc, **bc_kwargs)
+
+        annotate = kwargs.pop("annotate", True)
+        annotate = annotate and pyadjoint.annotate_tape()
+
+        FloatingType.__init__(
+            self,
+            g,
+            dtype=dtype,
+            block_class=kwargs.pop("block_class", DirichletBCBlock),
+            _ad_floating_active=False,
+            _ad_args=kwargs.pop("_ad_args", (g, dofs)),
+            annotate=annotate,
+            **kwargs,
+        )
+
+        if annotate:
+            self._ad_annotate_block()
+
+    def _ad_create_checkpoint(self):
+        return self
+
+    def _ad_restore_at_checkpoint(self, checkpoint):
+        return self
+
+
+def dirichletbc(value: Function, dofs: npt.NDArray[np.int32], **kwargs) -> DirichletBC:
+    """Overloaded DirichletBC constructor that creates an adjoint-aware DirichletBC
+
+    Args:
+        value: The value of the Dirichlet BC. Should be a :py:class:`dolfinx_adjoint.Function`.
+            This means you can also pass in a :py:class:`dolfinx_adjoint.Constant` but not
+            a :py:class:`dolfinx.fem.Constant`.
+        dofs: An array of degree-of-freedom indices in `V` where the BC should be applied.
+        **kwargs: Additional keyword arguments to pass to the
+            :py:class:`dolfinx_adjoint.types.dirichletbc.DirichletBC` constructor.
+
+
+    """
+    assert isinstance(value, Function), "value must be a dolfinx_adjoint.Function"
+    return DirichletBC(value, dofs, **kwargs)