PETSc TSAdjoint: a discrete adjoint ODE solver for first-order and second-order sensitivity analysis

TL;DR

PETSc TSAdjoint is a software system that efficiently computes first- and second-order sensitivities of time-dependent nonlinear differential equations using a high-level algorithmic differentiation approach, simplifying user effort and ensuring scalability.

Contribution

The paper introduces PETSc TSAdjoint, a novel software tool that automates the derivation of adjoint models for differential equations, avoiding manual differentiation and integrating with parallel infrastructure.

Findings

Demonstrates high efficiency and scalability in various applications.

Provides minimal user input for derivative computations.

Achieves accurate gradient and Hessian-vector products.

Abstract

We present a new software system PETSc TSAdjoint for first-order and second-order adjoint sensitivity analysis of time-dependent nonlinear differential equations. The derivative calculation in PETSc TSAdjoint is essentially a high-level algorithmic differentiation process. The adjoint models are derived by differentiating the timestepping algorithms and implemented based on the parallel infrastructure in PETSc. Full differentiation of the library code including MPI routines thus is avoided, and users do not need to derive their own adjoint models for their specific applications. PETSc TSAdjoint can compute the first-order derivative, that is, the gradient of a scalar functional, and the Hessian-vector product that carries second-order derivative information, while requiring minimal input (a few callbacks) from the users. Optimal checkpointing schemes are employed by the adjoint model in a manner that is transparent to users. Usability, efficiency, and scalability are demonstrated through examples from a variety of applications.