Accelerate Coastal Ocean Circulation Model with AI Surrogate

2024-10-19

AI Summary

To address the high computational cost of coastal ocean circulation models, which hinders real-time disaster early warning, this paper proposes a physics-constrained 4D Swin Transformer neural surrogate model to replace traditional ROMS-based tidal wave propagation simulation. The model takes four-dimensional spatiotemporal grids as input and incorporates a physics-consistency loss function, enabling end-to-end, human-verification-free, high-fidelity modeling. Evaluated on an NVIDIA DGX-2 (A100) platform, it reduces the runtime for 12-day hindcast and forecast from 9,908 seconds (using 512 CPU cores) to just 22 seconds—a speedup exceeding 450×—while maintaining accuracy comparable to ROMS. Key innovations include: (i) the first 4D Swin architecture tailored for ocean dynamics; (ii) a differentiable, physics-driven constraint mechanism; and (iii) a lightweight, high-reliability AI surrogate paradigm suitable for operational deployment.

Technology Category

• Planning, Routing, and Scheduling: Optimization of Spatio-temporal Systems
• Computer Vision: Diffusion Models for Vision
• Machine Learning: Hardware-aware ML

Application Category

• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Search and Retrieval-Augmented AI: Large language models for search
• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs

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$ imes$ 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.

Problem

• Accelerate coastal ocean circulation simulations using AI
• Replace time-consuming traditional models with faster AI surrogates
• Ensure accuracy in coastal tidal wave forecasting

Innovation

• AI surrogate with 4D Swin Transformer
• Physics-based constraint for reliability
• Fully GPU-accelerated workflow optimization