version 2

README.md

@@ -1,5 +1,7 @@
 # P2DFlow
 
+> ## ℹ️ The version 2 code of P2DFlow has been released!
+
 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:
@@ -37,43 +39,27 @@ python ./dataset/download.py
 ```
 We will use `.pdb` and `.xtc` files for the following calculation.
 
-#### 2. Calculate the 'approximate energy and select representative structures
-(i) Use `gaussian_kde` to calculate the 'approximate energy' (You need to put all files above in `./dataset`, include `ATLAS_filename.txt` for filenames of all proteins):
+#### 2. Calculate the 'approximate energy' and select representative structures
+(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](https://drive.google.com/drive/folders/11mdVfMi2rpVn7nNG2mQAGA5sNXCKePZj?usp=sharing)):
 ```
-python ./dataset/traj_analyse.py
+python ./dataset/traj_analyse_select.py
 ```
-And you will get `traj_info.csv`.
+And you will get selected representative structures in `select` dir and `traj_info_select.csv` for 'approximate energy'.
 
-(ii) Select representative structures at equal intervals based on the 'approximate energy':
-```
-python ./dataset/md_select.py
-```
 
 #### 3. Process selected dataset
 
-(i) Download the selected dataset (or get it from the two steps above) from [Google Drive](https://drive.google.com/drive/folders/11mdVfMi2rpVn7nNG2mQAGA5sNXCKePZj?usp=sharing) whose filename is `selected_dataset.tar`, and decompress it using:
+(i) Download the selected dataset (or get it from the two steps above) from [Google Drive](https://drive.google.com/drive/folders/11mdVfMi2rpVn7nNG2mQAGA5sNXCKePZj?usp=sharing) 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 -xvf select_dataset.tar
+tar -xzvf select_dataset_v1.tar
 ```
-(ii) Preprocess `.pdb` files to get `.pkl` files:
-```
-python ./data/process_pdb_files.py --pdb_dir ${pdb_dir} --write_dir ${write_dir}
-```
-And you will get `metadata.csv`.
 
-then compute node representation and pair representation using ESM-2 (`csv_path` is the path of `metadata.csv`):
-```
-python ./data/cal_repr.py --csv_path ${csv_path}
+(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:
 ```
-then compute predicted static structure using ESMFold (`csv_path` is the path of `metadata.csv`):
+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}
 ```
-python ./data/cal_static_structure.py --csv_path ${csv_path}
-```
-(iii) Provide the necessary `.csv` files for training
-
-If you are using the data we have preprocessed, download the `.csv` files from [Google Drive](https://drive.google.com/drive/folders/11mdVfMi2rpVn7nNG2mQAGA5sNXCKePZj?usp=sharing) whose filenames are `train_dataset.csv` and `train_dataset_energy.csv`(they correspond to `csv_path` and `energy_csv_path` in `./configs/base.yaml` during training).
-
-Or if you are using your own data, you can get `metadata.csv` from step 3 (correspond to `csv_path` in `./configs/base.yaml` during training, and you need to split a subset from it as the train dataset), and get `traj_info.csv` from step 2 (correspond to `energy_csv_path`).
+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](https://drive.google.com/drive/folders/11mdVfMi2rpVn7nNG2mQAGA5sNXCKePZj?usp=sharing)
 
 
 
@@ -82,7 +68,7 @@ Download the pretrained checkpoint from [Google Drive](https://drive.google.com/
 
 
 ## Training
-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` and `energy_csv_path`). After this, you can run:
+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
 ```
@@ -90,7 +76,7 @@ And you will get the checkpoints in `./ckpt`.
 
 
 ## Inference
-To infer for specified protein sequence, firstly modify `./inference/valid_seq.csv` and `./configs/inference.yaml` (especially for `validset_path`), then run:
+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
 ```
@@ -98,18 +84,15 @@ And you will get the results in `./inference_outputs/weights/`.
 
 
 ## Evaluation
-To evaluate metrics related to fidelity and dynamics, specify paths in `./analysis/eval_test.py`, then run:
-```
-python ./analysis/eval_test.py
-```
-To evaluate PCA, specify paths in `./analysis/pca_analyse.py`, then run:
+To evaluate metrics related to validity, fidelity and dynamics, run:
 ```
-python ./analysis/pca_analyse.py
+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 draw the ramachandran plots, specify paths in `./analysis/Ramachandran_plot.py`, then run:
+To evaluate PCA, run:
 ```
-python ./analysis/Ramachandran_plot.py
+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](https://drive.google.com/drive/folders/11mdVfMi2rpVn7nNG2mQAGA5sNXCKePZj?usp=sharing)
 
 ## License
 This project is licensed under the terms of the GPL-3.0 license.