The repository is the official implementation of the ICML24 paper Protein Conformation Generation via Force-Guided SE(3) Diffusion Models, which introduces ConfDiff, a force-guided SE(3) diffusion model for protein conformation generation. ConfDiff can generate protein conformations with rich diversity while preserving high fidelity. Physics-based energy and force guidance strategies effectively guide the diffusion sampler to generate low-energy conformations that better align with the underlying Boltzmann distribution.

With recent progress in protein conformation prediction, we extend ConfDiff to ConfDiff-FullAtom, diffusion models for full-atom protein conformation prediction. Current models include the following updates:

• Integrated a regression module to predict atomic coordinates for side-chain heavy atoms
• Provided models options with four folding model representations (ESMFold or OpenFold with recycling number of 0 and 3)
• Used all feature outputs from the folding model (node + edge) for diffusion model training.
• Released a version of sequence-conditional models fine-tuned on the Atlas MD dataset.

# clone project
git clone https://url/to/this/repo/ConfDiff.git
cd ConfDiff

# create conda virtual environment
conda env create -f env.yml
conda activate confdiff

# install openfold
git clone https://github.com/aqlaboratory/openfold.git
pip install -e openfold

We precompute ESMFold and OpenFold representations as inputs to the model. The detailed generation pipline can be referenced in the README of the pretrained_repr/ folder.

ConfDiff-BASE employs a sequence-based conditional score network to guide an unconditional score model using classifier-free guidance, enabling diverse conformation sampling while ensuring structural fidelity to the input sequence.

We train Confdiff-BASE using the protein structures from Protein Data Bank, and evaluate on various datasets including fast-folding, bpti, apo-holo and atlas. Details on dataset prepration and evaluation can be found in the dataset folder.

The following datasets and pre-computed representations are required to train Confdiff-BASE:

1. RCSB PDB dataset: See dataset/rcsb for details. Once prepared, specify the csv_path and pdb_dir in the configuration file configs/paths/default.yaml.
2. ESMFold or OpenFold representations: See pretrained_repr for details. Once prepared, specify the data_root of esmfold_repr/openfold_repr in the configuration file configs/paths/default.yaml.

ConfDiff-BASE consists of a sequence-conditional model and an unconditional model.

To train the conditional model:

python3 src/train.py \
        task_name=cond \
        experiment=full_atom \
        data/dataset=rcsb \
        data/repr_loader=openfold \
        data.repr_loader.num_recycles=3 \
        data.train_batch_size=4

The detailed training configuration can be found in configs/experiment/full_atom.yaml.

To train the unconditonal model:

python3 src/train.py \
        task_name=uncond \
        experiment=uncond_model \
        data/dataset=rcsb \
        data.train_batch_size=4 

The detailed training configuration can be found in configs/experiment/uncond.yaml.

Access pretrained models with different pretrained representations:

To sample conformations using the ConfDiff-BASE model:

#Please note that the model and representation need to be compatible.
python3 src/eval.py \
        task_name=eval_base_bpti \
        experiment=clsfree_guide \
        data/repr_loader=openfold \
        data.repr_loader.num_recycles=3 \
        paths.guidance.cond_ckpt=/path/to/your/cond_model \
        paths.guidance.uncond_ckpt=/path/to/your/uncond_model \
        data.dataset.test_gen_dataset.csv_path=/path/to/your/testset_csv \
        data.dataset.test_gen_dataset.num_samples=1000 \
        data.gen_batch_size=20 \
        model.score_network.cfg.clsfree_guidance_strength=0.8

By utilizing prior information from the MD force field, our model effectively reweights the generated conformations to ensure they better adhere to the equilibrium distribution

Protein conformations with force or energy labels are required to train the corresponding ConfDiff-FORCE or ConfDiff-ENERGY. We use OpenMM for energy and force evaluation of the conformation samples generated by ConfDiff-BASE

To evaluate force and energy labels using OpenMM and prepare the training data:

python3 src/utils/protein/openmm_energy.py \
                --input-root /path/to/your/generated/samples \
                --output-root /path/to/your/output_dir \

The output directory /path/to/your/output_dir contains force annotation files (with the suffix *force.npy), optimized energy PDB files, and train and validation CSV files with energy labels.

Before training, please ensure that the pretrained representations for the training proteins have been prepared.

To train ConfDiff-FORCE:

# case for training ConfDiff-FORCE
python3 src/train.py \
        experiment=force_guide \
        data/repr_loader=esmfold \
        data.repr_loader.num_recycles=3 \
        paths.guidance.cond_ckpt=/path/to/your/cond_model \
        paths.guidance.uncond_ckpt=/path/to/your/uncond_model \
        paths.guidance.train_csv=/path/to/your/output_dir/train.csv \
        paths.guidance.val_csv=/path/to/your/output_dir/val.csv \
        paths.guidance.pdb_dir=/path/to/your/output_dir/ \
        data.train_batch_size=4 

Similarly, the ConfDiff-ENERGY model can be trained by setting experiment=energy_guide. Detailed training configurations can be found in the file configs/experiment/force_guide(energy_guide).yaml.

Access pretrained ConfDiff-FORCE/ENERGY with ESMFold representations on different datasets.

We found that using only the node representation's on the fast-folding dataset yields better results. To train models with only node representation, set data.repr_loader.edge_size=0

To sample conformations using the ConfDiff-FORCE/ENERGY model:

# case for generating samples by ConfDiff-FORCE
python3 src/eval.py \
        task_name=eval_force \
        experiment=force_guide \
        data/repr_loader=esmfold \
        data.repr_loader.num_recycles=3 \
        ckpt_path=/path/to/your/model/ckpt/ \
        data.dataset.test_gen_dataset.csv_path=/path/to/your/testset_csv \
        data.dataset.test_gen_dataset.num_samples=1000 \
        data.gen_batch_size=20 \
        model.score_network.cfg.clsfree_guidance_strength=0.8 \
        model.score_network.cfg.force_guidance_strength=1.0
        # data.repr_loader.edge_size=0 for pretrained checkpoints on fast-folding

See datasets/atlas for ATLAS data preparation.

To fine-tune ConfDiff-BASE on ATLAS.

#Training
python3 src/train.py \
        task_name=finetune_atlas \
        experiment=full_atom \
        data/dataset=atlas \
        data/repr_loader=esmfold \
        ckpt_path=/path/to/your/cond_model

#Evaluation
python3 src/eval.py \
        task_name=eval_atlas \
        experiment=full_atom \
        data/dataset=atlas \
        data.dataset.test_gen_dataset.csv_path=/path/to/your/testset_csv \
        data.dataset.test_gen_dataset.num_samples=1000 \
        data.gen_batch_size=20 \
        ckpt_path=/path/to/your/atlas_model

The models fine-tuned on ATLAS MD dataset is shown in the table below:

We benchmark model performance on following datasets: BPTI, fast-folding, Apo-holo, and Atlas. Evaluation details can be found in the datasets folder and the notebook notebooks/analysis.ipynb.

ConfDiff-XXX-ClsFree refers to the ConfDiff-BASE model utilizing classifier-free guidance sampling with the ConfDiff-XXX-COND and ConfDiff-UNCOND models. As described in the paper, all results are based on ensemble sampling with varying levels of classifier-guidance strength. For the fast-folding dataset, the classifier-guidance strength values range from 0.5 to 1.0, while for other datasets, the range is 0.8 to 1.0. For BPTI and fast-folding, we also provide results from the pretrained FORCE and ENERGY models.

The guidance strength is set to 1.5 for the FORCE model and 1.0 for the ENERGY model.

The models here utilize only the pretrained node representation. The force guidance strength is set to 2.0 for the FORCE model and 1.0 for the ENERGY model.

@inproceedings{wang2024proteinconfdiff,
  title={Protein Conformation Generation via Force-Guided SE (3) Diffusion Models},
  author={Wang, Yan and Wang, Lihao and Shen, Yuning and Wang, Yiqun and Yuan, Huizhuo and Wu, Yue and Gu, Quanquan},
  booktitle={Forty-first International Conference on Machine Learning},
  year={2024}
}