diff --git a/pyproject.toml b/pyproject.toml index 2b46375..a27dab6 100644 --- a/pyproject.toml +++ b/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", ] diff --git a/src/dolfinx_adjoint/__init__.py b/src/dolfinx_adjoint/__init__.py index 00e96e7..7a0a688 100644 --- a/src/dolfinx_adjoint/__init__.py +++ b/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", diff --git a/src/dolfinx_adjoint/blocks/assembly.py b/src/dolfinx_adjoint/blocks/assembly.py index b288732..dc50824 100644 --- a/src/dolfinx_adjoint/blocks/assembly.py +++ b/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() diff --git a/src/dolfinx_adjoint/blocks/dirichletbc.py b/src/dolfinx_adjoint/blocks/dirichletbc.py new file mode 100644 index 0000000..a248430 --- /dev/null +++ b/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 diff --git a/src/dolfinx_adjoint/blocks/function_assigner.py b/src/dolfinx_adjoint/blocks/function_assigner.py index 7407de3..fca04b5 100644 --- a/src/dolfinx_adjoint/blocks/function_assigner.py +++ b/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): diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index b5db2e3..f05c0ee 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/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 diff --git a/src/dolfinx_adjoint/types/__init__.py b/src/dolfinx_adjoint/types/__init__.py index c1a57a0..a781b83 100644 --- a/src/dolfinx_adjoint/types/__init__.py +++ b/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 diff --git a/src/dolfinx_adjoint/types/dirichletbc.py b/src/dolfinx_adjoint/types/dirichletbc.py new file mode 100644 index 0000000..f17e6e3 --- /dev/null +++ b/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) diff --git a/tests/test_dirichlet_bc.py b/tests/test_dirichlet_bc.py new file mode 100644 index 0000000..7660efc --- /dev/null +++ b/tests/test_dirichlet_bc.py @@ -0,0 +1,138 @@ +from mpi4py import MPI + +import dolfinx +import numpy as np +import pyadjoint +import ufl +from pyadjoint.overloaded_type import Weakref + +from dolfinx_adjoint import Function, LinearProblem, assemble_scalar, assign, dirichletbc +from dolfinx_adjoint.blocks.dirichletbc import DirichletBCBlock + + +def test_dirichletbc_recording(): + """Test that creating an overloaded dirichletbc correctly registers a block and dependency on the tape.""" + pyadjoint.get_working_tape().clear_tape() + mesh = dolfinx.mesh.create_unit_interval(MPI.COMM_WORLD, 10) + V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) + + c = Function(V, name="boundary_value") + c.interpolate(lambda x: x[0]) + + dofs = dolfinx.fem.locate_dofs_geometrical(V, lambda x: np.isclose(x[0], 0.0)) + bc = dirichletbc(c, dofs) + + tape = pyadjoint.get_working_tape() + blocks = tape.get_blocks() + + # The tape should have 1 block: DirichletBCBlock + assert len(blocks) == 1 + assert isinstance(blocks[0], DirichletBCBlock) + + # The block should have exactly 1 dependency (the function 'c') + assert len(blocks[0].get_dependencies()) == 1 + assert blocks[0].get_dependencies()[0].output is c + + # The returned BC object should now possess the injected block_variable + assert hasattr(bc, "block_variable") + + +def test_dirichletbc_no_annotate(): + """Test that setting annotate=False bypasses tape recording entirely.""" + + pyadjoint.get_working_tape().clear_tape() + mesh = dolfinx.mesh.create_unit_interval(MPI.COMM_WORLD, 10) + V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) + + c = Function(V, name="boundary_value") + c.interpolate(lambda x: x[0]) + + dofs = dolfinx.fem.locate_dofs_geometrical(V, lambda x: np.isclose(x[0], 0.0)) + + # Run with annotation off + bc = dirichletbc(c, dofs, annotate=False) + + tape = pyadjoint.get_working_tape() + + assert len(tape.get_blocks()) == 0 + # FIX: Check the underlying weak reference rather than invoking the property + assert getattr(bc, "_block_variable", Weakref())() is None + + +def test_dirichletbc_recompute(): + """Test the PyAdjoint internal recompute logic specifically for the DirichletBCBlock.""" + pyadjoint.get_working_tape().clear_tape() + mesh = dolfinx.mesh.create_unit_interval(MPI.COMM_WORLD, 10) + V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) + + c = Function(V, name="boundary_value") + c.interpolate(lambda x: np.full_like(x[0], 5.0)) + + dofs = dolfinx.fem.locate_dofs_geometrical(V, lambda x: np.isclose(x[0], 0.0)) + bc = dirichletbc(c, dofs) + + tape = pyadjoint.get_working_tape() + block = tape.get_blocks()[0] + assert isinstance(block, DirichletBCBlock) + + # Simulate an optimizer changing the function value + c.interpolate(lambda x: np.full_like(x[0], 15.0)) + + # Replay the PyAdjoint mechanics manually + prepared = block.prepare_recompute_component([c], None) + new_bc = block.recompute_component([c], bc.block_variable, 0, prepared) + + # Assert that the re-instantiated C++ object captured the updated control value + assert isinstance(new_bc, dolfinx.fem.bcs.DirichletBC) + assert np.isclose(new_bc.g.x.array[0], 15.0) + + +def test_time_dependent_bc_replay(): + pyadjoint.get_working_tape().clear_tape() + + mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 8, 8) + V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) + + dt = 0.1 + num_steps = 3 + + m = Function(V, name="control") + m.interpolate(lambda x: np.sin(x[0] * np.pi)) + + u = ufl.TrialFunction(V) + v = ufl.TestFunction(V) + + uh = Function(V, name="state") + assign(0.0, uh) + + u_prev = Function(V, name="state_prev") + assign(0.0, u_prev) + + F = (u - u_prev) / dt * v * ufl.dx + ufl.inner(ufl.grad(u), ufl.grad(v)) * ufl.dx - m * v * ufl.dx + a, L = ufl.system(F) + + bc_func = Function(V, name="bc_func") + 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) + + # 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) + + J = 0.0 + + for i in range(num_steps): + assign(float(i + 1), bc_func) + problem.solve() + J += assemble_scalar(0.5 * ufl.inner(uh, uh) * ufl.dx) + assign(uh, u_prev) + + J_forward = float(J) + + control = pyadjoint.Control(m) + Jhat = pyadjoint.ReducedFunctional(J, control) + J_replay = Jhat(m) + + assert np.isclose(J_replay, J_forward, atol=1e-10, rtol=1e-10) diff --git a/tests/test_hessian.py b/tests/test_hessian.py new file mode 100644 index 0000000..a8a05c7 --- /dev/null +++ b/tests/test_hessian.py @@ -0,0 +1,188 @@ +from mpi4py import MPI + +import dolfinx +import numpy as np +import pyadjoint +import ufl + +import dolfinx_adjoint + + +def test_constant_hessian(): + pyadjoint.get_working_tape().clear_tape() + + domain = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 10, 10) + V = dolfinx.fem.functionspace(domain, ("Lagrange", 1)) + + # ========================================== + # 1. SETUP THE FORWARD PROBLEM + # PDE: -div(grad(u)) + m * u = f + # Where 'm' is our scalar control parameter + # ========================================== + u = ufl.TrialFunction(V) + v = ufl.TestFunction(V) + f = dolfinx_adjoint.Constant(domain, 1.0) + + # The true parameter value + m_val = 2.0 + m_control = dolfinx_adjoint.Constant(domain, m_val) + + # Weak form + a = ufl.inner(ufl.grad(u), ufl.grad(v)) * ufl.dx + m_control * ufl.inner(u, v) * ufl.dx + L = ufl.inner(f, v) * ufl.dx + + # Zero Dirichlet Boundary Conditions + domain.topology.create_connectivity(domain.topology.dim - 1, domain.topology.dim) + boundary_facets = dolfinx.mesh.exterior_facet_indices(domain.topology) + boundary_dofs = dolfinx.fem.locate_dofs_topological(V, domain.topology.dim - 1, boundary_facets) + uD = dolfinx_adjoint.Function(V) + uD.x.array[:] = 0.0 + bc = dolfinx_adjoint.dirichletbc(uD, boundary_dofs) + + # Solve and tape the PDE + u_sol = dolfinx_adjoint.Function(V, name="State") + problem = dolfinx_adjoint.LinearProblem(a, L, bcs=[bc], u=u_sol) + problem.solve() + + # ========================================== + # 2. SETUP DATA MISFIT + # J_data = 1/(2*var) * \int (u - u_obs)^2 dx + # ========================================== + u_obs = dolfinx_adjoint.Function(V) + u_obs.x.array[:] = 0.0 # Dummy observation + + J_form = 0.5 * ufl.inner(u_sol - u_obs, u_sol - u_obs) * ufl.dx + J_data = dolfinx_adjoint.assemble_scalar(J_form) + + # ========================================== + # 3. EXTRACT THE EXACT TRUE HESSIAN + # ========================================== + control = pyadjoint.Control(m_control) + Jhat = pyadjoint.ReducedFunctional(J_data, control) + + # To get the dense 1x1 Hessian matrix, we compute the Hessian Action + # in the standard basis direction (which for a scalar is simply 1.0) + direction = dolfinx_adjoint.Constant(domain, 1.0) + hessian_action = Jhat.hessian(direction) + + # Cast the action result to a standard Python float + H_misfit = hessian_action.x.array[0] + + # Ensure PyAdjoint actually computed a non-zero curvature! + assert H_misfit > 0.0, "Hessian computation failed or is zero!" + + +def test_constant_hessian_assemble_only(): + """ + Test 1: Does AssembleBlock support Hessians for Constants? + J(m) = 0.5 * (m - 5)**2 * vol + d2J/dm2 = 1.0 * vol + """ + pyadjoint.get_working_tape().clear_tape() + domain = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 3, 3) + + m = dolfinx_adjoint.Constant(domain, 3.0) + + # J = 0.5 * (m - 5)^2 * dx + J_form = 0.5 * (m - 5.0) ** 2 * ufl.dx(domain) + J = dolfinx_adjoint.assemble_scalar(J_form) + + control = pyadjoint.Control(m) + Jhat = pyadjoint.ReducedFunctional(J, control) + + # Direction m_t = 1.0 + direction = dolfinx_adjoint.Constant(domain, 1.0) + hessian_action = Jhat.hessian(direction) + + # Expected Hessian is simply the volume of the domain (1.0 for a unit square) + H_val = hessian_action.x.array[0] + + assert H_val > 0.0, f"AssembleBlock Hessian failed! Value is {H_val}" + assert np.isclose(H_val, 1.0), f"Expected 1.0, got {H_val}" + + +def test_constant_hessian_linear_source(): + """ + Test 2: Does LinearProblemBlock support TLM and SOA for linear parameters? + PDE: -div(grad(u)) = m + J(u) = 0.5 * u**2 * dx + Here, d2F/dudm = 0, so the Hessian is purely the pullback of the objective Hessian. + """ + pyadjoint.get_working_tape().clear_tape() + domain = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 3, 3) + V = dolfinx.fem.functionspace(domain, ("Lagrange", 1)) + + m = dolfinx_adjoint.Constant(domain, 2.0) + + u = ufl.TrialFunction(V) + v = ufl.TestFunction(V) + + # m is just a source term (linear dependence) + a = ufl.inner(ufl.grad(u), ufl.grad(v)) * ufl.dx + L = m * v * ufl.dx + + domain.topology.create_connectivity(domain.topology.dim - 1, domain.topology.dim) + boundary_facets = dolfinx.mesh.exterior_facet_indices(domain.topology) + boundary_dofs = dolfinx.fem.locate_dofs_topological(V, domain.topology.dim - 1, boundary_facets) + u_bc = dolfinx_adjoint.Function(V) + u_bc.x.array[:] = 0.0 + bc = dolfinx_adjoint.dirichletbc(u_bc, boundary_dofs) + + u_sol = dolfinx_adjoint.Function(V) + problem = dolfinx_adjoint.LinearProblem(a, L, bcs=[bc], u=u_sol) + problem.solve() + + J_form = 0.5 * ufl.inner(u_sol, u_sol) * ufl.dx + J = dolfinx_adjoint.assemble_scalar(J_form) + + control = pyadjoint.Control(m) + Jhat = pyadjoint.ReducedFunctional(J, control) + + direction = dolfinx_adjoint.Constant(domain, 1.0) + hessian_action = Jhat.hessian(direction) + H_val = hessian_action.x.array[0] + + assert H_val > 0.0, f"Linear source Hessian failed! Value is {H_val}" + + +def test_constant_hessian_linear_operator(): + """ + Test 3: Does LinearProblemBlock support cross-derivatives (d2F/dudm)? + PDE: -div(grad(u)) + m * u = f + """ + pyadjoint.get_working_tape().clear_tape() + domain = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 3, 3) + V = dolfinx.fem.functionspace(domain, ("Lagrange", 1)) + + m = dolfinx_adjoint.Constant(domain, 2.0) + f = dolfinx_adjoint.Constant(domain, 1.0) + + u = ufl.TrialFunction(V) + v = ufl.TestFunction(V) + + # m multiplies u (non-linear dependence on the parameter) + a = ufl.inner(ufl.grad(u), ufl.grad(v)) * ufl.dx + m * ufl.inner(u, v) * ufl.dx + L = f * v * ufl.dx + + domain.topology.create_connectivity(domain.topology.dim - 1, domain.topology.dim) + boundary_facets = dolfinx.mesh.exterior_facet_indices(domain.topology) + boundary_dofs = dolfinx.fem.locate_dofs_topological(V, domain.topology.dim - 1, boundary_facets) + u_bc = dolfinx_adjoint.Function(V) + u_bc.x.array[:] = 0.0 + bc = dolfinx_adjoint.dirichletbc(u_bc, boundary_dofs) + + u_sol = dolfinx_adjoint.Function(V) + problem = dolfinx_adjoint.LinearProblem(a, L, bcs=[bc], u=u_sol) + problem.solve() + + J_form = 0.5 * ufl.inner(u_sol, u_sol) * ufl.dx + J = dolfinx_adjoint.assemble_scalar(J_form) + + control = pyadjoint.Control(m) + Jhat = pyadjoint.ReducedFunctional(J, control) + + direction = dolfinx_adjoint.Constant(domain, 1.0) + hessian_action = Jhat.hessian(direction) + H_val = hessian_action.x.array[0] + + assert H_val > 0.0, f"Operator Hessian failed! Value is {H_val}"