# Riemannian Denoising Diffusion Probabilistic Models

This is the code Repository for the paper `Riemannian Denoising Diffusion Probabilistic Models`.


## Dependencies
RDDPM is built in Python 3.8 using PyTorch 2.2.0. Please use the following command to install the requirements:
```
conda env create -f environment.yaml
```

## Experiments

This repo supports experiments on the following datasets:
- Earth and climate science datasets: `Volcano`, `Earthquake`, `Flood`, and `Fire`
- Mesh datasets: `Standford Bunny` and `Spot the Cow`
- Special orthogonal group: `SO(10)` with `m=3` and `m=5`
- Alanine dipeptide with fixed dihedral angle

In each experiment, we first present the scripts used to generate the dataset, followed by the command to run the program. The commands for dataset generation are not required, as the training data is already available in the `./data` folder.

### Earth and climate science datasets
To run experiments on the Earth and climate science datasets, use the following commands:
```
python main.py --multirun experiment=fire seed=0,1,2,3,4
python main.py --multirun experiment=earthquake seed=0,1,2,3,4
python main.py --multirun experiment=flood seed=0,1,2,3,4
python main.py --multirun experiment=volcano seed=0,1,2,3,4
```


### Mesh data on learned manifolds

Create triangular mesh datasets with the following command:
```
python generate_data.py --manifold_name mesh_gen
```

Train the function $\xi$ by running the following command. The trained models are provided in the folder `./constraint/model`

```
python train_constraint.py --dataset_name bunny_whole
python train_constraint.py --dataset_name spot_whole
```

Refinement for the dataset such that all points belong to the learned manifold.
```
python generate_data.py --manifold_name mesh_refine
```

Run the following commands for training the models:
```
python main.py --multirun experiment=bunny_eigfn049 seed=0,1,2,3,4 
python main.py --multirun experiment=bunny_eigfn099 seed=0,1,2,3,4 
python main.py --multirun experiment=spot_eigfn049 seed=0,1,2,3,4
python main.py --multirun experiment=spot_eigfn099 seed=0,1,2,3,4 
```

### High-dimensional special orthogonal group

Generate synthetic dataset of the special orthogonal group using:
```
python generate_data.py --manifold_name SOn
```
Then, train the models using:
```
python main.py --multirun experiment=SO10_5w
python main.py --multirun experiment=SO10_3w
```


### Alanine dipeptide
Generate the dataset and obtain reference conformations and prior samples using:
```
python generate_data.py --manifold_name MD
```

Train the models with:
```
python main.py --multirun experiment=dipeptide
```

