Ecosyste.ms: Awesome
An open API service indexing awesome lists of open source software.
https://github.com/plainerman/Latent-TPS
Source code of the paper: Transition Path Sampling with Boltzmann Generator-based MCMC Moves
https://github.com/plainerman/Latent-TPS
Last synced: 11 days ago
JSON representation
Source code of the paper: Transition Path Sampling with Boltzmann Generator-based MCMC Moves
- Host: GitHub
- URL: https://github.com/plainerman/Latent-TPS
- Owner: plainerman
- License: mit
- Created: 2023-11-21T10:32:56.000Z (12 months ago)
- Default Branch: main
- Last Pushed: 2024-05-07T12:14:39.000Z (6 months ago)
- Last Synced: 2024-08-01T16:45:44.774Z (3 months ago)
- Language: Jupyter Notebook
- Homepage:
- Size: 4.92 MB
- Stars: 2
- Watchers: 1
- Forks: 0
- Open Issues: 0
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
README
# Transition Path Sampling with Boltzmann Generator-based MCMC Moves
[![arXiv](https://img.shields.io/badge/arXiv-2312.05340-b31b1b.svg?style=flat-square)](https://arxiv.org/abs/2312.05340)
[![python](https://img.shields.io/badge/language-python%20-%2300599C.svg?style=flat-square)](https://github.com/plainerman/Latent-TPS)
[![License](https://img.shields.io/github/license/plainerman/Latent-TPS?style=flat-square)](LICENSE)In this repository we provide code to train a [Boltzmann generator](https://www.science.org/doi/full/10.1126/science.aaw1147) on alanine dipeptide and use it to perform Latent TPS (Transition Path Sampling).
## Setting up the conda environment
We provide a conda environment file for CUDA and CPU. You can create it by using one of the files (with or without the cpu flag).```bash
conda env create -f environment[-cpu].yml
```## Training the Boltzmann Generator
To train the Boltzmann Generator, you can use the `train.py` script. It has a number of options, but for ALDP you can train the model like this:```bash
python -m train --system AlanineDipeptideImplicit --data_save_frequency 120 --num_frames 1000000 --print_freq 250 --ckpt_freq 250 --val_freq 250 --flow_type internal_coords --batch_size 1024 --lr 5.e-4 --weight_decay 1.e-5 --lr_schedule cosine --warmup_dur 1000 --grad_clip 1000 --kl_loss_weight 1 --rkl_loss_weight 0 --hidden_dim 256 --update_layers 12 --run_name ALDP_RKL0_KL1_h256_u12_warmup_lrcosine_rerun
```If you are working with cuda, you can add the flags
```bash
--torch_device cuda --md_device CUDA
```## Run Latent TPS
You can find all the different options in `inference.py`.
You can change the states to find paths between by changing the `--start_state_idx` and `--end_state_idx` flags.Here is an example using the gaussian kernel, which adds random gaussian noise to the frames in latent space.
```bash
python -m inference --run_name mcmc_prob_langevin_40_noise0.05_seed0 --sampling_method mcmc --model_dir ./workdir/best --ckpt model_4250.ckpt --path_density langevin --noise_scale 0.05 --num_steps 40 --langevin_timestep 40 --num_paths 100 --seed 0
```# Acknowledgements
We thank the authors of [Flow Annealed Importance Sampling Bootstrap](https://github.com/lollcat/fab-torch) and [normflows](https://github.com/VincentStimper/normalizing-flows) which our flow training uses.