Differentiable programming for ocean model design and calibration

Differentiable programming gives numerical models access to the exact gradients of their solution with respect to their inputs, and turns the calibration of model parameters and the training of hybrid machine-learning components into gradient-based optimization problems. Here we use a differentiable version of the ocean model Veros to examine what this capability brings to ocean model design and calibration. The model is written in JAX and computes sensitivities with forward- and reverse-mode automatic differentiation, with no change to its forward numerical behavior. We consider three applications. In the first, two physical parameters, lateral viscosity and bottom friction, are calibrated from satellite-like observables over a two-day control window; the calibrated model restores the simulated circulation over a long term, multi-year evaluation. In the second, a subgrid parameterization representing the backscatter of energy toward resolved scales is trained online and improves the kinetic energy and shear of a long term simulation. In the third, a neural correction is trained online to compensate a biased vertical mixing closure. The simulation is fitted over a two-day control window, and similarly to the previous applications, the correction remains stable and effective over a multi-year free-running simulation. Across the three applications, quantities optimized over windows of about one hundred solver steps yield skill that persists over integration times that are orders of magnitude longer, which indicates that solving online learning problems using differentiable programming is a practical tool for ocean model design and calibration.