Deep learning in the abyss: a stratified Physics Informed Neural Network for data assimilation

The reconstruction of deep ocean currents is a major challenge in data assimilation due to the scarcity of interior data. In this work, we present a proof of concept for deep ocean flow reconstruction using a Physics-Informed Neural Network (PINN), a machine learning approach that offers an alternative to traditional data assimilation methods. We introduce an efficient algorithm called StrAssPINN (for Stratified Assimilation PINNs), which assigns a separate network to each layer of the ocean model while allowing them to interact during training. The neural network takes spatiotemporal coordinates as input and predicts the velocity field at those points. Using a SIREN architecture (a multilayer perceptron with sine activation functions), which has proven effective in various contexts, the network is trained using both available observational data and dynamical priors enforced at several collocation points. We apply this method to pseudo-observed ocean data generated from a 3-layer quasi-geostrophic model, where the pseudo-observations include surface-level data akin to SWOT observations of sea surface height, interior data similar to ARGO floats, and a limited number of deep ARGO-like measurements in the lower layers. Our approach successfully reconstructs ocean flows in both the interior and surface layers, demonstrating a strong ability to resolve key ocean mesoscale features, including vortex rings, eastward jets associated with potential vorticity fronts, and smoother Rossby waves. This work serves as a prelude to applying StrAssPINN to real-world observational data.