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

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](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.