OneScience-Group/FluxNet-ShallowWater-LAP

FluxNet-ShallowWater-LAP

Model Introduction

FluxNet-ShallowWater-LAP reproduces the two-dimensional shallow-water FluxNet-LAP pushforward experiment in the FluxNet paper. It forecasts water depth and two momentum fields on a 64x64 grid.

Model Description

The model takes (h, mx, my) and predicts the next state. The configuration uses 64 channels, six residual blocks, 5x5 kernels, a 3x3 transport neighborhood and five-step pushforward training.

Use Cases

Use case	Description
Shallow-water forecasting	Forecast depth and momentum on a regular grid
Simulation acceleration	Compare surrogate rollouts against numerical reference trajectories

Usage

1. OneCode

Open OneCode

2. Manual Setup

Hardware requirements

• Training and inference require a Slurm-allocated HIP/DCU compute node; there is no silent CPU fallback.
• The formal run used one DCU, eight CPU cores and 30 GB of host memory. GPU execution was not validated.

Download the Model

hf download OneScience-Group/FluxNet-ShallowWater-LAP --local-dir ./FluxNet-ShallowWater-LAP
cd FluxNet-ShallowWater-LAP

Install Dependencies

DCU environment

conda create -n onescience311 python=3.11 -y
conda activate onescience311
pip install onescience[cfd-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
pip install numpy h5py PyYAML matplotlib

GPU environment

conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
conda activate onescience311
pip install onescience[cfd-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
pip install numpy h5py PyYAML matplotlib

The GPU commands describe dependency installation only; the provided scripts were validated on Slurm HIP/DCU.

Training Data

The OneScience/FluxNet-ShallowWater dataset on ModelScope contains 50 training, 20 validation, 50 short-test and 50 long-test trajectories produced with the paper's official numerical generator. Each HDF5 trajectory contains h, mx, my and t. The dataset is not mirrored on Hugging Face. Download it into the model directory:

modelscope download --dataset OneScience/FluxNet-ShallowWater --local_dir ./data

Run commands from the package root. data.root in config/config.yaml defaults to ./data.

Training

On a Slurm-allocated DCU node:

python scripts/train.py

The full run trains for 300 epochs and writes weight/best_model.pth. Back up the included best_model.pth and last_model.pth before retraining if they must be preserved.

Trained Weights

The included checkpoint is the best validation-MSE model at epoch 283 of a completed 300-epoch run. weight/best_model.pth is 15636928 bytes; SHA256: 47fb0a029b0c5baa10b2452c70b12f54c958b548cb9aff300b85183d55f1ae70.

Inference

python scripts/inference.py

Evaluation and Visualization

python scripts/result.py

Outputs are written to results/. At the long-horizon endpoint (4.8 time units), MAE for (h, mx, my) is (1.5568e-3, 1.9634e-3, 1.7648e-3); averaging over all predicted steps gives (1.4063e-3, 1.8558e-3, 1.7632e-3). The paper's Table 3 two-times extrapolation means are (1.75e-3, 2.38e-3, 2.30e-3). This run uses one seed rather than the paper's five independently trained models, so the figures are not a strict like-for-like reproduction.

OneScience Repositories

Platform	Main repository	Skills repository
Gitee	https://gitee.com/onescience-ai/onescience	https://gitee.com/onescience-ai/oneskills
GitHub	https://github.com/onescience-ai/OneScience	https://github.com/onescience-ai/oneskills

Citation and License

• Paper: FluxNet.
• Original public repository; public dataset and its provenance. This model package is an independent reproduction.