Deep Conditional Emulators for calibrating ocean vertical physics schemes

Differentiable programming has emerged as a powerful tool in geoscientific modelling, offering new possibilities for optimization and parameter calibration. However, this approach requires the underlying physical models to be differentiable in order to compute gradients and apply optimization algorithms. In practice, current-generation geoscientific models are generally not differentiable, which limits the use of variational approaches to calibrate their parameters. In the past few years, several strategies have been proposed to overcome this limitation.

Here, we explore the use of deep learning techniques for the calibration of vertical physics schemes of current-generation ocean models. We propose to build conditional emulators of single column ocean models to approximate the gradient of their solution with respect to their physical parameters. Our baseline is a single column ocean model, implemented in Jax, which provides a differentiable framework for the calibration of ocean vertical physics schemes. We leverage this framework to generate sets of simulations for the design of deep conditional emulators of the model, and assess their ability to approximate the gradient of the model in an inverse problem setting.

We focus on several idealized cases corresponding to different forcing conditions, starting from the Kato-Philips case. It describes the evolution of a water column with no heat flux and uniform wind friction velocity. We obtain various trajectories for uniformly sampled n-uplets defining the initial conditions, friction velocity, and physical parameters. With this dataset, we train and test different kinds of neural networks, exploring architectures and losses, to make the most of temporal and spatial dependencies.

Comparison with the fully differentiable baseline solution shows that deep conditional emulators are able to predict the system states both forward and backward, with different initial and forcing conditions, and can be used to calibrate ocean model  parameters. Our results therefore illustrate how deep emulators are a potential solution to take over the non-differentiability of existing geoscientific models, and  solve inverse problems for their calibration.