README.md

---
frameworks:
- ""
language:
- en
license: apache-2.0
tags:
- OneScience
- UNO
- neural-operator
- computational-fluid-dynamics
- Navier-Stokes
tasks: []
---
<p align="center">
  <strong>
    <span style="font-size: 30px;">UNO</span>
  </strong>
</p>

# 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](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)

## 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

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

### Set Up the Runtime Environment

**DCU Environment**

```bash
# 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**
```bash
# 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`:

```text
NavierStokes_V1e-5_N1200_T20.mat
```

Place the file in `data/`:

```text
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:

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

### Training

```bash
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

```bash
python scripts/inference.py
```

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

```text
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:

```text
[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.