Automatic differentiation
Bramble.jl is designed so that grid functions, operators, and form assemblies can be differentiated end-to-end using Julia's automatic differentiation (AD) ecosystem.
Separation of geometry and coefficients
A grid function's coefficients and the coordinates of the underlying mesh are kept distinct:
- Geometry: Mesh coordinates, element spacings, and quadrature nodes remain in their native floating-point type (typically
Float64). - Coefficients: Grid functions (
VectorElement) store values in whatever scalar type a computation requires, includingForwardDiff.Dualnumbers or tracked values from reverse-mode packages.
Because differentiation typically targets parameter sensitivities or solution fields rather than grid coordinates, the mesh geometry remains undifferentiated. Arithmetic operations, difference operators, and discrete inner products propagate sensitivities via standard Julia multiple dispatch.
Container and allocation rules
Propagating derivatives through discrete operators requires that allocations do not discard sensitivity types or revert to space defaults:
- Output allocation with
similar: Discrete operators allocate destination buffers usingsimilar(uₕ)on their operand. This preserves the operand's coefficient type (such asForwardDiff.Dual) rather than reverting to the grid space's default float type. - Explicit value types with
element: When allocating a grid function for in-place operations during an AD sweep, the element type can be specified directly:uₕ = element(Wₕ, typeof(dual_val)) - Quadrature evaluation in
avgₕ: Cell averaging builds Gauss-Legendre quadrature rules in the coordinate type of the mesh, while evaluating the integrand at those real nodes. Sensitivities present in the integrand propagate into the destination element without requiring dual-valued quadrature tables.
Supported automatic differentiation backends
Bramble's runtime kernels mutate destination buffers using setindex!. AD backends that support array mutation work out of the box:
- ForwardDiff: Dual numbers propagate through coefficients, difference operators, weak forms, and boundary conditions.
- PolyesterForwardDiff: The same forward-mode propagation, chunked across threads.
- ReverseDiff: Tracked arrays and scalars propagate through space operators and form assembly.
- Mooncake & Enzyme: Differentiate operator evaluations and space methods.
- Zygote: Not supported directly because it disallows array mutation (
setindex!).
In practice, using DifferentiationInterface.jl allows swapping between forward and reverse backends through a unified API. For setup instructions and runnable examples across all 5 supported backends, see the Automatic differentiation tutorial.
Differentiating forms and linear systems
Sensitivities can appear in the differential operator, the source term, or the boundary data:
Dual matrix systems
Bilinear forms assembled with assemble(a) or refilled in-place with assemble!(A, a) produce sparse matrices whose entries match the form's coefficient type. When solving A \ F, standard linear algebra solvers propagate sensitivities into the solution vector.
Dirichlet boundary conditions
Boundary conditions constructed with dirichlet_constraints store raw closures without type erasure. When boundary data depends on a differentiated parameter:
bcs = dirichlet_constraints(Ωₕ, :boundary => (x -> a * sin(x[1])))dirichlet_bc! writes sensitivity values directly into the target vector. When eliminating Dirichlet rows in symmetrize!, zero checks use iszero(d), which tests both the nominal value and all partial derivatives. A boundary perturbation that evaluates nominally to zero with non-zero sensitivity is therefore correctly retained rather than skipped.
Sparse Jacobians for nonlinear equations
For nonlinear PDEs, residuals $R(u) = A(u)u - F$ inherit the local stencil pattern of the discrete differential operators. The Jacobian $\partial R / \partial u$ has the same compact sparsity structure as the system matrix.
Using DifferentiationInterface.jl with SparseConnectivityTracer.jl and SparseMatrixColorings.jl, the sparsity pattern and column coloring are determined once. Each Newton step then evaluates the sparse Jacobian with a minimal number of directional derivatives. See the Nonlinear Poisson example for a complete implementation.