DOCSEARCHbyAlgolia
Submit Docs

DiffusionPDE

repository·main·Indexed 41 days ago

https://github.com/jhhuangchloe/diffusionpde

A PyTorch implementation for solving Partial Differential Equations (PDEs) under partial observation using generative diffusion models. It supports forward, inverse, and joint problems for physical equations including Darcy Flow, Navier-Stokes, Burgers', Poisson, and Helmholtz.

Tokens
936
Snippets
3
Records
5
Agent score
2%

What's inside DiffusionPDE

  1. Tuning Guidance Weights for PDE Solving

    main

    When using generate_pde.py, guidance weights significantly impact results.

    Best Practices:

    • Weight Range: For forward observation loss terms, weights typically range from 0.1 to 10 (before the loss terms are averaged).
    • Oversaturation: Using excessively large weights can cause artifacts and oversaturation.
    • Scaling: Weights are sensitive to the scaling of the specific loss terms associated with the PDE being solved.
    • Stochasticity: The diffusion model is stochastic; results may vary with different random seeds.

  2. Download and Prepare Datasets

    main

    Datasets for Darcy Flow, Poisson, Helmholtz, Navier-Stokes, and Burgers' equations are available via Google Drive.

    Data Formats:

    • Darcy Flow, Poisson, Helmholtz: Shape [N, X, Y] (N instances, X/Y spatial resolutions).
    • Navier-Stokes (bounded/non-bounded): Shape [N, X, Y, T] (T time steps).
    • Burgers' equation: Shape [N, X, T].

    Setup Steps:

    1. Download training.zip and testing.zip.
    2. Unzip them into the data/ directory.
    3. Ensure raw data is scaled to the range (-1, 1) before training.

  3. Train Diffusion Models

    main

    Training is performed using a script derived from EDM. You must first prepare the data using merge_data.py to ensure it is correctly formatted and scaled.

    To train a new model (e.g., for Darcy Flow), follow these steps:

    1. Run merge_data.py to prepare .npy files.
    2. Use torchrun to execute train.py with the desired hyperparameters.
    # Prepare the .npy files for training.
python3 merge_data.py # Darcy Flow

# Train the diffusion model.
torchrun --standalone --nproc_per_node=3 train.py --outdir=pretrained-darcy-new --data=/data/Darcy-merged/ --cond=0 --arch=ddpmpp --batch=60 --batch-gpu=20 --tick=10 --snap=50 --dump=100 --duration=20 --ema=0.05

  5. Solve PDE Problems using generate_pde.py

    main

    The generate_pde.py script is used to perform inference for various PDE solving tasks by passing a configuration YAML file. Supported tasks include:

    • Forward Problem: Solve with sparse observations on the coefficient (or initial state) space.
    • Inverse Problem: Solve with sparse observations on the solution (or final state) space.
    • Joint Recovery: Simultaneously solve coefficient/initial state and solution/final state spaces with observations on both sides.
    • Time-series Recovery: Recover the solution throughout a time interval with sparse sensors.
    # Solve Forward Problem
python3 generate_pde.py --config configs/darcy-forward.yaml

# Solve Inverse Problem
python3 generate_pde.py --config configs/darcy-inverse.yaml

# Recover Both Spaces (Coefficient and Solution)
python3 generate_pde.py --config configs/darcy.yaml

# Solve Solution Over Time (e.g., Burgers)
python3 generate_pde.py --config configs/burgers.yaml