OneScience/UNO

UNO

Model Overview

UNO (U-shaped Neural Operator) combines the multiscale encoder–decoder structure of U-Net with the spectral convolutions of a Fourier Neural Operator. It learns PDE solution operators across multiple spatial scales while preserving local flow-field details through skip connections.

This model package targets Navier–Stokes time-series prediction on two-dimensional regular grids. By default, it takes the first 10 time steps as input and predicts the subsequent 10 autoregressively.

Paper: U-NO: U-shaped Neural Operators
https://arxiv.org/abs/2204.11127

Repository Overview

This repository is a minimal, self-contained, runnable UNO model package maintained by OneScience for ModelScope downloads, automated OneCode execution, and rapid local validation.

Supported capabilities:

• Train a two-dimensional UNO model from a YAML configuration
• Perform multistep autoregressive prediction of Navier–Stokes flow fields
• Compute relative L2 errors and save predicted tensors and visualizations
• Override the sample count, temporal window, spatial downsampling, and model size from the command line
• Run on a CPU, GPU, or DCU

Unsupported capabilities:

• Pretrained weights are not bundled
• The approximately 394 MiB raw Navier–Stokes data file is not bundled
• The standalone model includes only the two-dimensional UNO implementation required by this example; the general-purpose OneScience 1D and 3D components are not included

Use Cases

Use Case	Description
Flow-field time-series prediction	Autoregressively predict future Navier–Stokes states from historical vorticity fields
Neural operator training	Evaluate the combination of Fourier spectral convolutions and a U-shaped multiscale architecture
CFD surrogate modeling	Learn mappings from historical to future fields on regular grids
Pipeline validation	Validate training and inference with a small sample set and a single epoch

File Structure

Path	Purpose	Notes
README.md	Project documentation	English
conf/config.yaml	Data, model, training, and output configuration	Paths are resolved relative to the model package root
model/uno.py	Standalone two-dimensional UNO model	Does not depend on onescience.modules
scripts/common.py	Shared configuration, device, metric, and autoregressive utilities	Used by both training and inference
scripts/train.py	Training and validation script	Saves the checkpoint with the best relative L2 error
scripts/inference.py	Inference, evaluation, and visualization script	Loads weight/*.pt
data/	Navier–Stokes data directory	Stores the benchmark .mat file
weight/	Model weight directory	Checkpoints are written automatically during training
result/	Inference result directory	Created automatically on first inference

Usage

1. OneCode

Use the online OneCode environment for an intelligent, one-click AI for Science (AI4S) programming experience:

Launch OneCode for one-click AI4S programming

2. Manual Setup

Hardware Requirements

• A GPU or DCU is recommended for full training.
• A CPU can be used for pipeline validation with a reduced model and dataset, but full training will be slow.
• DCU users must install DTK and a PyTorch environment compatible with the target cluster.

Download the Model Package

modelscope download --model OneScience/UNO --local_dir ./UNO
cd UNO

Set Up the Runtime Environment

DCU Environment

# Activate DTK and Conda first
conda create -n onescience311 python=3.11 -y
conda activate onescience311
# Installation with uv is also supported
pip install onescience[cfd-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai

GPU Environment

# Activate Conda first
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
# Installation with uv is also supported
pip install onescience[cfd-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai

3. Quick Start

Prepare the Data

Download the following file from the Transolver PDE-Solving-StandardBenchmark:

NavierStokes_V1e-5_N1200_T20.mat

Place the file in data/:

data/
  NavierStokes_V1e-5_N1200_T20.mat

Data downloads and documentation:

• https://github.com/thuml/Transolver/tree/main/PDE-Solving-StandardBenchmark
• https://drive.google.com/drive/folders/1UnbQh2WWc6knEHbLn-ZaXrKUZhp7pjt-

The OneScience community also provides the training data. Download it with the following command and verify that the data path in conf/config.yaml is configured correctly:

modelscope download --dataset OneScience/cfd_benchmark  data/ns/NavierStokes_V1e-5_N1200_T20.mat  --local_dir ./data

Training

python scripts/train.py

The default checkpoint is saved to weight/uno_navier_stokes.pt. The training script is controlled entirely by conf/config.yaml and does not accept command-line configuration arguments.

Model Weights

This repository provides weights trained on the standard Navier–Stokes dataset in the weight/ directory.

Inference, Evaluation, and Visualization

python scripts/inference.py

The script loads weight/uno_navier_stokes.pt by default and generates the following files under result/:

result/
  prediction_sample.pt
  prediction_sample.png

The inference script is likewise controlled entirely by conf/config.yaml and does not accept command-line configuration arguments. The weight path is determined jointly by training.weight_dir and training.checkpoint_name; the output directory and number of saved samples are controlled by inference.result_dir and inference.num_samples, respectively.

Configuration

conf/config.yaml contains five sections:

• common: device and random seed
• datapipe: data file, sample splits, temporal windows, downsampling, and DataLoader settings
• model: hidden channels, Fourier modes, normalization, and spatial padding
• training: optimizer, learning-rate schedule, early stopping, and weight directory
• inference: inference result directory and number of saved samples

The model automatically updates in_dim from t_in * out_dim in the configuration, so the model input dimension does not need to be synchronized manually when the history window changes.

Data Format

The .mat file must contain a variable named u with the following standard shape:

[1200, 64, 64, 20]

Dimension	Meaning
1200	Number of independent flow-field samples
64, 64	Height and width of the two-dimensional regular grid
20	Number of consecutive time steps

The data pipeline produces:

• pos: [H*W, 2], normalized two-dimensional coordinates
• x: [H*W, t_in*out_dim], historical states
• y: [H*W, t_out*out_dim], future states

Official OneScience Resources

Platform	OneScience 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

Citations and License

• Rahman, M. A., Ross, Z. E., and Azizzadenesheli, K. U-NO: U-shaped Neural Operators. arXiv:2204.11127, 2022.
• Li, Z. et al. Fourier Neural Operator for Parametric Partial Differential Equations. arXiv:2010.08895, 2020.
• This model package is released under the Apache-2.0 license and retains attribution to the original paper and data sources.