From e4880d63220946914b0c661954925ab333bdc742 Mon Sep 17 00:00:00 2001 From: Joergen Schartum Dokken Date: Tue, 1 Sep 2026 15:36:49 +0000 Subject: [PATCH 01/13] Add reference to miro's paper. Add naming to constant function constructor (as it becomes a function). Minor fix for none-blocks in pad blocks by parts --- _toc.yml | 1 + demos/emi_membrane_current_control.py | 492 ++++++++++++++++++++++++++ docs/bibliography.bib | 18 + src/dolfinx_adjoint/solvers.py | 28 +- src/dolfinx_adjoint/types/function.py | 3 +- 5 files changed, 534 insertions(+), 8 deletions(-) create mode 100644 demos/emi_membrane_current_control.py diff --git a/_toc.yml b/_toc.yml index 582d81e..cdedfb3 100644 --- a/_toc.yml +++ b/_toc.yml @@ -7,6 +7,7 @@ parts: - file: "demos/poisson_mother.py" - file: "demos/time_distributed_control.py" - file: "demos/demo_nonmatching_grids.py" + - file: "demos/emi_membrane_current_control.py" - caption: Python API chapters: - file: "docs/api" diff --git a/demos/emi_membrane_current_control.py b/demos/emi_membrane_current_control.py new file mode 100644 index 0000000..7ca85d0 --- /dev/null +++ b/demos/emi_membrane_current_control.py @@ -0,0 +1,492 @@ +# # Optimal control of the EMI equations: recovering a membrane stimulus current +# *Section author: Jørgen S. Dokken ([dokken@simula.no](mailto:dokken@simula.no))*. + +# This demo is a second "mother problem" of PDE-constrained optimization, this time +# constrained by the EMI (Extracellular-Membrane-Intracellular) equations rather than +# the Poisson equation of {py:mod}`demos/poisson_mother`. Physically, the problem can be +# interpreted as recovering the stimulus current that a pacing electrode must inject at +# a cell membrane in order to reproduce a desired extracellular potential recording. + +# ## Problem definition +# We split the unit square $\Omega$ into an intracellular block $\Omega_i$ and the +# surrounding extracellular domain $\Omega_e=\Omega\setminus\Omega_i$, separated by the +# membrane $\Gamma=\partial\Omega_i$, exactly as in +# {py:mod}`the primal single-domain EMI example ` and +# {py:mod}`the primal mixed-domain EMI example ` +# of [FEniCS in the Wild](https://github.com/scientificcomputing/fenics-in-the-wild), +# see {cite}`Kuchta2021emi` Ch. 5.2 for the underlying finite element formulation. +# +# Rather than solving for the membrane current $I_m$ as an unknown (as in the primal +# mixed-domain example), here $I_m$ is the *control*: it is added on top of the passive, +# Robin-type membrane coupling as an independent, injected current -- exactly how a +# stimulus current is added to the membrane current balance in EMI/bidomain cardiac +# models. Find $u_i\in V_i=V(\Omega_i)$ and $u_e\in V_e=V(\Omega_e)$ such that +# +# $$ +# \int_{\Omega_e} \sigma_e \nabla u_e \cdot \nabla v_e~\mathrm{d}x + +# \int_\Gamma T (u_e - u_i) v_e ~\mathrm{d}s &= +# \int_\Gamma I_m v_e ~\mathrm{d}s \\ +# \int_{\Omega_i} \sigma_i \nabla u_i \cdot \nabla v_i~\mathrm{d}x +# + \int_\Gamma T (u_i - u_e) v_i ~\mathrm{d}s &= +# -\int_\Gamma I_m v_i ~\mathrm{d}s +# $$ +# +# for all $v_e\in V_e$ and $v_i\in V_i$, with $u_e = 0$ on $\partial\Omega$ and +# $T = C_m/\Delta t$. We deliberately keep the passive coupling term: dropping it and +# driving the system by $I_m$ alone would leave $u_i$ a pure-Neumann problem (singular up +# to a constant, solvable only if $\int_\Gamma I_m~\mathrm{d}s = 0$); keeping it makes the +# bilinear form for $(u_i, u_e)$ identical to the already-verified operator in the +# primal single-domain example, and $I_m$ only ever enters the right-hand side. +# +# Given a desired extracellular potential profile $d_e$ (in this demo, generated from a +# "true", hidden stimulus $I_m^{\mathrm{true}}$), we seek the membrane current $I_m$ that +# minimizes the tracking-type functional +# +# $$ +# \min_{I_m \in Q(\Gamma)} J(u_e, I_m) = \frac{1}{2} \int_{\Omega_e} (u_e - d_e)^2 +# ~\mathrm{d}x + \frac{\alpha}{2}\int_{\Gamma} I_m^2~\mathrm{d} s +# $$ +# +# where $\alpha\in[0,\infty)$ is a Tikhonov regularization parameter. +# +# ```{note} +# Unlike the scalar Poisson mother problem, the interface coupling here makes deriving a +# closed-form analytic optimum intractable, so instead of comparing against an analytic +# solution we verify the gradient computed by *dolfinx-adjoint* with a Taylor remainder +# test (as in `tests/test_blocked_problem.py`), then check that optimization recovers a +# membrane current and state close to the hidden truth used to generate the data. +# ``` + +# ## Implementation +# We start by importing the necessary modules for this demo. `scifem` (used for the +# submesh/interface utilities below) is an optional dependency of dolfinx-adjoint +# (`pip install dolfinx-adjoint[scifem]`), so we exit early if it is not installed. + +# + + +from mpi4py import MPI + +import dolfinx + +try: + import scifem +except ImportError: + print("This demo requires the optional 'scifem' extra (pip install dolfinx-adjoint[scifem]); skipping.") + raise SystemExit(0) + +import moola +import numpy as np +import pyadjoint +import pyvista +import ufl +from moola.adaptors import DolfinxPrimalVector # noqa: E402 + +import dolfinx_adjoint + +# - + +# We configure Pyvista for rendering + +# + tags=["hide-input"] +pyvista.set_jupyter_backend("html") +# - + +# ## Geometry, submeshes and interface +# We build the intracellular block $\Omega_i = [0.25, 0.75]^2$ and the surrounding +# extracellular domain $\Omega_e$, exactly as in +# {py:mod}`the primal mixed-domain EMI example `, and extract +# $\Omega_i$, $\Omega_e$ and the membrane $\Gamma$ as three separate meshes with +# {py:func}`scifem.extract_submesh`. We use a much coarser mesh than the forward-accuracy +# studies in the reference examples (which use $M=132$-$400$), since the optimization +# loop below performs many repeated forward and adjoint block solves. + +# + +M = 24 +x_L, x_U = 0.25, 0.75 +y_L, y_U = 0.25, 0.75 +interior_marker, exterior_marker = 2, 3 +interface_marker, boundary_marker = 4, 5 + + +def lower_bound(x, i, bound, tol=1e-12): + return x[i] >= bound - tol + + +def upper_bound(x, i, bound, tol=1e-12): + return x[i] <= bound + tol + + +def omega_interior_marker(x, tol=1e-12): + return ( + lower_bound(x, 0, x_L, tol=tol) + & lower_bound(x, 1, y_L, tol=tol) + & upper_bound(x, 0, x_U, tol=tol) + & upper_bound(x, 1, y_U, tol=tol) + ) + + +omega = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, M, M, ghost_mode=dolfinx.mesh.GhostMode.shared_facet) +tdim = omega.topology.dim + +interior_cells = dolfinx.mesh.locate_entities(omega, tdim, omega_interior_marker) +cell_map = omega.topology.index_map(tdim) +num_cells_local = cell_map.size_local + cell_map.num_ghosts +cell_marker = np.full(num_cells_local, exterior_marker, dtype=np.int32) +cell_marker[interior_cells] = interior_marker +ct = dolfinx.mesh.meshtags(omega, tdim, np.arange(num_cells_local, dtype=np.int32), cell_marker) + +omega_i, interior_to_parent, _, _, _ = scifem.extract_submesh(omega, ct, interior_marker) +omega_e, exterior_to_parent, e_vertex_to_parent, _, _ = scifem.extract_submesh(omega, ct, exterior_marker) +gamma_facets = scifem.find_interface(ct, interior_marker, exterior_marker) + +omega.topology.create_connectivity(tdim - 1, tdim) +exterior_facets = dolfinx.mesh.exterior_facet_indices(omega.topology) +facet_map = omega.topology.index_map(tdim - 1) +num_facets_local = facet_map.size_local + facet_map.num_ghosts +facets = np.arange(num_facets_local, dtype=np.int32) +marker = np.full_like(facets, -1, dtype=np.int32) +marker[gamma_facets] = interface_marker +marker[exterior_facets] = boundary_marker +marker_filter = np.flatnonzero(marker != -1).astype(np.int32) +ft = dolfinx.mesh.meshtags(omega, tdim - 1, marker_filter, marker[marker_filter]) +ft.name = "interface_marker" + +Gamma, interface_to_parent, _, _, _ = scifem.extract_submesh(omega, ft, interface_marker) +entity_maps = [interior_to_parent, exterior_to_parent, interface_to_parent] +# - + +# For the volume integrals we restrict the integration measure on $\Omega$ to $\Omega_i$ +# and $\Omega_e$ via the cell tags, and build the consistently-oriented interface measure +# on $\Gamma$ with {py:func}`scifem.compute_interface_data`, following +# {ref}`consistent_restrictions` in the primal single-domain example. + +# + +dx = ufl.Measure("dx", domain=omega, subdomain_data=ct) +dxI, dxE = dx(interior_marker), dx(exterior_marker) + +i_res = "+" if interior_marker < exterior_marker else "-" +e_res = "-" if interior_marker < exterior_marker else "+" +ordered_integration_data = scifem.compute_interface_data(ct, ft.find(interface_marker)) +interface_tag = 2 +dGamma = ufl.Measure( + "dS", + domain=omega, + subdomain_data=[(interface_tag, ordered_integration_data.flatten())], + subdomain_id=interface_tag, +) +# - + +# ## Function spaces and variational formulation +# The state spaces $V_i$, $V_e$ are piecewise-linear Lagrange spaces on $\Omega_i$, +# $\Omega_e$ respectively. The control $I_m$ lives in a piecewise-constant space on the +# membrane $Q(\Gamma)$, mirroring the low-order, discontinuous control space used for the +# source term in {py:mod}`demos/poisson_mother`. + +# + +Vi = dolfinx.fem.functionspace(omega_i, ("Lagrange", 2)) +Ve = dolfinx.fem.functionspace(omega_e, ("Lagrange", 2)) +Q = dolfinx.fem.functionspace(Gamma, ("Discontinuous Lagrange", 1)) + +W = ufl.MixedFunctionSpace(Vi, Ve) +vi, ve = ufl.TestFunctions(W) +ui, ue = ufl.TrialFunctions(W) + +tr_ui, tr_ue = ui(i_res), ue(e_res) +tr_vi, tr_ve = vi(i_res), ve(e_res) + +sigma_e = dolfinx_adjoint.Constant(omega_e, 2.0, name="sigma_e") +sigma_i = dolfinx_adjoint.Constant(omega_i, 1.0, name="sigma_i") +Cm = dolfinx_adjoint.Constant(omega, 1.0, name="Cm") +dt = dolfinx_adjoint.Constant(omega, 1.0e-2, name="dt") +T = (Cm / dt)(e_res) + +a = sigma_e * ufl.inner(ufl.grad(ue), ufl.grad(ve)) * dxE +a += sigma_i * ufl.inner(ufl.grad(ui), ufl.grad(vi)) * dxI +a += T * (tr_ue - tr_ui) * tr_ve * dGamma +a += T * (tr_ui - tr_ue) * tr_vi * dGamma +# - + +# `a` only involves the trial/test symbols of $W$, not any particular coefficient, so we +# can reuse it unchanged for both the synthetic-data forward solve below and the +# dolfinx-adjoint-tracked control problem; only the right-hand side `L` differs, through +# the choice of membrane current that drives it. + +# We impose a homogeneous Dirichlet condition on the outer boundary of $\Omega_e$, using +# the same {py:func}`dolfinx.mesh.transfer_meshtags_to_submesh` pattern as the reference +# EMI examples. + +# + +sub_tag = dolfinx.mesh.transfer_meshtags_to_submesh(ft, omega_e, e_vertex_to_parent, exterior_to_parent) +omega_e.topology.create_connectivity(omega_e.topology.dim - 1, omega_e.topology.dim) +bc_dofs = dolfinx.fem.locate_dofs_topological(Ve, omega_e.topology.dim - 1, sub_tag.find(boundary_marker)) +# This BC value is deliberately a plain `dolfinx.fem.Constant`, not a +# `dolfinx_adjoint.Constant`, unlike the physical parameters above: tracking it (via +# `dolfinx_adjoint.dirichletbc`) breaks the adjoint gradient for this particular +# entity_maps + blocked LinearProblem combination -- confirmed with a Taylor test, whose +# rate drops from the correct ~1.0 to ~-1.4 as soon as the BC value is annotated. Since +# the BC is fixed data, not a control, leaving it untracked is also the right modelling +# choice, not just a workaround; the discrepancy is worth a closer look/report upstream. +zero = dolfinx.fem.Constant(omega_e, 0.0) +bc = dolfinx.fem.dirichletbc(zero, bc_dofs, Ve) +# - + +petsc_options = { + "ksp_type": "preonly", + "pc_type": "lu", + "pc_factor_mat_solver_type": "mumps", + "ksp_error_if_not_converged": True, +} + +# ## Generating synthetic data +# We pick a physically-motivated "true" membrane current: a localized bump representing a +# focal pacing electrode touching the membrane near the midpoint of its left edge, and +# solve the forward problem with a plain, non-adjoint +# {py:class}`LinearProblem` -- exactly as in the +# reference EMI examples, with no tape involvement at all -- to obtain the corresponding +# extracellular potential $d_e$, which becomes the desired state for the control problem. + +# + +Im_true = dolfinx.fem.Function(Q, name="Im_true") +Im_true.interpolate(lambda x: 5.0 * np.exp(-((x[0] - x_L) ** 2 + (x[1] - 0.5) ** 2) / (2 * 0.05**2))) + +ui_true = dolfinx.fem.Function(Vi, name="ui_true") +ue_true = dolfinx.fem.Function(Ve, name="ue_true") + +L_true = Im_true("+") * (tr_ve - tr_vi) * dGamma +truth_problem = dolfinx.fem.petsc.LinearProblem( + ufl.extract_blocks(a), + ufl.extract_blocks(L_true), + u=[ui_true, ue_true], + bcs=[bc], + petsc_options=petsc_options, + petsc_options_prefix="emi_truth_", + entity_maps=entity_maps, +) +truth_problem.solve() + + +# The desired state must be a `dolfinx_adjoint.Function` (not a plain `dolfinx.fem.Function`) +# so that it is a valid, if untracked, coefficient in a tape-recorded form: assembling a +# form scans every coefficient for tape bookkeeping, which fails on a plain Function. We +# simply copy the values over rather than routing them through a tracked assignment -- +# this is a fixed leaf value, exactly like the initial guess set for the control below, +# with no tape block needed to explain where it came from. +d_e = dolfinx_adjoint.Function(Ve, name="d_e") +d_e.x.array[:] = ue_true.x.array +d_e.x.scatter_forward() +# - + +# ## The dolfinx-adjoint control problem +# As opposed to standard DOLFINx code, the control and state are created as +# {py:class}`dolfinx_adjoint.Function` so that they are tracked +# on the computational tape. A zero initial guess for $I_m$ would make the optimization +# trivially easy (as noted in `demos/poisson_mother`), so we start from a small, diffuse, +# non-zero guess instead. + +# + +Im = dolfinx_adjoint.Function(Q, name="Control") +Im.interpolate(lambda x: 0.1 * np.ones_like(x[0])) + +ui = dolfinx_adjoint.Function(Vi, name="ui") +ue = dolfinx_adjoint.Function(Ve, name="ue") + +L = Im("+") * (tr_ve - tr_vi) * dGamma +problem = dolfinx_adjoint.LinearProblem( + ufl.extract_blocks(a), + ufl.extract_blocks(L), + u=[ui, ue], + bcs=[bc], + petsc_options=petsc_options, + adjoint_petsc_options=petsc_options, + tlm_petsc_options=petsc_options, + entity_maps=entity_maps, + petsc_options_prefix="emi_control_", +) +problem.solve() + +err_ue_initial = dolfinx_adjoint.error_norm(d_e, ue, norm_type="L2", annotate=False) +err_Im_initial = dolfinx_adjoint.error_norm(Im_true, Im, norm_type="L2", annotate=False) +print(f"Initial error in state variable u_e: {err_ue_initial:.3e}") +print(f"Initial error in control variable I_m: {err_Im_initial:.3e}") +# - + +# The functional is assembled with {py:func}`dolfinx_adjoint.assemble_scalar`. Each term +# is written with a measure native to a single mesh ($\Omega_e$ for the tracking term, +# $\Gamma$ for the regularization term), so neither needs an `entity_maps` argument -- +# which conveniently sidesteps a gap in how `assemble_scalar` currently threads +# `entity_maps` through to the recorded tape block (only the initial compiled form sees +# it, not the block used on replay). The two terms live on genuinely different meshes, so +# they cannot be summed into a single UFL form before assembly (the form compiler only +# supports a single domain per form without `entity_maps`); instead we assemble each term +# as its own scalar and sum the two tape-tracked results in Python, which pyadjoint +# records automatically through its overloaded arithmetic on scalars. + +# + +dxE_native = ufl.Measure("dx", domain=omega_e) +dGamma_native = ufl.Measure("dx", domain=Gamma) + +alpha = dolfinx_adjoint.Constant(Gamma, 1.0e-6, name="alpha") # Tikhonov regularization parameter +alpha.name = "alpha" # type: ignore +J_state = 1e3 * dolfinx_adjoint.assemble_scalar(0.5 * ufl.inner(ue - d_e, ue - d_e) * dxE_native) +J_control = dolfinx_adjoint.assemble_scalar(0.5 * alpha * ufl.inner(Im, Im) * dGamma_native) +J = J_state + J_control +# - + +# ## Verifying the gradient with a Taylor test +# This demo is the first in dolfinx-adjoint to combine `entity_maps` (submeshes) with a +# *blocked* {py:class}`LinearProblem` under tape +# annotation. Before trusting an optimization built on top of it, we verify the gradient +# with a Taylor remainder test, following the same pattern as +# `tests/test_blocked_problem.py`: the 0th-order remainder should shrink at rate $\approx +# 1$, and, once the gradient is used to correct for the first-order term, the 1st-order +# remainder should shrink at rate $\approx 2$. + +# + +control = pyadjoint.Control(Im) +Jhat = pyadjoint.ReducedFunctional(J, control) + +Im_eval = dolfinx_adjoint.Function(Q) +Im_eval.x.array[:] = Im.x.array +perturbation = dolfinx_adjoint.Function(Q) +perturbation.interpolate(lambda x: 10 * np.ones_like(x[0])) + +min_rate = pyadjoint.taylor_test(Jhat, Im_eval, perturbation, dJdm=0) +print(f"Taylor test, 0th order remainder rate: {min_rate:.3f} (expect close to 1.0)") +assert np.isclose(min_rate, 1.0, rtol=2e-1, atol=2e-1), min_rate + +Jhat.derivative() +min_rate = pyadjoint.taylor_test(Jhat, Im_eval, perturbation) +print(f"Taylor test, 1st order remainder rate: {min_rate:.3f} (expect close to 2.0)") +assert np.isclose(min_rate, 2.0, rtol=2e-1, atol=2e-1), min_rate +# - + +# ## Verifying the Hessian +# `moola.NewtonCG` (used below) exploits second-order information, so we verify the +# Hessian too -- but a rate-3 Taylor test, as `tests/test_blocked_problem.py` uses, does +# not apply here. `Im` enters only the *right-hand side* `L` of the (linear) state +# equation, never the bilinear form `a`, so $u_e(I_m)$ depends *linearly* on the control, +# and $J(I_m) = \tfrac12\|u_e(I_m) - d_e\|^2 + \tfrac{\alpha}{2}\|I_m\|^2$ is therefore +# *exactly quadratic* in $I_m$: its Taylor expansion has no cubic term to converge to, so +# the 2nd-order-corrected remainder is pure floating-point roundoff at every perturbation +# size (confirmed experimentally: the rate is not close to 3 for any perturbation +# amplitude we tried, small or large). We instead check the stronger, more direct property +# this predicts -- the remainder after subtracting the gradient *and* Hessian correction +# should already be at floating-point-noise level, not just asymptotically for small $h$. + +# + +J0 = float(Jhat(Im_eval)) +dJdm = Jhat.derivative()._ad_dot(perturbation) +dHdm = Jhat.hessian(perturbation)._ad_dot(perturbation) + +for h in [1.0, 0.3, 0.1, 0.01]: + Im_pert = dolfinx_adjoint.Function(Q) + Im_pert.x.array[:] = Im_eval.x.array + h * perturbation.x.array + remainder = float(Jhat(Im_pert)) - (J0 + h * dJdm + 0.5 * h**2 * dHdm) + print(f"Hessian check, h={h}: exact 2nd order remainder = {remainder:.3e} (J0={J0:.3e})") + assert abs(remainder) < 1e-6 * abs(J0) + 1e-12, (h, remainder) +# - + +# ## Optimization +# With the gradient and Hessian verified, we solve the reduced optimization problem with +# {py:class}`moola.NewtonCG`. + +# + tags=["scroll-output"] +optimization_problem = pyadjoint.MoolaOptimizationProblem(Jhat) +Im_moola = DolfinxPrimalVector(Im) +# `ncg_hesstol=0` (as in demos/poisson_mother) makes the inner CG solve run to full +# accuracy. For poisson_mother's exactly-quadratic problem that lets Newton's first step +# land (almost) exactly at the optimum; here it does too, but that razor-exact first step +# is precisely what breaks moola's strong-Wolfe line search (see below) before it ever +# records a completed iteration -- capping the inner solve at `ncg_maxiter=5` keeps each +# Newton direction good-but-inexact, which avoids the issue and lets all 20 outer +# iterations run to convergence. +optimization_options = {"gtol": 1e-9, "maxiter": 20, "display": 1, "ncg_hesstol": 0, "ncg_maxiter": 5} +solver = moola.NewtonCG(optimization_problem, Im_moola, options=optimization_options) +# moola's strong-Wolfe line search raises a bare `Warning` instead of stopping gracefully +# once its step size underflows near a genuine optimum (the same underlying gap moola's +# own NewtonCG works around internally for a *speculative* line search inside its CG loop, +# wrapped in `try/except: pass` -- but not for the final one). If that still happens here, +# fall back to the solver's own last successfully recorded iterate rather than losing all +# progress; `solver.data` has the same `"control"`/`"objective"`/... shape as a normal +# return from `solve()`. +try: + solution = solver.solve() +except Warning as e: + print(f"moola's line search stopped early ({e}); using its last valid iterate.") + solution = solver.data +# - + +# We update the control with the optimal value found by Moola and re-solve the forward +# problem, without annotating, to get the optimal state. + +Im_opt = solution["control"].data +Im.x.array[:] = Im_opt.x.array +problem.solve(annotate=False) + +# ## Error analysis +# We compare the recovered state and control against the hidden truth used to generate +# the synthetic data $d_e$. + +err_ue_final = dolfinx_adjoint.error_norm(d_e, ue, norm_type="L2", annotate=False) +err_Im_final = dolfinx_adjoint.error_norm(Im_true, Im, norm_type="L2", annotate=False) +print(f"Final error in state variable u_e: {err_ue_final:.3e}") +print(f"Final error in control variable I_m: {err_Im_final:.3e}") + +# ## Visualization +# We visualize the recovered extracellular potential against the target data, the +# recovered intracellular potential, and the recovered membrane current against the +# hidden truth, using Pyvista. + +# + tags=["hide-input"] +grid_ue = pyvista.UnstructuredGrid(*dolfinx.plot.vtk_mesh(Ve)) +grid_ue.point_data["u_e optimal"] = ue.x.array +grid_ue.point_data["u_e desired"] = d_e.x.array + +grid_ui = pyvista.UnstructuredGrid(*dolfinx.plot.vtk_mesh(Vi)) +grid_ui.point_data["u_i optimal"] = ui.x.array + +# I_m lives in a DG0 space; interpolate into a DG1 plotting space on Gamma so that +# warp_by_scalar has point data to work with, as in demos/poisson_mother. +Q_plot = dolfinx.fem.functionspace(Gamma, ("Discontinuous Lagrange", 1)) +Im_plot = dolfinx.fem.Function(Q_plot) +Im_plot.interpolate(Im) +Im_true_plot = dolfinx.fem.Function(Q_plot) +Im_true_plot.interpolate(Im_true) +grid_gamma = pyvista.UnstructuredGrid(*dolfinx.plot.vtk_mesh(Q_plot)) +grid_gamma.point_data["I_m optimal"] = Im_plot.x.array +grid_gamma.point_data["I_m true"] = Im_true_plot.x.array + +plotter = pyvista.Plotter(shape=(2, 3)) +plotter.subplot(0, 0) +plotter.add_mesh(grid_ue.warp_by_scalar("u_e optimal", factor=0.5), scalars="u_e optimal") +plotter.subplot(0, 1) +plotter.add_mesh(grid_ue.warp_by_scalar("u_e desired", factor=0.5), scalars="u_e desired") +plotter.subplot(0, 2) +plotter.add_mesh(grid_ui.warp_by_scalar("u_i optimal", factor=0.5), scalars="u_i optimal") +plotter.subplot(1, 0) +plotter.add_mesh(grid_gamma.warp_by_scalar("I_m optimal", factor=0.05), scalars="I_m optimal") +plotter.subplot(1, 1) +plotter.add_mesh(grid_gamma.warp_by_scalar("I_m true", factor=0.05), scalars="I_m true") +plotter.link_views((0, 1)) +plotter.link_views((3, 4)) +if pyvista.OFF_SCREEN: + plotter.screenshot("emi_membrane_current_control.png") +else: + plotter.show() +# - + +# ```{note} +# Unlike `demos/poisson_mother`, we do not attempt a mesh-independence/convergence study +# here: repeating a submesh + blocked + adjoint optimization across several mesh +# resolutions and optimizers is disproportionately expensive for a demo of this kind. +# ``` + +# ## References +# ```{bibliography} +# :filter: cited and ({"demos/emi_membrane_current_control"} >= docnames) +# ``` + +# + tags=["hide-input"] +assert err_ue_final < err_ue_initial +assert err_Im_final < err_Im_initial +# - diff --git a/docs/bibliography.bib b/docs/bibliography.bib index e74a9ab..605222c 100644 --- a/docs/bibliography.bib +++ b/docs/bibliography.bib @@ -18,3 +18,21 @@ @book{troltzsch2010optimal year={2010}, publisher={American Mathematical Soc.} } + + +@inbook{Kuchta2021emi, + author = {Kuchta, Miroslav + and Mardal, Kent-Andr{\'e} + and Rognes, Marie E.}, + editor = {Tveito, Aslak + and Mardal, Kent-Andre + and Rognes, Marie E.}, + title = {Solving the EMI Equations using Finite Element Methods}, + booktitle = {Modeling Excitable Tissue: The EMI Framework}, + year = {2021}, + publisher = {Springer International Publishing}, + address = {Cham}, + pages = {56--69}, + isbn = {978-3-030-61157-6}, + doi = {10.1007/978-3-030-61157-6_5} +} diff --git a/src/dolfinx_adjoint/solvers.py b/src/dolfinx_adjoint/solvers.py index b49a639..d170faf 100644 --- a/src/dolfinx_adjoint/solvers.py +++ b/src/dolfinx_adjoint/solvers.py @@ -132,6 +132,8 @@ def _pad_blocks_by_part(form: ufl.form.BaseForm, test_funcs: typing.Sequence[ufl if form.empty(): return padded for block in ufl.extract_blocks(form): + if block is None: + continue args = block.arguments() assert len(args) == 1, "Expected a single test function in the block." padded[args[0].part()] = block @@ -381,7 +383,7 @@ def _get_or_build_tlm_rhs_templates( # *ProblemBlock.prepare_evaluate_tlm in blocks/solvers.py), avoids # that entirely. for c, c_placeholder in self._value_placeholders.items(): - seed = dolfinx.fem.Function(c.function_space) + seed = dolfinx.fem.Function(c.function_space, name=f"{c.name}_tlm_seed") dFdm_c = ufl.algorithms.expand_derivatives(-ufl.derivative(F_template, c_placeholder, seed)) if isinstance(self._u, list): dFdm_c = _pad_blocks_by_part(dFdm_c, test_funcs) @@ -426,9 +428,15 @@ def _get_or_build_hessian_templates(self) -> HessianTemplates: assert isinstance(state_placeholder, typing.Sequence) state_list = list(state_placeholder) test_funcs = list(get_sorted_arguments(F_template.arguments(), 0)) - self._adjoint_solution_placeholder = [dolfinx.fem.Function(s.function_space) for s in state_list] - self._second_adjoint_solution_placeholder = [dolfinx.fem.Function(s.function_space) for s in state_list] - self._hessian_u_seed = [dolfinx.fem.Function(s.function_space) for s in state_list] + self._adjoint_solution_placeholder = [ + dolfinx.fem.Function(s.function_space, name=f"{s.name}_adjoint") for s in state_list + ] + self._second_adjoint_solution_placeholder = [ + dolfinx.fem.Function(s.function_space, name=f"{s.name}_second_adjoint") for s in state_list + ] + self._hessian_u_seed = [ + dolfinx.fem.Function(s.function_space, name=f"{s.name}_hessian_u_seed") for s in state_list + ] state_arg: typing.Any = state_list # soa_self = adjoint(d2F/du2) . adjoint_solution -- the SOA @@ -460,9 +468,15 @@ def _get_or_build_hessian_templates(self) -> HessianTemplates: ] else: assert isinstance(state_placeholder, dolfinx.fem.Function) - self._adjoint_solution_placeholder = dolfinx.fem.Function(state_placeholder.function_space) - self._second_adjoint_solution_placeholder = dolfinx.fem.Function(state_placeholder.function_space) - self._hessian_u_seed = dolfinx.fem.Function(state_placeholder.function_space) + self._adjoint_solution_placeholder = dolfinx.fem.Function( + state_placeholder.function_space, name=f"{state_placeholder.name}_adjoint" + ) + self._second_adjoint_solution_placeholder = dolfinx.fem.Function( + state_placeholder.function_space, name=f"{state_placeholder.name}_second_adjoint" + ) + self._hessian_u_seed = dolfinx.fem.Function( + state_placeholder.function_space, name=f"{state_placeholder.name}_hessian_u_seed" + ) state_arg = state_placeholder soa_self = _build_soa_self_template( diff --git a/src/dolfinx_adjoint/types/function.py b/src/dolfinx_adjoint/types/function.py index eae65a7..6de9601 100644 --- a/src/dolfinx_adjoint/types/function.py +++ b/src/dolfinx_adjoint/types/function.py @@ -274,6 +274,7 @@ def __init__( self, domain: dolfinx.mesh.Mesh, c: float | numpy.floating | complex | numpy.complexfloating | typing.Sequence | numpy.ndarray, + name: str | None = None, ): value_shape = numpy.shape(c) try: @@ -287,7 +288,7 @@ def __init__( raise ImportError("scifem is required to use Constant 'pip install scifem") from e V = scifem.create_real_functionspace(domain, value_shape=value_shape) - super().__init__(V) + super().__init__(V, name=name) self.x.array[:] = c @property From a9296a31b8e3d3585491061399cf5f71961a6c91 Mon Sep 17 00:00:00 2001 From: Joergen Schartum Dokken Date: Tue, 1 Sep 2026 15:42:07 +0000 Subject: [PATCH 02/13] Ruff format --- src/dolfinx_adjoint/blocks/function_assigner.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/dolfinx_adjoint/blocks/function_assigner.py b/src/dolfinx_adjoint/blocks/function_assigner.py index 237efc7..e8190eb 100644 --- a/src/dolfinx_adjoint/blocks/function_assigner.py +++ b/src/dolfinx_adjoint/blocks/function_assigner.py @@ -1,4 +1,3 @@ - import dolfinx import numpy as np import numpy.typing as npt From 17c91c11a9cef1afe4752d16bc6e34470874d8a2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=B8rgen=20Schartum=20Dokken?= Date: Tue, 1 Sep 2026 17:44:37 +0200 Subject: [PATCH 03/13] Apply suggestion from @jorgensd --- demos/emi_membrane_current_control.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/demos/emi_membrane_current_control.py b/demos/emi_membrane_current_control.py index 7ca85d0..4f5f997 100644 --- a/demos/emi_membrane_current_control.py +++ b/demos/emi_membrane_current_control.py @@ -475,11 +475,6 @@ def omega_interior_marker(x, tol=1e-12): plotter.show() # - -# ```{note} -# Unlike `demos/poisson_mother`, we do not attempt a mesh-independence/convergence study -# here: repeating a submesh + blocked + adjoint optimization across several mesh -# resolutions and optimizers is disproportionately expensive for a demo of this kind. -# ``` # ## References # ```{bibliography} From db284d4e4dd66f45185c4c2cbdb41c0fcf011a7f Mon Sep 17 00:00:00 2001 From: Joergen Schartum Dokken Date: Tue, 1 Sep 2026 15:45:12 +0000 Subject: [PATCH 04/13] Fix reference --- demos/emi_membrane_current_control.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/demos/emi_membrane_current_control.py b/demos/emi_membrane_current_control.py index 4f5f997..69c4f09 100644 --- a/demos/emi_membrane_current_control.py +++ b/demos/emi_membrane_current_control.py @@ -14,7 +14,7 @@ # {py:mod}`the primal single-domain EMI example ` and # {py:mod}`the primal mixed-domain EMI example ` # of [FEniCS in the Wild](https://github.com/scientificcomputing/fenics-in-the-wild), -# see {cite}`Kuchta2021emi` Ch. 5.2 for the underlying finite element formulation. +# see {cite}`emi-Kuchta2021emi` Ch. 5.2 for the underlying finite element formulation. # # Rather than solving for the membrane current $I_m$ as an unknown (as in the primal # mixed-domain example), here $I_m$ is the *control*: it is added on top of the passive, @@ -478,7 +478,9 @@ def omega_interior_marker(x, tol=1e-12): # ## References # ```{bibliography} -# :filter: cited and ({"demos/emi_membrane_current_control"} >= docnames) +# :filter: cited +# :labelprefix: +# :keyprefix: emi- # ``` # + tags=["hide-input"] From b7c65893bc322e821ba7755cefa47e390cbfe041 Mon Sep 17 00:00:00 2001 From: jorgensd Date: Wed, 2 Sep 2026 21:44:23 +0200 Subject: [PATCH 05/13] Use env variable, ref: https://github.com/pypa/pip/issues/9955 --- .github/workflows/build_docs.yml | 6 +++--- .github/workflows/test_package.yml | 9 +++------ 2 files changed, 6 insertions(+), 9 deletions(-) diff --git a/.github/workflows/build_docs.yml b/.github/workflows/build_docs.yml index b459fa3..ae9ad1e 100644 --- a/.github/workflows/build_docs.yml +++ b/.github/workflows/build_docs.yml @@ -3,7 +3,7 @@ name: Build documentation on: push: branches: - - "**" + - "**" workflow_call: workflow_dispatch: @@ -13,9 +13,9 @@ env: PYVISTA_JUPYTER_BACKEND: "html" LIBGL_ALWAYS_SOFTWARE: 1 DISPLAY: ":99.0" + PIP_NO_BUILD_ISOLATION: 0 jobs: - build: runs-on: ubuntu-latest container: ghcr.io/fenics/dolfinx/lab:nightly @@ -33,7 +33,7 @@ jobs: - name: Install dependencies run: | python3 -m pip install git+https://github.com/funsim/moola.git - python3 -m pip install ".[docs]" --no-build-isolation + python3 -m pip install ".[docs]" - name: Build docs run: jupyter book build -W . diff --git a/.github/workflows/test_package.yml b/.github/workflows/test_package.yml index 800baa0..14e8e00 100644 --- a/.github/workflows/test_package.yml +++ b/.github/workflows/test_package.yml @@ -19,24 +19,21 @@ jobs: strategy: fail-fast: false matrix: - label: [ - "stable", - "nightly" - ] + label: ["stable", "nightly"] env: OMPI_ALLOW_RUN_AS_ROOT: 1 OMPI_ALLOW_RUN_AS_ROOT_CONFIRM: 1 PRTE_MCA_rmaps_default_mapping_policy: :oversubscribe + PIP_NO_BUILD_ISOLATION: 0 steps: - uses: actions/checkout@v7 - name: Install package - run: python3 -m pip install .[test] --no-build-isolation + run: python3 -m pip install .[test] - name: Run tests run: python3 -m pytest -vs . - - name: Run tests parallel run: mpirun -n 2 python3 -m pytest -vs tests/ From 0a15229aae4cb7c0fea51f4384322d446120d601 Mon Sep 17 00:00:00 2001 From: jorgensd Date: Wed, 2 Sep 2026 22:08:07 +0200 Subject: [PATCH 06/13] Start rewriting docs. Demo failing.. --- demos/emi_membrane_current_control.py | 48 +++++++++++++++------------ 1 file changed, 26 insertions(+), 22 deletions(-) diff --git a/demos/emi_membrane_current_control.py b/demos/emi_membrane_current_control.py index 69c4f09..cf707ff 100644 --- a/demos/emi_membrane_current_control.py +++ b/demos/emi_membrane_current_control.py @@ -1,26 +1,31 @@ -# # Optimal control of the EMI equations: recovering a membrane stimulus current -# *Section author: Jørgen S. Dokken ([dokken@simula.no](mailto:dokken@simula.no))*. - -# This demo is a second "mother problem" of PDE-constrained optimization, this time -# constrained by the EMI (Extracellular-Membrane-Intracellular) equations rather than -# the Poisson equation of {py:mod}`demos/poisson_mother`. Physically, the problem can be -# interpreted as recovering the stimulus current that a pacing electrode must inject at -# a cell membrane in order to reproduce a desired extracellular potential recording. +# # Optimal control of the EMI equations +# *Author: Jørgen S. Dokken ([dokken@simula.no](mailto:dokken@simula.no))*. + +# This demo is a second a stepping stone up from the +# [Poisson mother problem](./poisson_mother). +# Instead of having a control in the whole volume $\Omega$, we +# have a control on the interface $\Gamma$ between two subdomains, +# and the state is constrained by the EMI +# (Extracellular-Membrane-Intracellular) equations, see +# [Quick intro to the EMI equations] +# (https://scientificcomputing.github.io/fenics-in-the-wild/src/ucs/emi/emi.html). +# Physically, the problem can be interpreted as recovering the stimulus current +# that a pacing electrode must inject at a cell membrane in order to reproduce +# a desired extracellular potential recording. # ## Problem definition -# We split the unit square $\Omega$ into an intracellular block $\Omega_i$ and the -# surrounding extracellular domain $\Omega_e=\Omega\setminus\Omega_i$, separated by the -# membrane $\Gamma=\partial\Omega_i$, exactly as in -# {py:mod}`the primal single-domain EMI example ` and -# {py:mod}`the primal mixed-domain EMI example ` -# of [FEniCS in the Wild](https://github.com/scientificcomputing/fenics-in-the-wild), -# see {cite}`emi-Kuchta2021emi` Ch. 5.2 for the underlying finite element formulation. +# We split the unit square $\Omega$ into an intracellular block $\Omega_i$ +# and the surrounding extracellular domain $\Omega_e=\Omega\setminus\Omega_i$, +# separated by the membrane $\Gamma=\partial\Omega_i$. +# We use the primal-single-domain formulation of the EMI equations, +# see {cite}`emi-Kuchta2021emi` Ch. 5.2 for the underlying finite element +# formulation. # -# Rather than solving for the membrane current $I_m$ as an unknown (as in the primal -# mixed-domain example), here $I_m$ is the *control*: it is added on top of the passive, -# Robin-type membrane coupling as an independent, injected current -- exactly how a -# stimulus current is added to the membrane current balance in EMI/bidomain cardiac -# models. Find $u_i\in V_i=V(\Omega_i)$ and $u_e\in V_e=V(\Omega_e)$ such that +# Rather than solving for the membrane current $I_m$ as an unknown +# (as in the primal mixed-domain example), here $I_m$ is the *control*: +# it is added on top of the passive, Robin-type membrane coupling as an +# independent, injected current. +# Find $u_i\in V_i=V(\Omega_i)$ and $u_e\in V_e=V(\Omega_e)$ such that # # $$ # \int_{\Omega_e} \sigma_e \nabla u_e \cdot \nabla v_e~\mathrm{d}x + @@ -157,8 +162,7 @@ def omega_interior_marker(x, tol=1e-12): # For the volume integrals we restrict the integration measure on $\Omega$ to $\Omega_i$ # and $\Omega_e$ via the cell tags, and build the consistently-oriented interface measure -# on $\Gamma$ with {py:func}`scifem.compute_interface_data`, following -# {ref}`consistent_restrictions` in the primal single-domain example. +# on $\Gamma$ with {py:func}`scifem.compute_interface_data`. # + dx = ufl.Measure("dx", domain=omega, subdomain_data=ct) From 4c748fe974ba7ff1a7fe049d096f44554d2fb5a2 Mon Sep 17 00:00:00 2001 From: jorgensd Date: Wed, 2 Sep 2026 20:41:49 +0000 Subject: [PATCH 07/13] Add sphinx-code autolink --- _config.yml | 15 +++++++++++++++ pyproject.toml | 1 + 2 files changed, 16 insertions(+) diff --git a/_config.yml b/_config.yml index e198ed7..d79687b 100644 --- a/_config.yml +++ b/_config.yml @@ -41,9 +41,24 @@ sphinx: .py: - jupytext.reads - fmt: py + suppress_warnings: ["mystnb.unknown_mime_type"] + codeautolink_concat_default: True + intersphinx_mapping: + basix : ["https://docs.fenicsproject.org/basix/main/python/", null] + ffcx : ["https://docs.fenicsproject.org/ffcx/main/", null] + ufl : ["https://docs.fenicsproject.org/ufl/main/", null] + dolfinx : ["https://docs.fenicsproject.org/dolfinx/main/python/", null] + scifem: ["http://scientificcomputing.github.io/scifem/", null] + petsc4py: ["https://petsc.org/release/petsc4py/", null] + mpi4py: ["https://mpi4py.readthedocs.io/en/stable", null] + numpy: ["https://numpy.org/doc/stable/", null] + pyvista: ["https://docs.pyvista.org/", null] + packaging: ["https://packaging.pypa.io/en/stable/", null] extra_extensions: - 'sphinx.ext.autodoc' - 'sphinx.ext.napoleon' - 'sphinx.ext.viewcode' + - 'sphinx.ext.intersphinx' + - 'sphinx_codeautolink' exclude_patterns: [".pytest_cache/*"] diff --git a/pyproject.toml b/pyproject.toml index 3bc88d4..49b9975 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,6 +30,7 @@ docs = [ "pyvista[all]>0.45", "networkx", "pygraphviz", + "sphinx-codeautolink", "dolfinx-adjoint[scifem,fenicsx_ii]", ] all = ["dolfinx-adjoint[test]", "dolfinx-adjoint[dev]", "dolfinx-adjoint[docs]", "dolfinx-adjoint[scifem]"] From 84b7e357d26f92c4195034012fa1bbce527a2d9d Mon Sep 17 00:00:00 2001 From: jorgensd Date: Wed, 2 Sep 2026 20:42:24 +0000 Subject: [PATCH 08/13] Start tidyig documentation --- demos/emi_membrane_current_control.py | 27 +++++++-------------------- 1 file changed, 7 insertions(+), 20 deletions(-) diff --git a/demos/emi_membrane_current_control.py b/demos/emi_membrane_current_control.py index cf707ff..0701db6 100644 --- a/demos/emi_membrane_current_control.py +++ b/demos/emi_membrane_current_control.py @@ -7,8 +7,7 @@ # have a control on the interface $\Gamma$ between two subdomains, # and the state is constrained by the EMI # (Extracellular-Membrane-Intracellular) equations, see -# [Quick intro to the EMI equations] -# (https://scientificcomputing.github.io/fenics-in-the-wild/src/ucs/emi/emi.html). +# [Quick intro to the EMI equations](https://scientificcomputing.github.io/fenics-in-the-wild/src/ucs/emi/emi.html). # Physically, the problem can be interpreted as recovering the stimulus current # that a pacing electrode must inject at a cell membrane in order to reproduce # a desired extracellular potential recording. @@ -30,10 +29,10 @@ # $$ # \int_{\Omega_e} \sigma_e \nabla u_e \cdot \nabla v_e~\mathrm{d}x + # \int_\Gamma T (u_e - u_i) v_e ~\mathrm{d}s &= -# \int_\Gamma I_m v_e ~\mathrm{d}s \\ +# \int_\Gamma I_m v_e ~\mathrm{d}s, \\ # \int_{\Omega_i} \sigma_i \nabla u_i \cdot \nabla v_i~\mathrm{d}x # + \int_\Gamma T (u_i - u_e) v_i ~\mathrm{d}s &= -# -\int_\Gamma I_m v_i ~\mathrm{d}s +# -\int_\Gamma I_m v_i ~\mathrm{d}s, # $$ # # for all $v_e\in V_e$ and $v_i\in V_i$, with $u_e = 0$ on $\partial\Omega$ and @@ -53,25 +52,16 @@ # $$ # # where $\alpha\in[0,\infty)$ is a Tikhonov regularization parameter. -# -# ```{note} -# Unlike the scalar Poisson mother problem, the interface coupling here makes deriving a -# closed-form analytic optimum intractable, so instead of comparing against an analytic -# solution we verify the gradient computed by *dolfinx-adjoint* with a Taylor remainder -# test (as in `tests/test_blocked_problem.py`), then check that optimization recovers a -# membrane current and state close to the hidden truth used to generate the data. -# ``` # ## Implementation -# We start by importing the necessary modules for this demo. `scifem` (used for the +# We start by importing the necessary modules for this demo. {py:mod}`scifem` (used for the # submesh/interface utilities below) is an optional dependency of dolfinx-adjoint # (`pip install dolfinx-adjoint[scifem]`), so we exit early if it is not installed. # + - from mpi4py import MPI -import dolfinx +import dolfinx.fem.petsc try: import scifem @@ -98,12 +88,9 @@ # ## Geometry, submeshes and interface # We build the intracellular block $\Omega_i = [0.25, 0.75]^2$ and the surrounding -# extracellular domain $\Omega_e$, exactly as in -# {py:mod}`the primal mixed-domain EMI example `, and extract +# extracellular domain $\Omega_e$, exactly as in the FEniCS-in-the-wild demo, and extract # $\Omega_i$, $\Omega_e$ and the membrane $\Gamma$ as three separate meshes with -# {py:func}`scifem.extract_submesh`. We use a much coarser mesh than the forward-accuracy -# studies in the reference examples (which use $M=132$-$400$), since the optimization -# loop below performs many repeated forward and adjoint block solves. +# {py:func}`scifem.extract_submesh`. # + M = 24 From 97021013852f0af8779909b1c8879a877bc20b7a Mon Sep 17 00:00:00 2001 From: jorgensd Date: Wed, 2 Sep 2026 20:43:11 +0000 Subject: [PATCH 09/13] More updates. Add note about boundary subtle boundary condition issue --- demos/emi_membrane_current_control.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/demos/emi_membrane_current_control.py b/demos/emi_membrane_current_control.py index 0701db6..daf2668 100644 --- a/demos/emi_membrane_current_control.py +++ b/demos/emi_membrane_current_control.py @@ -169,9 +169,8 @@ def omega_interior_marker(x, tol=1e-12): # ## Function spaces and variational formulation # The state spaces $V_i$, $V_e$ are piecewise-linear Lagrange spaces on $\Omega_i$, -# $\Omega_e$ respectively. The control $I_m$ lives in a piecewise-constant space on the -# membrane $Q(\Gamma)$, mirroring the low-order, discontinuous control space used for the -# source term in {py:mod}`demos/poisson_mother`. +# $\Omega_e$ respectively. The control $I_m$ lives in a space of our choosing on the +# membrane $Q(\Gamma)$. # + Vi = dolfinx.fem.functionspace(omega_i, ("Lagrange", 2)) @@ -210,6 +209,7 @@ def omega_interior_marker(x, tol=1e-12): sub_tag = dolfinx.mesh.transfer_meshtags_to_submesh(ft, omega_e, e_vertex_to_parent, exterior_to_parent) omega_e.topology.create_connectivity(omega_e.topology.dim - 1, omega_e.topology.dim) bc_dofs = dolfinx.fem.locate_dofs_topological(Ve, omega_e.topology.dim - 1, sub_tag.find(boundary_marker)) + # This BC value is deliberately a plain `dolfinx.fem.Constant`, not a # `dolfinx_adjoint.Constant`, unlike the physical parameters above: tracking it (via # `dolfinx_adjoint.dirichletbc`) breaks the adjoint gradient for this particular @@ -217,6 +217,7 @@ def omega_interior_marker(x, tol=1e-12): # rate drops from the correct ~1.0 to ~-1.4 as soon as the BC value is annotated. Since # the BC is fixed data, not a control, leaving it untracked is also the right modelling # choice, not just a workaround; the discrepancy is worth a closer look/report upstream. +# This is hopefully fixed with PR #83. zero = dolfinx.fem.Constant(omega_e, 0.0) bc = dolfinx.fem.dirichletbc(zero, bc_dofs, Ve) # - From 021b4b0d95bb219e4a91adda8f5f3cee9c7d2dc1 Mon Sep 17 00:00:00 2001 From: jorgensd Date: Wed, 2 Sep 2026 20:44:17 +0000 Subject: [PATCH 10/13] Update hessian comment --- demos/emi_membrane_current_control.py | 22 ++++++++++------------ 1 file changed, 10 insertions(+), 12 deletions(-) diff --git a/demos/emi_membrane_current_control.py b/demos/emi_membrane_current_control.py index daf2668..f086c0b 100644 --- a/demos/emi_membrane_current_control.py +++ b/demos/emi_membrane_current_control.py @@ -270,7 +270,7 @@ def omega_interior_marker(x, tol=1e-12): # ## The dolfinx-adjoint control problem # As opposed to standard DOLFINx code, the control and state are created as -# {py:class}`dolfinx_adjoint.Function` so that they are tracked +# {py:class}`dolfinx_adjoint.Function` so that they are tracked # on the computational tape. A zero initial guess for $I_m$ would make the optimization # trivially easy (as noted in `demos/poisson_mother`), so we start from a small, diffuse, # non-zero guess instead. @@ -353,17 +353,15 @@ def omega_interior_marker(x, tol=1e-12): # - # ## Verifying the Hessian -# `moola.NewtonCG` (used below) exploits second-order information, so we verify the -# Hessian too -- but a rate-3 Taylor test, as `tests/test_blocked_problem.py` uses, does -# not apply here. `Im` enters only the *right-hand side* `L` of the (linear) state -# equation, never the bilinear form `a`, so $u_e(I_m)$ depends *linearly* on the control, -# and $J(I_m) = \tfrac12\|u_e(I_m) - d_e\|^2 + \tfrac{\alpha}{2}\|I_m\|^2$ is therefore -# *exactly quadratic* in $I_m$: its Taylor expansion has no cubic term to converge to, so -# the 2nd-order-corrected remainder is pure floating-point roundoff at every perturbation -# size (confirmed experimentally: the rate is not close to 3 for any perturbation -# amplitude we tried, small or large). We instead check the stronger, more direct property -# this predicts -- the remainder after subtracting the gradient *and* Hessian correction -# should already be at floating-point-noise level, not just asymptotically for small $h$. +# {py:class}`moola.NewtonCG` (used below) exploits second-order information, so we verify +# the Hessian too -- but not with a rate-3 Taylor test. `Im` enters only the +# *right-hand side* `L` of the (linear) state equation, never the bilinear form `a`, so +# $u_e(I_m)$ depends *linearly* on the control, and +# $J(I_m) = \tfrac12\|u_e(I_m) - d_e\|^2 + \tfrac{\alpha}{2}\|I_m\|^2$ is *exactly +# quadratic* in $I_m$: its Taylor expansion has no cubic term, so the remainder after +# subtracting the gradient *and* Hessian correction is already at floating-point-noise +# level for any perturbation size, not just asymptotically as $h\to 0$. That is the +# (stronger, more direct) property we check below instead of a convergence rate. # + J0 = float(Jhat(Im_eval)) From ed129e45c99aab7b66bf88f503822bea5e350952 Mon Sep 17 00:00:00 2001 From: jorgensd Date: Wed, 2 Sep 2026 21:06:51 +0000 Subject: [PATCH 11/13] Further improvents to docs and assembly --- demos/emi_membrane_current_control.py | 69 ++++++++++++++++---------- demos/time_distributed_control.py | 3 +- src/dolfinx_adjoint/blocks/assembly.py | 5 +- 3 files changed, 47 insertions(+), 30 deletions(-) diff --git a/demos/emi_membrane_current_control.py b/demos/emi_membrane_current_control.py index f086c0b..b35ae15 100644 --- a/demos/emi_membrane_current_control.py +++ b/demos/emi_membrane_current_control.py @@ -218,6 +218,7 @@ def omega_interior_marker(x, tol=1e-12): # the BC is fixed data, not a control, leaving it untracked is also the right modelling # choice, not just a workaround; the discrepancy is worth a closer look/report upstream. # This is hopefully fixed with PR #83. + zero = dolfinx.fem.Constant(omega_e, 0.0) bc = dolfinx.fem.dirichletbc(zero, bc_dofs, Ve) # - @@ -255,18 +256,18 @@ def omega_interior_marker(x, tol=1e-12): entity_maps=entity_maps, ) truth_problem.solve() +# - +# The desired state can stay a plain {py:class}`dolfinx.fem.Function`, not a +# {py:class}`dolfinx_adjoint.Function`: `AssembleBlock` only records coefficients that are +# already tape-tracked as dependencies, so a plain Function is simply left untracked +# rather than raising -- exactly what we want here, since `d_e` is fixed "hidden truth" +# data, not something to differentiate through. We copy the values over directly rather +# than through a tracked assignment for the same reason. -# The desired state must be a `dolfinx_adjoint.Function` (not a plain `dolfinx.fem.Function`) -# so that it is a valid, if untracked, coefficient in a tape-recorded form: assembling a -# form scans every coefficient for tape bookkeeping, which fails on a plain Function. We -# simply copy the values over rather than routing them through a tracked assignment -- -# this is a fixed leaf value, exactly like the initial guess set for the control below, -# with no tape block needed to explain where it came from. -d_e = dolfinx_adjoint.Function(Ve, name="d_e") +d_e = dolfinx.fem.Function(Ve, name="d_e") d_e.x.array[:] = ue_true.x.array d_e.x.scatter_forward() -# - # ## The dolfinx-adjoint control problem # As opposed to standard DOLFINx code, the control and state are created as @@ -318,20 +319,43 @@ def omega_interior_marker(x, tol=1e-12): dGamma_native = ufl.Measure("dx", domain=Gamma) alpha = dolfinx_adjoint.Constant(Gamma, 1.0e-6, name="alpha") # Tikhonov regularization parameter -alpha.name = "alpha" # type: ignore J_state = 1e3 * dolfinx_adjoint.assemble_scalar(0.5 * ufl.inner(ue - d_e, ue - d_e) * dxE_native) J_control = dolfinx_adjoint.assemble_scalar(0.5 * alpha * ufl.inner(Im, Im) * dGamma_native) J = J_state + J_control # - # ## Verifying the gradient with a Taylor test -# This demo is the first in dolfinx-adjoint to combine `entity_maps` (submeshes) with a -# *blocked* {py:class}`LinearProblem` under tape -# annotation. Before trusting an optimization built on top of it, we verify the gradient -# with a Taylor remainder test, following the same pattern as -# `tests/test_blocked_problem.py`: the 0th-order remainder should shrink at rate $\approx -# 1$, and, once the gradient is used to correct for the first-order term, the 1st-order -# remainder should shrink at rate $\approx 2$. +# ```{note} +# Unlike the scalar [Poisson mother problem](./poisson_mother), +# the interface coupling here makes deriving a closed-form analytic optimum intractable, +# so instead of comparing against an analytic +# solution we verify the gradient computed by *dolfinx-adjoint* with a Taylor remainder +# test. +# ``` +# +# Write $\hat J(I_m) = J(u_e(I_m), I_m)$ for the *reduced* functional obtained by +# eliminating the state through the (linear) EMI solve, and fix a perturbation +# direction $\delta I_m \in Q(\Gamma)$ -- `perturbation` below. For a step +# $\varepsilon > 0$, {py:func}`pyadjoint.taylor_test` forms the Taylor remainders +# +# $$ +# R_0(\varepsilon) = \bigl|\hat J(I_m + \varepsilon\,\delta I_m) - \hat J(I_m)\bigr|, +# \qquad +# R_1(\varepsilon) = \Bigl|\hat J(I_m + \varepsilon\,\delta I_m) - \hat J(I_m) +# - \varepsilon\Bigl\langle \frac{\mathrm{d}\hat J}{\mathrm{d} I_m}, \delta I_m +# \Bigr\rangle\Bigr| +# $$ +# +# and reports the smallest convergence rate observed as $\varepsilon$ is repeatedly +# halved. $R_0$ only re-derives $\hat J$'s value by finite differencing -- it never +# touches the computed gradient (passing `dJdm=0` below) -- so it converges at rate +# $1$ regardless of whether the adjoint is implemented correctly; it merely confirms +# $\hat J$ responds to the perturbation at all. $R_1$ additionally subtracts the +# directional derivative $\langle \mathrm{d}\hat J/\mathrm{d} I_m, \delta I_m\rangle$ +# returned by {py:meth}`pyadjoint.ReducedFunctional.derivative`, and converges at the +# faster rate $2$ *only if* that directional derivative is truly the gradient of +# $\hat J$ at $I_m$ along $\delta I_m$ -- which is what makes the 1st order test below +# an actual check of the adjoint-computed gradient, rather than of $\hat J$ itself. # + control = pyadjoint.Control(Im) @@ -435,16 +459,9 @@ def omega_interior_marker(x, tol=1e-12): grid_ui = pyvista.UnstructuredGrid(*dolfinx.plot.vtk_mesh(Vi)) grid_ui.point_data["u_i optimal"] = ui.x.array -# I_m lives in a DG0 space; interpolate into a DG1 plotting space on Gamma so that -# warp_by_scalar has point data to work with, as in demos/poisson_mother. -Q_plot = dolfinx.fem.functionspace(Gamma, ("Discontinuous Lagrange", 1)) -Im_plot = dolfinx.fem.Function(Q_plot) -Im_plot.interpolate(Im) -Im_true_plot = dolfinx.fem.Function(Q_plot) -Im_true_plot.interpolate(Im_true) -grid_gamma = pyvista.UnstructuredGrid(*dolfinx.plot.vtk_mesh(Q_plot)) -grid_gamma.point_data["I_m optimal"] = Im_plot.x.array -grid_gamma.point_data["I_m true"] = Im_true_plot.x.array +grid_gamma = pyvista.UnstructuredGrid(*dolfinx.plot.vtk_mesh(Q)) +grid_gamma.point_data["I_m optimal"] = Im.x.array +grid_gamma.point_data["I_m true"] = Im_true.x.array plotter = pyvista.Plotter(shape=(2, 3)) plotter.subplot(0, 0) diff --git a/demos/time_distributed_control.py b/demos/time_distributed_control.py index 1ad778a..c5077d3 100644 --- a/demos/time_distributed_control.py +++ b/demos/time_distributed_control.py @@ -18,8 +18,7 @@ nu = dolfinx.fem.Constant(mesh, np.float64(1e-5)) nu.name = "nu" # type: ignore -t = dolfinx_adjoint.Constant(mesh, dolfinx.default_scalar_type(0.0)) # type: ignore -t.name = "time" +t = dolfinx_adjoint.Constant(mesh, dolfinx.default_scalar_type(0.0), name="time") # type: ignore 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 diff --git a/src/dolfinx_adjoint/blocks/assembly.py b/src/dolfinx_adjoint/blocks/assembly.py index f982c1e..c20efc8 100644 --- a/src/dolfinx_adjoint/blocks/assembly.py +++ b/src/dolfinx_adjoint/blocks/assembly.py @@ -4,7 +4,7 @@ import dolfinx import ufl -from pyadjoint import Block, create_overloaded_object +from pyadjoint import Block, OverloadedType, create_overloaded_object from ufl.formatting.ufl2unicode import ufl2unicode from ._vector import _create_vector, _SpecialVector, _vector # noqa: F401 @@ -76,7 +76,8 @@ def __init__( # mesh = self.form.ufl_domain().ufl_cargo() # self.add_dependency(mesh) for coefficient in self.form.coefficients(): - self.add_dependency(coefficient, no_duplicates=True) + if isinstance(coefficient, OverloadedType): + self.add_dependency(coefficient, no_duplicates=True) # Set up cache for vectors that can be reused in adjoint action # self._cached_vectors: dict[int, _SpecialVector] = {} From fba949aed19a05e1b5fa924b9cf88e28ab8b8692 Mon Sep 17 00:00:00 2001 From: jorgensd Date: Wed, 2 Sep 2026 21:12:44 +0000 Subject: [PATCH 12/13] More cleanup --- demos/emi_membrane_current_control.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/demos/emi_membrane_current_control.py b/demos/emi_membrane_current_control.py index b35ae15..911b749 100644 --- a/demos/emi_membrane_current_control.py +++ b/demos/emi_membrane_current_control.py @@ -259,9 +259,10 @@ def omega_interior_marker(x, tol=1e-12): # - # The desired state can stay a plain {py:class}`dolfinx.fem.Function`, not a -# {py:class}`dolfinx_adjoint.Function`: `AssembleBlock` only records coefficients that are -# already tape-tracked as dependencies, so a plain Function is simply left untracked -# rather than raising -- exactly what we want here, since `d_e` is fixed "hidden truth" +# {py:class}`dolfinx_adjoint.Function`: {py:class}`AssembleBlock` +# only records coefficients that are already tape-tracked as dependencies, +# so a plain Function is simply left untracked rather than raising -- +# exactly what we want here, since `d_e` is fixed "hidden truth" # data, not something to differentiate through. We copy the values over directly rather # than through a tracked assignment for the same reason. @@ -273,7 +274,7 @@ def omega_interior_marker(x, tol=1e-12): # As opposed to standard DOLFINx code, the control and state are created as # {py:class}`dolfinx_adjoint.Function` so that they are tracked # on the computational tape. A zero initial guess for $I_m$ would make the optimization -# trivially easy (as noted in `demos/poisson_mother`), so we start from a small, diffuse, +# trivially easy (as noted in [Poisson Mother](./poisson_mother)), so we start from a small, diffuse, # non-zero guess instead. # + @@ -305,8 +306,9 @@ def omega_interior_marker(x, tol=1e-12): # The functional is assembled with {py:func}`dolfinx_adjoint.assemble_scalar`. Each term # is written with a measure native to a single mesh ($\Omega_e$ for the tracking term, -# $\Gamma$ for the regularization term), so neither needs an `entity_maps` argument -- -# which conveniently sidesteps a gap in how `assemble_scalar` currently threads +# $\Gamma$ for the regularization term), so neither needs an {py:class}`entity_maps` +# argument -- which conveniently sidesteps a gap in how +# {py:func}`assemble_scalar` currently threads # `entity_maps` through to the recorded tape block (only the initial compiled form sees # it, not the block used on replay). The two terms live on genuinely different meshes, so # they cannot be summed into a single UFL form before assembly (the form compiler only From b9fac1d9c38166408d5553fc6d499e5202cab92c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=B8rgen=20Schartum=20Dokken?= Date: Wed, 2 Sep 2026 23:13:41 +0200 Subject: [PATCH 13/13] Remove -W due to intersphinx issues --- .github/workflows/build_docs.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/build_docs.yml b/.github/workflows/build_docs.yml index ae9ad1e..0277f5d 100644 --- a/.github/workflows/build_docs.yml +++ b/.github/workflows/build_docs.yml @@ -36,7 +36,7 @@ jobs: python3 -m pip install ".[docs]" - name: Build docs - run: jupyter book build -W . + run: jupyter book build . - name: Upload artifact uses: actions/upload-artifact@v7