WDNO: Wavelet Diffusion Neural Operator (ICLR 2025)

arXiv ｜ Paper

Official repo for the paper Wavelet Diffusion Neural Operator.<br /> Peiyan Hu*, Rui Wang*, Xiang Zheng, Tao Zhang, Haodong Feng, Ruiqi Feng, Long Wei, Yue Wang, Zhi-Ming Ma, Tailin Wu†.<br /> ICLR 2025.

We introduce Wavelet Diffusion Neural Operator (WDNO), a novel method for generative PDE simulation and control, to address diffusion models' challenges of modeling system states with abrupt changes and generalizing to higher resolutions.

Framework of WDNO:

<a href="url"><img src="https://github.com/AI4Science-WestlakeU/wdno/blob/main/fig/figures1.png" align="center" width="900" ></a>

News

We develop a WDNO that shares a unified input–output interface with standard neural operators in the RealPDEBench repository.

Environment

Run the following commands to install dependencies. In particular, when running the 2D control task, the Python version must be 3.8 due to the requirement of the Phiflow software.

bash env.sh

Datasets

The datasets can be downloaded respectively in link. Please place the corresponding datasets in the data folder under the experiment directory.

Checkpoints

The checkpoints can be downloaded respectively in link. Please place the corresponding checkpoints in the results folder under the experiment directory.

1D Burgers' Equation Simulation

Prepare data for WDNO:

cd burgers
python wave_trans.py

Training of Base-Resolution Model:

bash /scripts/burgers/train_base_sim.sh

Inference of Base-Resolution Model:

bash /scripts/burgers/eval_base_sim.sh

Training of Super-Resolution Model:

bash /scripts/burgers/train_super_sim.sh

Inference of Super-Resolution Model:

bash /scripts/burgers/eval_super_sim.sh

2D Smoke Simulation

Prepare data for WDNO:

cd smoke
python wave_trans_2d.py

Training of Base-Resolution Model:

bash /scripts/smoke/train_base_sim.sh

Inference of base-resolution:

bash /scripts/smoke/inf_base_sim.sh

Training of Super-Resolution Model:

bash /scripts/smoke/train_super_sim.sh

Inference of super-resolution:

bash /scripts/smoke/inf_super_sim.sh

1D Burgers' Equation Control:

Prepare data for WDNO:

cd burgers
python wave_trans.py

Training:

bash /scripts/burgers/train_base_control.sh

Inference:

bash /scripts/burgers/eval_base_control.sh

2D Smoke Control:

Prepare data for WDNO:

cd smoke
python wave_trans_2d.py

training:

bash /scripts/smoke/train_base_control.sh

inference:

bash /scripts/smoke/inf_base_control.sh

Related Projects

• SafeDiffCon (ICML 2025): We propose safe diffusion models for PDE Control, which introduces the uncertainty quantile as model uncertainty quantification to achieve optimal control under safety constraints through both post-training and inference phases.

• CL-DiffPhyCon (ICLR 2025): We introduce an improved, closed-loop version of DiffPhyCon. It has an asynchronous denoising schedule for physical systems control tasks and achieves closed-loop control with significant speedup of sampling efficiency.

• DiffPhyCon (NeurIPS 2024): We introduce DiffPhyCon which uses diffusion generative models to jointly model control and simulation of complex physical systems as a single task.

• CinDM (ICLR 2024 spotlight): We introduce a method that uses compositional generative models to design boundaries and initial states significantly more complex than the ones seen in training for physical simulations.

Citation

If you find our work and/or our code useful, please cite us via:

@inproceedings{
  hu2025wavelet,
  title={Wavelet Diffusion Neural Operator},
  author={Peiyan Hu and Rui Wang and Xiang Zheng and Tao Zhang and Haodong Feng and Ruiqi Feng and Long Wei and Yue Wang and Zhi-Ming Ma and Tailin Wu},
  booktitle={The Thirteenth International Conference on Learning Representations},
  year={2025},
  url={https://openreview.net/forum?id=FQhDIGuaJ4}
}