NeuralOM: Neural Ocean Model for Subseasonal-to-Seasonal Simulation

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, Kun Wang, Xiaomeng Huang^*
(† Equal contribution, * Corresponding Author)

Abstract: Long-term, high-fidelity simulation of slow-changing physical systems, such as the ocean and climate, presents a fundamental challenge in scientific computing. Traditional autoregressive machine learning models often fail in these tasks as minor errors accumulate and lead to rapid forecast degradation. To address this problem, we propose NeuralOM, a general neural operator framework designed for simulating complex, slow-changing dynamics. NeuralOM's core consists of two key innovations: (1) a Progressive Residual Correction Framework that decomposes the forecasting task into a series of fine-grained refinement steps, effectively suppressing long-term error accumulation; and (2) a Physics-Guided Graph Network whose built-in adaptive messaging mechanism explicitly models multi-scale physical interactions, such as gradient-driven flows and multiplicative couplings, thereby enhancing physical consistency while maintaining computational efficiency. We validate NeuralOM on the challenging task of global Subseasonal-to-Seasonal (S2S) ocean simulation. Extensive experiments demonstrate that NeuralOM not only surpasses state-of-the-art models in forecast accuracy and long-term stability, but also excels in simulating extreme events. For instance, at a 60-day lead time, NeuralOM achieves a 13.3% lower RMSE compared to the best-performing baseline, offering a stable, efficient, and physically-aware paradigm for data-driven scientific computing. Codes link: https://github.com/YuanGao-YG/NeuralOM.

News 🚀

2025.07.28: Inference codes for global ocean forecasting are released.
2025.06.01: Inference codes for global ocean simulation are released.
2025.05.27: Paper is released on ArXiv.

Notes

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.

Quick Start

Installation

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

Inference for Global Ocean Simulation

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_simulation.sh

Inference for Global Ocean Forecasting

Preparing the test data as follows:

./data/
|--test
|  |--2020.h5
|--test_atmos
|  |--2020.h5
|--mean_s_t_ssh.npy
|--std_s_t_ssh.npy
|--mean_atmos.npy
|--std_atmos.npy
|--climate_mean_s_t_ssh.npy
|--land_mask.h5

Run the following script:

sh inference_forecasting.sh

Training

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

Performance

Global Ocean Simulation

Global Ocean Forecasting and Extreme Event Assessment

Citation

@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 Wang, Kun and Huang, Xiaomeng},
  journal={arXiv preprint arXiv:2505.21020},
  year={2025}
}

YuanGao-YG
/

NeuralOM