OneScience-Group/DyMixOp-2dBurgers

DyMixOp-2dBurgers

Model Introduction

DyMixOp-2dBurgers is a two-dimensional Burgers-equation flow predictor based on the DyMixOp paper. It learns a surrogate time-marching solver for velocity fields.

Paper: DyMixOp: Dynamic Mixture of Operators for Learning Across Heterogeneous PDEs

Model Description

DyMixOp-2dBurgers combines local and global operators with time-scale-adaptive dynamics layers to roll out the two-channel u/v velocity field. It takes the first 10 frames as input and predicts the next 10 frames at a spatial resolution of 64 x 64.

Intended Uses

Use case	Description
Fluid-equation surrogate	Predict velocity-field sequences for the two-dimensional Burgers equation.
Neural-operator reproduction	Train, evaluate, and visualize the DyMixOp local-global operator architecture.

Usage

1. OneCode

Launch the OneCode AI-for-Science environment

2. Manual Setup

Hardware requirements

• A GPU or DCU is recommended.
• A CPU can run imports and small connectivity checks, but full training and inference will be slow.
• DCU users must install a DTK version compatible with the target cluster.

Download the model repository from Hugging Face

pip install -U huggingface_hub
hf download OneScience-Group/DyMixOp-2dBurgers --local-dir ./DyMixOp-2dBurgers
cd DyMixOp-2dBurgers

Install the runtime environment

DCU environment

# Activate DTK first.
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

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

Download the training dataset from Hugging Face

hf download OneScience-Group/DyMixOp-Benchmarks \
  "2dBurgers_1200x20x2x64x64_dt0.0025_t[0_0.5]_nu0.005.mat" \
  --repo-type dataset \
  --local-dir ./data

The data variable is uv with shape (1200, 21, 2, 64, 64). The first 1,000 trajectories are used for training and the remaining 200 for testing.

Train

python scripts/train.py

The best checkpoint is saved to weight/best_model.pth, which is also included for direct inference.

Inference

python scripts/inference.py

Inference on 200 test trajectories produced a normalized relative MSE of 1.0623358757584356e-05, a physical relative MSE of 2.31623665895313e-04, an RMSE of 0.007739236151728116, and an MAE of 0.005377683386206627. The Large configuration in the paper reports a reference relative MSE of 9.18e-04; the values here were measured with the current OneScience implementation and run configuration.

Evaluation and visualization

python scripts/result.py

Generated files include:

results/final_time_first_sample.png
results/final_time_last_sample.png
results/training_curves.png

OneScience

Platform	OneScience repository	OneSkills 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: DyMixOp: Dynamic Mixture of Operators for Learning Across Heterogeneous PDEs.
• This model repository uses the Hugging Face-compatible Apache License 2.0 identifier (apache-2.0). Dataset files and other third-party assets remain subject to their original terms.