Accelerate Coastal Ocean Circulation Model with AI Surrogate

TL;DR

This study presents a GPU-accelerated AI surrogate for ROMS coastal circulation simulations using a 4D Swin Transformer, achieving up to ~450× speedups for 12-day forecasts while enforcing mass conservation through a physics-based verifier. The approach leverages historical ROMS data, a patch-based encoder–decoder architecture, and a dual-model strategy to deliver near real-time forecasts with a fallback to ROMS when physics criteria are violated. Key contributions include the design of a high-dimensional spatiotemporal transformer for coastal domains, a rigorous data-preparation and optimization pipeline on NVIDIA DGX hardware, and a verification framework that ensures physically plausible results. The work demonstrates potential for real-time coastal hazard forecasting and rapid disaster response, with future work on extending to more coastal processes, uncertainty quantification, and larger-scale parallelism.

Abstract

Nearly 900 million people live in low-lying coastal zones around the world and bear the brunt of impacts from more frequent and severe hurricanes and storm surges. Oceanographers simulate ocean current circulation along the coasts to develop early warning systems that save lives and prevent loss and damage to property from coastal hazards. Traditionally, such simulations are conducted using coastal ocean circulation models such as the Regional Ocean Modeling System (ROMS), which usually runs on an HPC cluster with multiple CPU cores. However, the process is time-consuming and energy expensive. While coarse-grained ROMS simulations offer faster alternatives, they sacrifice detail and accuracy, particularly in complex coastal environments. Recent advances in deep learning and GPU architecture have enabled the development of faster AI (neural network) surrogates. This paper introduces an AI surrogate based on a 4D Swin Transformer to simulate coastal tidal wave propagation in an estuary for both hindcast and forecast (up to 12 days). Our approach not only accelerates simulations but also incorporates a physics-based constraint to detect and correct inaccurate results, ensuring reliability while minimizing manual intervention. We develop a fully GPU-accelerated workflow, optimizing the model training and inference pipeline on NVIDIA DGX-2 A100 GPUs. Our experiments demonstrate that our AI surrogate reduces the time cost of 12-day forecasting of traditional ROMS simulations from 9,908 seconds (on 512 CPU cores) to 22 seconds (on one A100 GPU), achieving over 450$\times$ speedup while maintaining high-quality simulation results. This work contributes to oceanographic modeling by offering a fast, accurate, and physically consistent alternative to traditional simulation models, particularly for real-time forecasting in rapid disaster response.

Figures (10)

• Figure 1: The overall workflow of our method. ROMS historical simulations are pre-processed and used to train the AI surrogate. The trained AI surrogate takes initial and boundary conditions as inputs and predicts the interior values. A verification module checks physical law adherence of AI predictions and switches to ROMS if the surrogate fails the check.
• Figure 2: The architecture of the AI surrogate. The encoder processes 2D and 3D variables via patch partitioning and patch embedding, then merges them along the depth dimension. The 4D Swin Transformer with Positional Embedding and Patch Merging downsamples the 4D data. The decoder employs Transposed Convolutions, Batch Normalization, GELU activation, and Skip Connection from the corresponding encoder layers to upsample the variables and finally recover them to their original spatial and temporal size.
• Figure 3: Swin transformer block. (a) In this example, the size of input patches and the windows are $8 \times 8\times 8 \times 4$ and $4 \times 4 \times 4 \times 2$, respectively. Following the window partitioning and shifting methods, the numbers of grouped windows in W-MSA and SW-MSA become $2\times 2\times 2\times 2 =16$ and $3 \times 3 \times 3 \times 3 = 81$, respectively. (Source: kim2024swiftliu2022video)
• Figure 4: Patch merging is applied to the spatial dimensions, while the temporal dimension remains unchanged. In this example, a tensor of size 8$\times$8$\times$8$\times$3 is merged to 4$\times$4$\times$4$\times$3.
• Figure 5: Visualization of ROMS and AI surrogate forecasting results for $u, v$, and $\zeta$ at 12:30 p.m. on April 1, 2012 ($u,v$ are on the surface level), using the initial condition from 12 p.m. on March 26, 2012. The AI surrogate accurately captures some of the spatial patterns in coastal circulation.