README.md

---
license: mit
language:
- en
library_name: pytorch
tags:
- OneScience
- fluid-dynamics
- neural-operator
- multiwavelet
datasets:
- OneScience-Group/fno
---

<p align="center">
  <strong><span style="font-size: 30px;">MWT</span></strong>
</p>

# Model Introduction

MWT (Multiwavelet-based Operator Learning), proposed by Gaurav Gupta, Xiongye Xiao, and Paul Bogdan, is a multiwavelet operator-learning framework. It builds fixed decomposition and reconstruction filters from orthogonal polynomials and learns differential-equation solution operators in a multiscale space, enabling data-efficient physical-field prediction and generalization across resolutions.

This repository is an independent OneScience reproduction of the two-dimensional Navier–Stokes vorticity experiment described in the paper. On a periodic unit torus, the model combines the first 10 vorticity frames with spatial and temporal coordinates and predicts the remaining \(T-10\) frames in one shot. The experiment uses a regular `64 x 64` grid downsampled from `256 x 256` data.

Paper: [Multiwavelet-based Operator Learning for Differential Equations](https://arxiv.org/abs/2109.13459)

# Model Description

MWT is a multiscale neural architecture for learning differential operators. Its pipeline is input lifting, multiwavelet decomposition, multiscale operator mapping, multiwavelet reconstruction, and vorticity projection. A linear layer first lifts the 13-dimensional input to \(c k^2=36\) features, where \(c=4\) and the Legendre multiwavelet order is \(k=3\). Fixed Legendre filter matrices from the paper and their two-dimensional Kronecker products recursively decompose both spatial dimensions.

At each scale, learnable \(A\), \(B\), and \(C\) operators transform detail and smooth coefficients. This implementation uses three-dimensional Fourier spectral and pointwise convolutions over \((x,y,t)\). The coarsest scale is processed by the \(\bar{T}\) map and reconstructed to the original resolution with fixed filters.

For the two-dimensional Navier–Stokes experiment, four MWT blocks are stacked with BatchNorm3d and ReLU between blocks. A `36 -> 128 -> 1` output head maps reconstructed features to vorticity at each grid point and prediction time.

## Intended Uses

| Use case | Description |
| --- | --- |
| Navier–Stokes vorticity prediction | Predict the remaining vorticity trajectory from the first 10 two-dimensional frames. |
| Regular-grid operator learning | Learn mappings between input and output functions on periodic regular grids. |
| Multiscale physical-field modeling | Combine fixed multiwavelet decomposition with learnable within-scale operators. |
| Fast spatiotemporal inference | Approximate a numerical solver for batched prediction within the training distribution and viscosity regime. |

# Usage

## 1. OneCode

[Launch the OneCode AI-for-Science environment](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)

## 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 should install DTK 25.04.2 or later, or the OneScience-recommended version for the cluster.

### Download the model repository from Hugging Face

```bash
pip install -U huggingface_hub
hf download OneScience-Group/MWT --local-dir ./MWT
cd MWT
```

### Install the runtime environment

**DCU environment**

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

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

```bash
hf download OneScience-Group/fno --repo-type dataset --local-dir ./data
```

Set `paths.data_root` in `config/config.yaml` to the downloaded directory. The experiments use:

- `ns_V1e-3_N5000_T50.mat`: \(\nu=10^{-3}\), 5,000 samples, and 50 frames.
- `ns_V1e-4_N10000_T30.mat`: \(\nu=10^{-4}\), 10,000 samples; the experiment uses the first 30 frames.
- `NavierStokes_V1e-5_N1200_T20.mat`: \(\nu=10^{-5}\), 1,200 samples, and 20 frames.

The main MAT variables are:

- `u`: the Navier–Stokes vorticity trajectory, converted to `[num_samples, 64, 64, T]` after loading.
- `t`: temporal coordinates, converted to `[T]`.
- `a`: the initial condition used by the numerical solver. The model uses the first 10 frames of `u` directly and does not load `a` as a separate input.

### Train

`config/config.yaml` defines four MWT Navier–Stokes experiments. The default is `ns_1e-3_t50`; choose another with `--experiment`.

```bash
python scripts/train.py \
  --config config/config.yaml \
  --experiment ns_1e-3_t50 \
  --seed 0
```

Supported experiments:

- `ns_1e-3_t50`: \(\nu=10^{-3}\), \(T=50\), nominally 1,000 training samples, 500 epochs.
- `ns_1e-4_t30_n1000`: \(\nu=10^{-4}\), \(T=30\), nominally 1,000 training samples, 500 epochs.
- `ns_1e-4_t30_n10000`: \(\nu=10^{-4}\), \(T=30\), nominally 10,000 training samples, 200 epochs.
- `ns_1e-5_t20`: \(\nu=10^{-5}\), \(T=20\), nominally 1,000 training samples, 500 epochs.

### Pretrained weights

`weight/best_model.pt` contains an MWT checkpoint trained on Navier–Stokes data and can be used directly for inference and numerical evaluation.

### Inference

Run one-shot inference on the fixed 200-sample test set recorded by the checkpoint and report the mean relative L2 error in physical space:

```bash
python scripts/inference.py \
  --config config/config.yaml \
  --checkpoint weight/best_model.pt
```

For explicit device, batch-size, and output options:

```bash
python scripts/inference.py \
  --config config/config.yaml \
  --checkpoint weight/best_model.pt \
  --device auto \
  --batch-size 1 \
  --output-dir results
```

Outputs:

```text
results/
├── inference_metrics.json
├── predictions.npy
├── targets.npy
└── sample_indices.npy
```

### Evaluation and visualization

```bash
python scripts/result.py --config config/config.yaml --sample 0
```

Generated files:

```text
results/
├── field_comparison.png
├── relative_l2_over_time.png
└── result_summary.json
```

- `field_comparison.png` compares ground-truth vorticity, MWT prediction, and absolute error at the beginning, middle, and end of the forecast interval.
- `relative_l2_over_time.png` plots spatial relative L2 over the full one-shot forecast interval.
- `result_summary.json` records experiment settings, sample counts, array shapes, error metrics, paper-reference values, comparability notes, and visualization paths.

# 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: [Multiwavelet-based Operator Learning for Differential Equations](https://proceedings.neurips.cc/paper/2021/file/c9e5c2b59d98488fe1070e744041ea0e-Paper.pdf), NeurIPS 2021; [arXiv:2109.13459](https://arxiv.org/abs/2109.13459).
- Official implementation: [gaurav71531/mwt-operator](https://github.com/gaurav71531/mwt-operator).
- This repository uses the Hugging Face-compatible MIT identifier (`mit`). The paper, upstream implementation, datasets, and all other third-party materials remain subject to their original licenses and terms.