Fully differentiable transport operators enable gradient-based parameter tuning and data assimilation of marine biogeochemical models

Marine biogeochemical models typically contain tens to hundreds of parameters and are notoriously challenging to tune to sparse and noisy observations, in particular for specific regional conditions. While ensemble-based methods can automate this process and are also used for data assimilation, they do not scale well to large numbers of unknown parameters. Gradient-based methods, on the other hand, scale well with high dimensionalities but require adjoint models. However, state-of-the-art differentiable programming frameworks such as PyTorch eliminate the need for manual adjoint implementations through automatic differentiation, that is, by using the chain rule to automatically compute analytic derivatives.

We introduce a fully differentiable framework for tracer transport and marine biogeochemical (BGC) processes in PyTorch. We implement advection and diffusion operators based on popular models written in Fortran, e.g. the General Ocean Turbulence Model (GOTM) for water columns. As GOTM's vertical mixing formulation requires implicit time stepping, we provide efficient differentiable solvers for batched tridiagonal systems with custom backward methods derived by implicit differentiation. Furthermore, our framework includes a PyTorch base class for differentiable BGC models with an interface similar to the Framework for Aquatic Biogeochemical Models (FABM). We provide several examples, including a re-implementation of the popular ecosystem model ECOSMO. As our operators are implemented in PyTorch, they can easily be combined with established neural network layers and optimizers.

We demonstrate our framework by performing model tuning and data assimilation in BGC models using 4DVar on sparse and noisy observations. We investigate the scaling behaviour of our tridiagonal solver for various batch and system sizes with both GPU and CPU computation. Our contribution has the potential to enhance data assimilation, speed up parameter tuning workflows and improve the accuracy of biogeochemical modelling.