Fast automated adjoints for spectral PDE solvers
This work provides a practical tool for researchers using spectral PDE solvers, enabling efficient gradient computation for inverse problems and optimization.
The authors present a general and automated approach for computing model gradients for PDE solvers built on sparse spectral methods, implemented in the Dedalus framework. This enables users to perform gradient-based optimization and sensitivity analyses for a wide range of time-dependent and nonlinear systems without writing additional code.
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.