Automatic differentiation for DOLFINx using the adjoint method. Efficient sensitivity analysis and gradient-based optimization for finite element simulations.
Full documentation is available at unistuttgart-cmcs.github.io/dolfinx-graph-ad.
git clone https://github.com/unistuttgart-cmcs/dolfinx-graph-ad.git
cd dolfinx-graph-ad
pip install -e .Dependency requirements are declared in pyproject.toml including:
- Python >= 3.12
- DOLFINx >= 0.10.0
We refer to the DOLFINx installation instructions for installing DOLFINx.
CI targets DOLFINx 0.11.0. Compatibility with the declared minimum, 0.10.0,
is not currently checked in CI.
For demos and tests, use pip install -e ".[all]".
Compute dJ/df for a Poisson problem:
from dolfinx_graph_ad import Graph
from dolfinx_graph_ad import fem as fem_ad
graph_ = Graph()
# Set up your DOLFINx problem, passing graph= to track operations
f = fem_ad.Function(W, name="f", graph=graph_)
uh = fem_ad.Function(V, name="u", graph=graph_)
problem = fem_ad.petsc.LinearProblem(a, L, u=uh, bcs=bcs, graph=graph_)
problem.solve(graph=graph_)
J = fem_ad.assemble_scalar(fem_ad.form(J_form, graph=graph_), graph=graph_)
# Compute the gradient
(dJdf,) = graph_.backprop(J, f)A graph records a single forward evaluation: record a new graph for every new control value. See Recording operations for details.
backprop accepts recorded objects or exact nodes, and always returns a tuple when
controls are requested. For example, graph_.backprop(J, [f, nu]) selects the latest
versions of f and nu. Pass graph_.get_node(f, version=0) to select a particular
recorded version.
Multiple outputs produce one gradient of their weighted sum per requested control,
following the vector–Jacobian product semantics of PyTorch's autograd.grad.
See Multiple Objectives and Controls
for details on multiple objectives and controls.
The demos/ directory contains comprehensive Jupyter notebook examples demonstrating various PDE problems:
- poisson.ipynb - Poisson equation with gradient computation w.r.t. source term, diffusion coefficient, and boundary conditions
- heat_equation.ipynb - Time-dependent heat equation with sensitivity analysis
- linear_elasticity.ipynb - Linear elasticity problem with parameter sensitivities
- stokes.ipynb - Stokes flow with obstacle, computing gradients w.r.t. viscosity and boundary conditions
This library builds upon the excellent work of the FEniCS Project and uses DOLFINx as its foundation.