diff --git a/demos/time_distributed_control.py b/demos/time_distributed_control.py index abef605..1ad778a 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() @@ -42,7 +42,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) mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim) @@ -119,22 +118,64 @@ 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] - for hi in h: - hi.x.array[:] = np.random.random(hi.x.array.shape) + # Insert this diagnostic section into your demo right after J and m are defined: - # 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}" + # 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) - # Prove the Hessian is mathematically exact - rf(list(ctrls.values())) + 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(h) + + 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)) - pyadjoint.taylor_test(rf, list(ctrls.values()), h, dJdm=dJdm, Hm=dHddu) - rf(list(ctrls.values())) + 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") diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index b4d6753..5120b24 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,43 +372,59 @@ 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( - 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, @@ -423,30 +448,33 @@ 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 + self._forward_solver.solve() + return self._u 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.""" + """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]] @@ -466,7 +494,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 +509,36 @@ 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()] + 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() + # 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, outputs, strict=True) + for arg, output in zip(test_functions, r_funcs, strict=True) ] - dFdu = ufl.derivative(F_form, outputs, trial_functions) - return ufl.extract_blocks(dFdu) + + dFdu = ufl.derivative(F_form, r_funcs, trial_functions) + + if isinstance(self._u, list): + return ufl.extract_blocks(dFdu) + 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() + 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 +563,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 +631,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 +641,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,10 +656,18 @@ 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] + 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, @@ -636,7 +675,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 +683,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 +699,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 +715,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 +731,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 +839,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 +875,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)