DOLFINx-ADJOINT
Automatic differentiation for DOLFINx using the adjoint method. Efficient sensitivity analysis and gradient-based optimization for finite element simulations.
Key Features
Computational graph construction that automatically tracks operations and dependencies.
Efficient gradient computation through backpropagation on the computational graph.
Seamless integration into existing DOLFINx workflows with minimal code changes.
Support for nonlinear problems and coupled multi-physics simulations.
Built-in tools for visualizing and debugging the computational graph.
Drop-in replacements for DOLFINx objects – just add a graph argument.
How It Works
Create a Graph object and pass it to DOLFINx operations.
The graph automatically records every Function, Constant,
and boundary condition.
Run your simulation as usual. Solvers, forms, and assemblies are recorded as edges connecting the nodes in the graph.
Call graph.backprop(J, x) to compute the gradient of any
quantity J with respect to any parameter x.
The computational graph is a directed acyclic graph (DAG) where:
Nodes represent DOLFINx objects (Functions, Constants, boundary conditions, solver states)
Edges represent operations and store the corresponding adjoint equations
During backpropagation, edges compute derivatives using UFL automatic differentiation for forms and implicit differentiation for PDE solves:
Supported Operations
The following DOLFINx objects and functions have drop-in replacements that record operations on the computational graph:
DOLFINx object |
Adjoint capability |
|---|---|
|
Graph-tracked initialization, |
|
Graph-tracked constants with adjoint support |
|
UFL differentiation w.r.t. coefficients and constants |
|
Tracked scalar assembly for objective functions |
|
Boundary condition adjoint via DOF restriction |
|
Implicit adjoint solve for linear systems |
|
Implicit adjoint solve for nonlinear systems |
|
Newton solver with graph-tracked iterations |
All overloaded objects accept an optional graph= keyword argument. When omitted,
they behave identically to their DOLFINx counterparts.
Quick Start
Install
git clone https://github.com/unistuttgart-cmcs/dolfinx-adjoint.git
cd dolfinx-adjoint && pip install -e .
Requires Python >= 3.12 and DOLFINx >= 0.10.0.
For demos and tests, use pip install -e ".[all]".
Usage – compute dJ/df for a Poisson problem:
from dolfinx_adjoint import Graph, fem
graph_ = Graph()
# Set up your DOLFINx problem, passing graph= to track operations
f = fem.Function(W, name="f", graph=graph_)
uh = fem.Function(V, name="u", graph=graph_)
problem = fem.petsc.LinearProblem(a, L, u=uh, bcs=bcs, graph=graph_)
problem.solve(graph=graph_)
J = fem.assemble_scalar(fem.form(J_form, graph=graph_), graph=graph_)
# Compute the gradient
dJdf = graph_.backprop(id(J), id(f))
Demos
Jupyter notebook examples in the demos/ directory:
Stationary Poisson equation on a unit square. Computes gradients of a tracking-type objective w.r.t. the source term f, the diffusion coefficient \(\nu\), and the Dirichlet boundary values.
Time-dependent heat equation solved with backward Euler. Demonstrates adjoint backpropagation through multiple time steps to compute sensitivities w.r.t. the initial condition.
3D cantilever beam under gravity. Computes gradients of the deformation energy w.r.t. the Lame parameters \(\lambda\) and \(\mu\).
Stokes flow around a cylindrical obstacle using P2-P1 Taylor-Hood elements. Computes gradients of viscous dissipation w.r.t. the viscosity and the obstacle boundary condition.
Acknowledgments
This library builds upon the FEniCS Project and uses DOLFINx as its foundation.