P2DFlow is a protein ensemble generative model with SE(3) flow matching based on ESMFold, the ensembles generated by P2DFlow could aid in understanding protein functions across various scenarios.
Technical details and evaluation results are provided in our paper ( 🎉 It has been accepted by JCTC ):
The codes have also been uploaded to Huggingface !
In an environment with cuda 11.7, run:
conda env create -f environment.yml
To activate the environment, run:
conda activate P2DFlow
(tips: If you want to use the data we have preprocessed, please go directly to 3. Process selected dataset; if you prefer to process the data from scratch or work with your own data, please start from the beginning)
(i) Download the Analysis & MDs dataset from ATLAS, or you can use ./dataset/download.py by running:
python ./dataset/download.py
We will use .pdb and .xtc files for the following calculation.
(i) Use gaussian_kde to calculate the 'approximate energy' (You need to put all files above in ./dataset, just like ATLAS_init_example in Google Drive):
python ./dataset/traj_analyse_select.py
And you will get selected representative structures in select dir and traj_info_select.csv for 'approximate energy'.
(i) Download the selected dataset (or get it from the two steps above) from Google Drive whose filename is selected_dataset_v1.tar or selected_dataset_v2.tar ('v1' selects ~10 structures from MD, 'v2' selects ~100 structures from MD), and decompress it using:
tar -xzvf select_dataset_v1.tar
(ii) Preprocess .pdb files to get .pkl files, compute node representation and pair representation using ESM-2, predict static structure using ESMFold, and get merged .csv file:
python ./data/process_pdb_files.py --pdb_dir ${pdb_dir} --write_dir ${write_dir} --traj_info_file ${traj_info_file} --valid_seq_file ${valid_seq_file} --merged_output_file ${merged_output_file}
And you will get .pkl files (large file size) and metadata_merged.csv. (if you are using your own data, you need to split dataset to get validation set as ${valid_seq_file} first, an example is ./inference/valid_seq.csv).
Processed data will be similar to ATLAS_processed_example.tar.gz in Google Drive
Download the pretrained checkpoint from Google Drive whose filename is pretrained.ckpt, and put it into ./weights folder. You can use the pretrained weight for inference.
To train P2DFlow, firstly make sure you have prepared the dataset according to Prepare Dataset, and put it in the right folder, then modify ./configs/base.yaml (especially for csv_path). After this, you can run:
python experiments/train_se3_flows.py
And you will get the checkpoints in ./ckpt.
To infer for specified protein sequence, firstly modify ./configs/inference.yaml (especially for ckpt_path and validset_path), then run:
python experiments/inference_se3_flows.py
And you will get the results in ./inference_outputs/weights/.
To evaluate metrics related to validity, fidelity and dynamics, run:
python ./analysis/eval_result.py --pred_org_dir ${pred_org_dir} --valid_csv_file ${valid_csv_file} --pred_merge_dir ${pred_merge_dir} --target_dir ${target_dir} --crystal_dir ${crystal_dir}
To evaluate PCA, run:
python ./analysis/pca_analyse.py --pred_pdb_dir ${pred_pdb_dir} --target_dir ${target_dir} --crystal_dir ${crystal_dir}
Evaluation results will be similar to evaluation_example in Google Drive
This project is licensed under the terms of the GPL-3.0 license.
@article{jin2025p2dflow,
title={P2DFlow: A Protein Ensemble Generative Model with SE(3) Flow Matching},
author={Yaowei Jin, Qi Huang, Ziyang Song, Mingyue Zheng, Dan Teng, Qian Shi},
journal={Journal of Chemical Theory and Computation},
year={2025}
}