Support blocked linear systems.

.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"]

demos/time_distributed_control.py

@@ -38,23 +38,25 @@ 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)
     exterior_dofs = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim - 1, exterior_facets)
 
     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 +83,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 +93,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 +114,75 @@ 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():
+    # 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.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}")
+    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])
+
+    # 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
+    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))
+
+    # 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))
+
+    # Central Finite Difference
+    fd_dir_deriv = (J_plus - J_minus) / (2 * eps)
+
+    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}"
+
+# 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")
 
 opt_ctrls = pyadjoint.minimize(
     rf,
     method="BFGS",
-    # method="Newton-CG",
     options={"maxiter": 100, "disp": True},
 )
 

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
 

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