NeuralOM: Neural Ocean Model for Subseasonal-to-Seasonal Simulation
Yuan Gao† , Ruiqi Shu† , Hao Wu† ,Fan Xu, Yanfei Xiang, Ruijian Gou, Qingsong Wen, Xian Wu, Xiaomeng Huang*
(† Equal contribution, * Corresponding Author)
Abstract: Accurate Subseasonal-to-Seasonal (S2S) ocean simulation is critically important for marine research, yet remains challenging due to its substantial thermal inertia and extended time delay. Machine learning (ML)-based models have demonstrated significant advancements in simulation accuracy and computational efficiency compared to traditional numerical methods. Nevertheless, a significant limitation of current ML models for S2S ocean simulation is their inadequate incorporation of physical consistency and the slow-changing properties of the ocean system. In this work, we propose a neural ocean model (NeuralOM) for S2S ocean simulation with a multi-scale interactive graph neural network to emulate diverse physical phenomena associated with ocean systems effectively. Specifically, we propose a multi-stage framework tailored to model the ocean's slowly changing nature. Additionally, we introduce a multi-scale interactive messaging module to capture complex dynamical behaviors, such as gradient changes and multiplicative coupling relationships inherent in ocean dynamics. Extensive experimental evaluations confirm that our proposed NeuralOM outperforms state-of-the-art models in S2S and extreme event simulation. The codes are available at \url{https://github.com/YuanGao-YG/NeuralOM}.
- 2025.06.01: Codes for inference are released.
- 2025.05.27: Paper is released on ArXiv.
The intact project is avilable on Hugging Face, you can find the pretrained models, test data on Hugging Face and put them in the same location.
- cuda 11.8
# git clone this repository
git clone https://github.com/YuanGao-YG/NeuralOM.git
cd NeuralOM
# create new anaconda env
conda env create -f environment.yml
conda activate neuralom
Preparing the test data as follows:
./data/
|--test
| |--2020.h5
|--mean_s_t_ssh.npy
|--std_s_t_ssh.npy
|--climate_mean_s_t_ssh.npy
|--land_mask.h5
Run the following script:
sh inference.sh
The training codes will be released after the paper is accepted.
1. Prepare Data
Preparing the train, valid, and test data as follows:
./data/
|--train
| |--1993.h5
| |--1994.h5
| |--......
| |--2016.h5
| |--2017.h5
|--valid
| |--2018.h5
| |--2019.h5
|--test
| |--2020.h5
|--mean_s_t_ssh.npy
|--std_s_t_ssh.npy
|--climate_mean_s_t_ssh.npy
|--land_mask.h5
For data ranging from 1993 to 2020, each h5 file includes a key named 'fields' with the shape [T, C, H, W] (T=365/366, C=97, H=361, W=720)
2. Model Training
-
Single GPU Training
Continue update
-
Single-node Multi-GPU Training
Continue update
-
Multi-node Multi-GPU Training
Continue update
@article{gao2025neuralom,
title={NeuralOM: Neural Ocean Model for Subseasonal-to-Seasonal Simulation},
author={Gao, Yuan and Shu, Ruiqi and Wu, Hao and Xu, Fan and Xiang, Yanfei and Gou, Ruijian and Wen, Qingsong and Wu, Xian and Huang, Xiaomeng},
journal={arXiv preprint arXiv:2505.21020},
year={2025}
}