OneScience-Group/NODM-Kolmogorov

NODM

Model Overview

NODM (Neural Operator Diffusion Model) combines a neural operator with a conditional diffusion model for turbulence prediction. The neural operator first learns the deterministic spatiotemporal evolution of a flow field. Conditioned on that prediction, the diffusion model then generates the target flow field, recovering high-frequency, small-scale structures that tend to be lost by a neural operator alone and improving the representation of the turbulent energy spectrum.

This repository is an independent reproduction of the NODM Kolmogorov-flow experiment implemented through the OneScience workflow from the paper description.

Paper: Integrating Neural Operators with Diffusion Models Improves Spectral Representation in Turbulence Modeling

Model Description

NODM consists of a deterministic three-dimensional Fourier Neural Operator (FNO) and a conditional diffusion model. The FNO receives historical vorticity fields together with spatial and temporal coordinates. Four spectral layers learn the spatiotemporal evolution of two-dimensional Kolmogorov flow and predict future vorticity fields. The FNO is then frozen, and its predictions condition a U-Net-based EDM diffusion model. During inference, the model produces its final output through 32 diffusion sampling steps, preserving large-scale dynamics while restoring high-wavenumber energy and fine-grained vortex structures.

Use Cases

Use Case	Description
Kolmogorov-flow prediction	Predict future vorticity evolution from historical two-dimensional vorticity fields
Two-dimensional turbulence surrogate modeling	Learn the spatiotemporal map of the two-dimensional Navier-Stokes vorticity equation on a periodic domain
Turbulent energy-spectrum recovery	Improve high-wavenumber spectral representation with a conditional diffusion model
Generative correction of neural operators	Study the combination of deterministic operator prediction and probabilistic diffusion-based generation

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.
• A CPU can be used for import checks and small-scale pipeline validation, but full training and inference will be slow.
• DCU users must install DTK in advance. DTK 25.04.2 or later, or the OneScience-recommended version for the target cluster, is recommended.

Download the Model Package from Hugging Face

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

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

Training Data

The project generates its training data locally with a Fourier pseudospectral solver configured for the Kolmogorov-flow experiment in the paper. Run:

python scripts/generate_data.py

The generated data is saved to data/kolmogorov/traj.npy by default, so no external dataset download is required.

The dataset contains 1,000 two-dimensional vorticity trajectories and is randomly split into 80% training, 10% validation, and 10% test samples. Numerical simulation is performed on a $512 \times 512$ grid over the periodic domain $[0, 2\pi]^2$, with kinematic viscosity $10^{-5}$ and an internal time step of $10^{-4}$. Each trajectory stores the initial state and 200 time intervals and is uniformly downsampled to $128 \times 128$. The main generated files are:

• traj.npy: vorticity trajectories in NHWT layout with shape (1000, 128, 128, 201) and float32 data type
• traj.npy.generation.json: generation parameters, completion status, temporal coordinates, and provenance information for validation and resumable generation

The model samples the first 80 stored time points at a temporal stride of 2, producing 40 frames. The first 20 frames are inputs, and the remaining 20 frames are prediction targets.

Training

Full training:

• Training parameters are read from the data, model, diffusion, and training sections of config/config.yaml.
• Stage one trains the FNO for 500 epochs. Stage two freezes the FNO and trains the conditional diffusion model for 10,000 epochs.
• Before training, generate data/kolmogorov/traj.npy, or use --data-path to specify a trajectory file that follows the NHWT data contract.
• --smoke-test is only for software and pipeline validation. Its synthetic data and metrics must not be used for paper-level accuracy comparisons.

python scripts/train.py

Stage-by-stage training:

python scripts/train.py --stage operator
python scripts/train.py --stage diffusion --init-checkpoint weight/best_model.pt

Minimal end-to-end validation:

python scripts/train.py --smoke-test

Training reports loss, flow-field relative MSE, and energy-spectrum relative MSE. The best FNO checkpoint is selected by validation flow-field relative MSE; the best diffusion checkpoint is selected by validation energy-spectrum relative MSE. The checkpoint is saved to:

./weight/best_model.pt

Model Weights

The best fully trained NODM checkpoint is stored at weight/best_model.pt. A checkpoint intended for full inference must have completed the diffusion-training stage.

Inference

python scripts/inference.py

Predictions, per-sample metrics, and aggregate metrics are saved under results/.

Evaluation and Visualization

python scripts/result.py

The evaluation report, paper-metric comparison plots, training curves, and flow-field visualizations are saved under results/.

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

• Original NODM paper: Integrating Neural Operators with Diffusion Models Improves Spectral Representation in Turbulence Modeling
• This repository retains attribution to the original paper and the Kolmogorov-flow data-generation method. Use of this project must comply with the licenses and terms that apply to the original paper and related resources.