PaperlayerView exact revision on arXiv

Source-linked AI summary

Automated derivation of the adjoint of high-level transient finite element programs

Patrick E. Farrell, David A. Ham, Simon F. Funke, Marie E. Rognes

arXiv:1204.5577v2cs.MS

TL;DR

Complex nonlinear and time-dependent finite element adjoints are difficult to derive and implement, especially with parallel communication and checkpointing. The paper introduces dolfin-adjoint, which combines runtime equation-level annotation, libadjoint symbolic manipulation, and FEniCS code generation to automate discrete adjoint and tangent linear models. The framework requires only minor forward-model changes and delivers robust, efficient derivations with optimal checkpointing and inherited parallel support.

  • Problem

    Deriving adjoint models for nonlinear or time-dependent forward models is difficult, with parallel communication reversal and hand-coded checkpointing adding implementation challenges.

  • Method

    dolfin-adjoint records equation-level temporal structure, symbolically derives adjoint and tangent linear systems through libadjoint, and uses FEniCS form compilation to generate executable models.

  • Results

    The framework automatically derives adjoint and tangent linear models for a wide variety of FEniCS finite element models with minimal code changes, achieving reliable and efficient automation.

  • Takeaways & Limitations

    Operating at a higher abstraction level relieves developers of adjoint development while providing optimal checkpointing strategies and inherited parallelism.

  • Takeaways & Limitations

    dolfin-adjoint cannot fully automate differentiation with respect to mesh parameters, although shape calculus can use its automatically computed adjoint solutions.

Abstract

from arXiv · show

In this paper we demonstrate a new technique for deriving discrete adjoint and tangent linear models of finite element models. The technique is significantly more efficient and automatic than standard algorithmic differentiation techniques. The approach relies on a high-level symbolic representation of the forward problem. In contrast to developing a model directly in Fortran or C++, high-level systems allow the developer to express the variational problems to be solved in near-mathematical notation. As such, these systems have a key advantage: since the mathematical structure of the problem is preserved, they are more amenable to automated analysis and manipulation. The framework introduced here is implemented in a freely available software package named dolfin-adjoint, based on the FEniCS Project. Our approach to automated adjoint derivation relies on run-time annotation of the temporal structure of the model, and employs the FEniCS finite element form compiler to automatically generate the low-level code for the derived models. The approach requires only trivial changes to a large class of forward models, including complicated time-dependent nonlinear models. The adjoint model automatically employs optimal checkpointing schemes to mitigate storage requirements for nonlinear models, without any user management or intervention. Furthermore, both the tangent linear and adjoint models naturally work in parallel, without any need to differentiate through calls to MPI or to parse OpenMP directives. The generality, applicability and efficiency of the approach are demonstrated with examples from a wide range of scientific applications.

1. Introduction.

The paper presents a higher-level alternative to hand-coded or instruction-level algorithmic differentiation for deriving adjoint models of complex finite element systems. Its FEniCS/libadjoint framework records equation-level structure, automates derivative-model generation, and requires only minor forward-code changes.

  • Motivation: Nonlinear and time-dependent adjoint models are difficult to implement because operators, reversed parallel communication, and implementation details complicate differentiation.Forward sends become adjoint receives, while forward receives become adjoint sends in parallel computations.
  • Motivation: Selective algorithmic differentiation can restore some mathematical structure, but requires more expertise than naive black-box AD.The approach differentiates small code sections and manually assembles the resulting routines into discrete adjoint equations.
  • Approach: libadjoint treats a model as a sequence of equation solves rather than elementary programming-language instructions, enabling symbolic derivation of adjoint and tangent linear systems.This higher-level abstraction is implemented in an open-source library.
  • Approach: Runtime annotations record the temporal structure of low-level forward models, allowing libadjoint to manipulate the annotated system and assemble derived equations through operator callbacks.The library also manages forward and adjoint variable lifecycles and checkpointing.
  • FEniCS integration: For FEniCS models, the framework uses near-mathematical variational formulations, UFL linearisation, and the same form compiler to generate forward and adjoint code.The derived adjoint remains valid FEniCS input and inherits parallel support from the forward model.

2. The fundamental abstraction of libadjoint.

libadjoint represents discretised PDE models as equation solves and derives tangent linear and adjoint systems from that structured representation. The abstraction supports stationary or transient, linear or nonlinear systems, while callbacks and symbolic manipulation drive model assembly and checkpointing.

  • Mathematical framework: The fundamental abstraction represents a discretised PDE system through the prognostic state u, source term b(u), and discretisation matrix A(u).For transient systems, u, A, and b are block-structured across time levels, with A block-lower-triangular because information propagates forward in time.
  • Mathematical framework: The formulation does not require assembling the entire matrix A or storing the complete state u; forward solvers can process one block-row and block-component at a time.This preserves the abstraction while allowing memory-efficient implementation.
  • Derived systems: The tangent linear model computes the Jacobian du/dm of the solution with respect to parameters, including terms arising from nonlinear dependencies in A and b.Differentiating A with respect to u produces a rank-3 tensor contracted with u, while the right-hand-side contribution uses the Jacobian of b with respect to u.
  • Derived systems: The adjoint associated with a scalar functional J uses an adjoint solution z and reverses the temporal information flow of the forward system.Because the adjoint system is block-upper-triangular, it is typically solved by backward substitution from the final time.
  • Automation: libadjoint derives, assembles, and solves tangent linear and adjoint systems block-row by block-row after the forward code is annotated.Operator callbacks and their derivatives supply the information needed to assemble the derived equations.
  • Automation: The library can automate variable deallocation, consistency checks, and checkpointing, while FEniCS supplies components for automated differential-equation solution.Checkpointing is implemented within the library using information available for reassembling forward equations.

3. The FEniCS system.

FEniCS represents transient finite element models as sequences of variational problems solved within a handwritten temporal loop. This abstraction matches libadjoint and supports automatic integration with dolfin-adjoint.

  • High-level formulation: FEniCS lets users specify variational problems in UFL, whose mathematical-like syntax supports generated low-level finite element code.The form compiler generates optimized C++ code for local element tensors, while DOLFIN performs global assembly and numerical solution.
  • Transient models: Transient DOLFIN solvers typically use a handwritten temporal loop that solves one or more discrete variational problems at each timestep.DOLFIN abstracts spatial discretization but leaves temporal discretization to the model code.
  • Transient models: At the highest level, developers define timestep forms, unknowns, and boundary conditions before calling DOLFIN’s solve function.Linear problems are assembled and solved internally, while nonlinear problems invoke Newton iteration with automatically derived Jacobians.
  • Solver control: DOLFIN also permits explicit assembly and solver control, including matrix reuse and direct LU, Krylov, or matrix-free solvers.These options provide more prescriptive control over linear-system construction and solution.
  • Integration with libadjoint: DOLFIN’s explicit sequence of variational problems exactly matches libadjoint’s equation-solve abstraction, enabling the integration developed here.The integration and related annotation, recording, and callback processes occur automatically without model-developer intervention.

4. Applying libadjoint to DOLFIN.

dolfin-adjoint integrates DOLFIN with libadjoint by automatically annotating equation solves and registering callbacks from runtime variational-form data. Minimal model changes then support adjoint and tangent-linear construction, gradient computation, and nonlinear-solve handling, subject to annotation correctness.

  • 4.1. Annotations: dolfin-adjoint overloads DOLFIN functions that change variable values, annotates events, and requires only minimal source modifications.The Burgers example adds a timestep marker while overloaded solve and assign functions record relevant operations.
  • 4.1. Annotations: Runtime variational forms let dolfin-adjoint inspect solved variables, operators, and dependencies, then register the required information with libadjoint.This automatic annotation is possible because the Python interface represents equations as runtime data.
  • 4.1. Annotations: For nonlinear equations F(u) = 0, the adjoint annotates one equivalent linearized equation instead of replaying every Newton iteration.This avoids rewinding through each nonlinear-solver iteration during the adjoint run.
  • 4.1. Annotations: Overloaded assembly preserves form associations when models pre-assemble tensors, allowing matrix-vector solves to recover and annotate the original equations.The association restores semantic form information that would otherwise be absent from low-level matrices and vectors.
  • 4.1. Annotations: The adjoint’s correctness depends on annotations exactly recording the forward model, and direct underlying vector edits can violate this requirement.libadjoint replay can check annotations by comparing replayed values with those from the original run.
  • 4.2. Callbacks: Automatically generated callbacks capture operator forms and time-dependent boundary conditions or forcing terms before libadjoint composes adjoint or tangent-linear systems.Recorded parameter values are retained in callback closures so forward-solve conditions can be recovered.
  • 4.2. Callbacks: The highest-level interface computes dJ/dm from a functional and parameter using the adjoint solution, while direct adjoint and tangent-linear solutions are also accessible.The tangent-linear solutions advance in time, whereas adjoint solutions proceed backwards.

5. Discussion.

The discussion presents high-level equation-based derivation as a way to automate checkpointing, parallel adjoints, and matrix-free models, while identifying important scope limitations.

  • Checkpointing: libadjoint automatically recomputes unavailable forward solutions from checkpoints and uses revolve to place checkpoints under storage limits.It selects offline or online checkpointing algorithms depending on whether the timestep count is known in advance.
  • Parallelism: Parallel communication remains a developer responsibility when callbacks are supplied directly to libadjoint.Callbacks must implement the reversed information flow through transpose actions, although DOLFIN can handle these patterns automatically.
  • Parallelism: DOLFIN automatically derives adjoint communication patterns for MPI and OpenMP, eliminating parallel-specific code in dolfin-adjoint.The adjoint inherits the forward model’s parallel scalability properties through the same runtime system.
  • Matrix-free models: High-level equation-solve abstraction makes automatic adjoint derivation independent of whether matrix-free solves assemble matrices.Matrix-free operation is treated as an implementation detail that does not change the discrete problem or its adjoint.
  • Limitations: Direct raw-memory modifications to DOLFIN function values can make derived adjoint and tangent linear models inconsistent with the forward model.Replay can detect some inconsistencies by comparing replayed annotations with recorded forward-run values.
  • Limitations: dolfin-adjoint cannot fully automate differentiation with respect to mesh parameters, limiting direct discrete differentiation for typical shape-optimization problems.Shape calculus can use the automatically computed adjoint solutions without differentiating through mesh generation.

6. Examples.

The examples apply dolfin-adjoint to nonlinear and large-scale finite element models, verifying adjoint and tangent linear correctness while demonstrating low overhead, checkpointing, and parallel execution.

  • 6.1. Cahn-Hilliard.: The Cahn-Hilliard example reformulates a fourth-order PDE as two coupled second-order equations, using variational linear finite elements and Crank-Nicolson timestepping.The implementation changes involved fewer than ten lines of code.
  • 6.1. Cahn-Hilliard.: Second-order Taylor remainder convergence verifies the Willmore functional gradient computed by the adjoint and the corresponding tangent linear model.The adjoint test used a fine mesh with more than one million degrees of freedom, while the tangent linear verification ran on 24 processors.
  • 6.1. Cahn-Hilliard.: The Cahn-Hilliard adjoint adds less than 1% annotation overhead and costs approximately 1.22 times the forward model, close to the theoretical ideal ratio of 1.2.The benchmark used a lower-resolution mesh with 40328 degrees of freedom.
  • 6.3. Viscoelasticity.: The viscoelasticity example uses a locking-free mixed finite element discretisation and a two-step TR-BDF2 temporal scheme, with adjoint efficiency approaching the theoretical ideal ratio of 2.The implementation applies a periodically varying pressure boundary condition and directly approximates stresses.

7. Conclusion.

The paper presents dolfin-adjoint as a high-level framework for automatically deriving adjoint and tangent linear models of FEniCS finite element programs. It reports greater automation and efficiency than algorithmic differentiation, with minimal model changes, optimal checkpointing, and parallel execution.

  • 7. Conclusion.: dolfin-adjoint automatically derives adjoint and tangent linear models for a wide variety of FEniCS finite element models with minimal forward-model changes.The framework inherits checkpointing from libadjoint and parallelism from FEniCS.
  • 7. Conclusion.: Operating at a higher abstraction level than algorithmic differentiation achieves greater automation and efficiency while relieving developers of adjoint development.The approach derives models automatically, reliably, and robustly.
Loading 1204.5577v2…