Automatic Differentiation and Adjoint Looping for Geophysical Flows in Dedalus

The adjoint state method is used to perform PDE-based optimization across geophysical fluid dynamics, including parameter fitting and nonlinear stability analyses for many types of flows. Classical adjoint methods require manually writing a model's adjoint equations, which becomes cumbersome and error-prone for complex models. Reverse-mode automatic differentiation performs the same computations at the compiler level and has become a mainstay of differential programming for machine learning. Here, we present a high-level automatic differentiation system that automatically computes the discrete adjoints of general spectral PDE models in the Dedalus framework. This system leverages the symbolic equation representation in Dedalus, inherently supports MPI parallelization, and does not require users to rewrite their code for a differentiable compiler. We will briefly review the implementation of the system and illustrate its capabilities with various geophysical applications, including transient growth in kinematic dynamos and optimal mixing in stratified shear flows.