OneScience-Group/MWT

MWT

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

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

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

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

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/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.

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:

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

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

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

Outputs:

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

Evaluation and visualization

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

Generated files:

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, NeurIPS 2021; arXiv:2109.13459.
• Official implementation: 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.