Leveraging automatic differentiation for calibrating vertical mixing parameterizations

Fine-scale turbulence in the upper ocean boundary layer (OSBL) governs ocean surface stratification, and vertical exchanges of heat, momentum and matter in the ocean, which are key in the response of the oceans to changing environmental conditions. However, these turbulent processes are not explicitly represented in ocean models and their parameterization remains a significant source of uncertainty in climate models and operational prediction systems. Increasingly, systematically leveraging diverse data sources is becoming standard practice for developing and assessing OSBL parameterizations. Over the past years, data-driven automated procedures have for instance been used for calibrating the parameters of physics-based models, for developing parameterizations embedding ML components, and for proposing pure ML-based parameterizations of OSBL processes. 

This study explores the advantages of the emerging paradigm of differentiable programming for the calibration of OSBL parameterizations . We developed a benchmark tool, Tunax, implemented in JAX, a differentiable framework for Python. This benchmark includes a fully differentiable single-column model with various possible OSBL parameterizations, alongside a calibration module which tunes the coefficients of these parameterizations against a reference database. The differentiability of the model enables the application of variational techniques for parameter calibration. The reference database is a collection of  Large Eddy Simulations (LES) covering a range of typical physical conditions.

Here, we focus on the k-ε closure (Umlauf and Burchard, 2005), widely used in global ocean circulation models, and calibrate its parameters using a dataset of LES. These simulations have been designed to model the evolution of the oceanic mixed layer under various surface conditions (wind, heat fluxes and rotation). This work highlights the potential of differentiable calibration techniques to address uncertainties inherent to turbulence closures by enabling more flexible and data-informed parameterizations. Although this approach does not yet consistently outperform traditional calibration methods, it provides a promising avenue for reducing model biases associated with sub-grid scale parameterizations.