Fast automated adjoints for spectral PDE solvers

Abstract

We present an automated procedure for computing model gradients for partial differential equation (PDE) solvers built on sparse spectral methods, a broad class of numerical techniques widely used in the study of fluid dynamics, continuum mechanics, waves, and pattern formation across disciplines. Our approach applies reverse-mode automatic differentiation to symbolic graph representations of PDEs and efficiently constructs adjoint solvers that retain the speed and flexibility of modern Fourier and polynomial spectral methods. We demonstrate the advantages of this approach with a comprehensive implementation in the open-source Dedalus framework. This work uniquely provides a differentiable spectral solver that supports a broad class of equations, geometries, and boundary conditions, and runs efficiently in parallel. It enables users to compute gradients and perform PDE-based optimization for a wide range of time-dependent and nonlinear models with minimal additional code. We demonstrate this system’s capabilities using canonical problems from the literature, showing both strong performance and practical utility for a wide variety of optimization tasks. By integrating automatic adjoints into a flexible solver, our work enables researchers to perform sensitivity analyses in spectral simulations with ease and efficiency.