Fast automated adjoints for spectral PDE solvers

TL;DR

This paper introduces an automated method for computing model gradients in spectral PDE solvers, integrated into Dedalus, enabling efficient optimization and sensitivity analysis for complex systems.

Contribution

It presents a general, automated reverse-mode automatic differentiation approach for spectral PDE solvers, supporting diverse equations and geometries with high efficiency.

Findings

Efficient adjoint solvers retain spectral methods' speed and memory benefits.

The framework supports a broad class of PDEs, geometries, and boundary conditions.

Demonstrated strong performance on canonical inverse problems.

Abstract

We present a general and automated approach for computing model gradients for PDE solvers built on sparse spectral methods, and implement this capability in the widely used open-source Dedalus framework. We apply reverse-mode automatic differentiation to symbolic graph representations of PDEs, efficiently constructing adjoint solvers that retain the speed and memory efficiency of this important class of modern numerical methods. This approach enables users to compute gradients and perform optimization for a wide range of time-dependent and nonlinear systems without writing additional code. The framework supports a broad class of equations, geometries, and boundary conditions, and runs efficiently in parallel using MPI. We demonstrate the flexibility and capabilities of this system using canonical problems from the literature, showing both strong performance and practical utility for a wide variety of inverse problems. By integrating automatic adjoints into a flexible high-level solver, our approach enables researchers to perform gradient-based optimization and sensitivity analyses in spectral simulations with ease and efficiency.