Oceanfourcast: Emulating Ocean Models with Transformers for Adjoint-based Data Assimilation

Adjoints have become a staple of the oceanic and atmospheric numerical modeling community over the past couple of decades as they are useful for tuning of dynamical models, sensitivity analyses, and data assimilation. One such application is generation of reanalysis datasets, which provide an optimal record of our past weather, climate, and ocean. For example, the state-of-the-art ocean-ice renanalysis dataset, ECCO, is created by optimally combining a numerical ocean model with heterogeneous observations through a technique called data assimilation. Data assimilation in ECCO minimizes the distance between model and observations by calculating adjoints, i.e., gradients of the loss w.r.t. simulation forcing fields (wind and surface heat fluxes). The forcing fields are iteratively updated and the model is rerun until the loss is minimized to ensure that the numerical model does not drastically deviate from the observations. Calculating adjoints, however, either requires  disproportionately high computational resources  or rewriting the dynamical model code to be autodifferentiable. 

Therefore, we ask if deep learning-based emulators can provide fast and accurate adjoints. Ocean data is smooth, high-dimensional, and has complex spatiotemporal correlations. Therefore, as an initial foray into ocean emulators, we leverage a combination of neural operators and transformers. Specifically, we have adapted the FourCastNet architecture, which has successfully emulated ERA5 weather data in seconds rather than hours, to emulate an idealized ocean simulation.

We generated a ground-truth dataset by simulating a double-gyre, an idealized representation of the North Atlantic Ocean, using MITgcm, a state-of-the-art dynamical model. The model was forced by zonal wind at the surface and relaxation to a meridional profile of temperature — warm/cold temperatures at low/high latitudes. This simulation produced turbulent western boundary currents embedded in the large-scale gyre circulation. We performed 4 additional simulations by modifying the magnitude of SST relaxation and wind forcing to introduce diversity in the dataset. From these simulations, we used 4 state variables (meridional and zonal surface velocities, pressure, and temperature) as well as the forcing fields (zonal wind velocity and relaxation SST profile) sampled in 10-day steps. The dataset was split into training, validation, and test datasets such that validation and test datasets were unseen during training. These datasets provide an ideal testbed for evaluating and comparing the performance of data-driven ocean emulators.

We used this data to train and evaluate Oceanfourcast. Our initial results in the following figure show that our model, Oceanfourcast, can successfully predict the streamfunction and pressure for a lead time of 1 month. 

We are currently working on generating adjoints from Oceanfourcast.  We expect the adjoint calculation to require significantly less compute time than that from a full-scale dynamical model like MITgcm.  Our work shows a promising path towards deep-learning augmented data assimilation and uncertainty quantification.