Authors: Sifan Wang*, Zehao Dou*, Tongrui Liu, Lu Lu
*Equal contribution

Affiliation: Yale University

This work is implemented using JAX and Flax. Please ensure you have the latest versions of these libraries installed.

To install the FunDiff package, clone the repository and install:

git clone https://github.com/sifanexisted/fundiff.git
cd fundiff
pip install -e .

Training datasets are available for download at dataset link. Download and extract the datasets to the appropriate directory and change the data path accordinly in config files before training.

FunDiff employs a two-stage training procedure:

1. Training Function Autoencoder: Maps functions to a latent representation
2. Training Latent Diffusion Model: Learns the diffusion process in the latent space via rectified flow

Navigate to the Kolmogorov flow example directory:

cd examples/burgers/diffusion

python main.py --config=configs/autoencoder.py:fae --config.training.use_pde=True

This stage learns to encode physical functions into a compact latent representation while preserving essential structural information.

python main.py --config=configs/diffusion.py:fae,dit --config.training.use_pde=True

This stage trains the diffusion model to generate new samples in the learned latent space.

Go to the baselines directory:

bash cd examples/burgers/baselines

Run the following commands to train other baselines (e.g. FNO, Unet, ViT):

python main.py --config=configs/base.py:fno

To evaluate the trained FunDiff model and generate new samples, please check our jupyter notebooks for each example.

FunDiff/
├── function_diffusion/         # source code of function diffusion implementation
├── burgers/                    # Burgers equation examples
├── damp_sine/                  # Damped sine wave examples
├── figures/                    # pipline figure
├── kf_generation/              # Kolmogorov flow generation
├── kf_reconstruction/          # Kolmogorov flow reconstruction
├── linear_elasticity/          # Linear elasticity problems
├── turbulence_mass_transfer/   # Turbulence and mass transfer examples
├── requirements.txt            # Python dependencies
└── setup.py                    # Package installation script

If you use FunDiff in your research, please cite:

@article{wang2025fundiff,
  title={FunDiff: Diffusion Models over Function Spaces for Physics-Informed Generative Modeling},
  author={Wang, Sifan and Dou, Zehao and Liu, Tong-Rui and Lu, Lu},
  journal={arXiv preprint arXiv:2506.07902},
  year={2025}
}