Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Towards fully differentiable neural ocean model with Veros

Etienne Meunier, Said Ouala, Hugo Frezat, Julien Le Sommer, Ronan Fablet

TL;DR

This work presents a differentiable extension of the VEROS ocean model by introducing modifications that make its step function compatible with JAX automatic differentiation. The authors provide forward and backward validation, demonstrating accurate gradient propagation and negligible impact on forward dynamics. They showcase two applications: gradient-based correction of an initial state and calibration of key physical parameters (Ah and r_bot) in an idealized ACC setup using BSF observations. The results illustrate the feasibility and potential of differentiable ocean modeling for data assimilation, parameter estimation, and physics-informed learning, while outlining future directions to improve efficiency and scalability for long-horizon simulations.

Abstract

We present a differentiable extension of the VEROS ocean model, enabling automatic differentiation through its dynamical core. We describe the key modifications required to make the model fully compatible with JAX autodifferentiation framework and evaluate the numerical consistency of the resulting implementation. Two illustrative applications are then demonstrated: (i) the correction of an initial ocean state through gradient-based optimization, and (ii) the calibration of unknown physical parameters directly from model observations. These examples highlight how differentiable programming can facilitate end-to-end learning and parameter tuning in ocean modeling. Our implementation is available online.

Towards fully differentiable neural ocean model with Veros

TL;DR

This work presents a differentiable extension of the VEROS ocean model by introducing modifications that make its step function compatible with JAX automatic differentiation. The authors provide forward and backward validation, demonstrating accurate gradient propagation and negligible impact on forward dynamics. They showcase two applications: gradient-based correction of an initial state and calibration of key physical parameters (Ah and r_bot) in an idealized ACC setup using BSF observations. The results illustrate the feasibility and potential of differentiable ocean modeling for data assimilation, parameter estimation, and physics-informed learning, while outlining future directions to improve efficiency and scalability for long-horizon simulations.

Abstract

We present a differentiable extension of the VEROS ocean model, enabling automatic differentiation through its dynamical core. We describe the key modifications required to make the model fully compatible with JAX autodifferentiation framework and evaluate the numerical consistency of the resulting implementation. Two illustrative applications are then demonstrated: (i) the correction of an initial ocean state through gradient-based optimization, and (ii) the calibration of unknown physical parameters directly from model observations. These examples highlight how differentiable programming can facilitate end-to-end learning and parameter tuning in ocean modeling. Our implementation is available online.

Paper Structure

This paper contains 18 sections, 1 equation, 4 figures, 1 table.

Table of Contents

  1. Introduction
  2. Differentiability properties in Veros
  3. Modifications to Veros
  4. Functionally pure step function.
  5. Handling functions with singular gradients.
  6. Computational graph preservation.
  7. Forward behaviour
  8. Backward behaviour
  9. Toy experiment: gradient-based correction of initial states
  10. Application to model calibration problems
  11. Idealised Antarctic Circumpolar Current (ACC) model configuration
  12. Sensitivity analysis for parameter calibration
  13. Calibration experiment analysis
  14. Conclusion and future work
  15. Experimental configuration
  16. ...and 3 more sections

Figures (4)

  • Figure 1: Results from the initial field reconstruction experiment. Left: Loss and distance metrics during optimization. Right: Temperature field snapshots showing the reference field, initial perturbed field, and final reconstructed field after optimization.
  • Figure 2: Left: Loss landscape for the calibration task. Color indicates the loss for each combination of parameter and vector field gradients direction. Right: Results of calibration experiment. ACC transport (Sverdrup) and Snapshots of the barotropic streamfunction.
  • Figure 3: Gradient of the loss with respect to $r_{\rm bot}$ and the accuracy of automatic differentiation methods over iterations. Left: Evolution of $\partial \mathcal{L} / \partial r_{\rm bot}$ for different methods. Right: Accuracy of forward-mode (JVP) and reverse-mode (VJP) automatic differentiation, computed as $1 - |\text{autodiff grad} - \text{finite diff}| / |\text{finite diff}|$, highlighting the agreement with numerical gradients over iterations.
  • Figure 4: Gradient computation time using vector-Jacobian products (VJP) for optimizing the initial temperature field with respect to a loss function evaluated after multiple Veros integration steps.