File size: 8,864 Bytes
96f168d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 | #!/bin/bash
cat << 'EOF'
____ _ _ _ _ _ _ ____ _ _____ ____ _ _ ____ _ _ ____ _ ____ _ _
____ __ ____ _ ____ _ __ __ _
/ ___| _ _ _ __ / _| _ \ ___ ___| | __ | __ ) ___| |_ __ _ \ \ / /__ _ __ ___(_) ___ _ __
\___ \| | | | '__| |_| | | |/ _ \ / __| |/ / | _ \ / _ \ __/ _` | \ \ / / _ \ '__/ __| |/ _ \| '_ \
___) | |_| | | | _| |_| | (_) | (__| < | |_) | __/ || (_| | \ V / __/ | \__ \ | (_) | | | |
|____/ \__,_|_| |_| |____/ \___/ \___|_|\_\ |____/ \___|\__\__,_| \_/ \___|_| |___/_|\___/|_| |_|
____ _ _ _ _ _ _ ____ _ _____ ____ _ _ ____ _ _ ____ _ ____ _ _
EOF
# This script is used to run SurfDock on test samples
source ~/miniforge3/bin/activate SurfDock
path=$(readlink -f "$0")
SurfDockdir="$(dirname "$(dirname "$(dirname "$(dirname "$path")")")")"
SurfDockdir=${SurfDockdir}
echo SurfDockdir : ${SurfDockdir}
temp="$(dirname "${SurfDockdir}")"
model_temp=${SurfDockdir}
#------------------------------------------------------------------------------------------------#
#------------------------------------ Step0 : Setup Params --------------------------------------#
#------------------------------------------------------------------------------------------------#
echo '------------------------------------ Step1 : Setup Params --------------------------------------'
export precomputed_arrays="${temp}/precomputed/precomputed_arrays"
gpu_string="0"
echo "Using GPU devices: ${gpu_string}"
IFS=',' read -ra gpu_array <<< "$gpu_string"
NUM_GPUS=${#gpu_array[@]}
export CUDA_VISIBLE_DEVICES=${gpu_string}
main_process_port=2951${gpu_array[-1]}
n_jobs=30
project_name='SurfDock_eval_samples_skip_target_processed'
## Please set the path to the input data
data_dir=${SurfDockdir}/model/data/Screen_sample_dirs/test_samples
esmbedding_dir=${temp}/Eval_result/processed_data/${project_name}/test_samples_esmbedding
# Set default value for target_have_processed if not already set
target_have_processed=${target_have_processed:-true}
## Please set the path to save the surface file and pocket file
surface_out_dir=${temp}/Eval_result/processed_data/${project_name}/test_samples_8A_surface
## Please set the path to the output csv file
out_csv_dir=${temp}/Eval_result/processed_data/${project_name}/input_csv_files/
out_csv_file=${out_csv_dir}/test_samples.csv
# project_name='SurfDock_Screen_samples/repeat5'
docking_out_dir=${temp}/Eval_result/docking_result/${project_name}
mkdir -p $surface_out_dir
if [ "$target_have_processed" = true ]; then
echo "Target structure has been processed, skipping this step."
else
echo "Processing target structure with OpenBabel..."
export BABEL_LIBDIR=~/miniforge3/envs/SurfDock/lib/openbabel/3.1.0
command=`
python ${SurfDockdir}/model/comp_surface/protein_process/openbabel_reduce_openbabel.py \
--data_path ${data_dir} \
--save_path ${surface_out_dir} \
--n_jobs ${n_jobs} \
`
state=$command
fi
#------------------------------------------------------------------------------------------------#
#----------------------------- Step1 : Compute Target Surface -----------------------------------#
#------------------------------------------------------------------------------------------------#
echo '----------------------------- Step1 : Compute Target Surface -----------------------------------'
mkdir -p $surface_out_dir
cd $surface_out_dir
command=`
python ${SurfDockdir}/model/comp_surface/prepare_target/computeTargetMesh_test_samples.py \
--data_dir ${data_dir} \
--out_dir ${surface_out_dir} \
--n_jobs ${n_jobs} \
`
state=$command
#------------------------------------------------------------------------------------------------#
#-------------------------------- Step2 : Get Input CSV File -----------------------------------#
#------------------------------------------------------------------------------------------------#
echo '-------------------------------- Step2 : Get Input CSV File -----------------------------------'
command=` python \
${SurfDockdir}/model/inference_utils/construct_csv_input.py \
--data_dir ${data_dir} \
--surface_out_dir ${surface_out_dir} \
--output_csv_file ${out_csv_file} \
`
state=$command
#------------------------------------------------------------------------------------------------#
#-------------------------------- Step3 : Get Pocket ESM Embedding ----------------------------#
#------------------------------------------------------------------------------------------------#
echo '-------------------------------- Step3 : Get Pocket ESM Embedding ----------------------------'
esm_dir=${SurfDockdir}/model/esm
sequence_out_file="${esmbedding_dir}/test_samples.fasta"
protein_pocket_csv=${out_csv_file}
full_protein_esm_embedding_dir="${esmbedding_dir}/esm_embedding_output"
pocket_emb_save_dir="${esmbedding_dir}/esm_embedding_pocket_output"
pocket_emb_save_to_single_file="${esmbedding_dir}/esm_embedding_pocket_output_for_train/esm2_3billion_pdbbind_embeddings.pt"
# get faste sequence
command=`python ${SurfDockdir}/model/datasets/esm_embedding_preparation.py \
--out_file ${sequence_out_file} \
--protein_ligand_csv ${protein_pocket_csv}`
state=$command
# esm embedding preprateion
command=`python ${esm_dir}/scripts/extract.py \
"esm2_t33_650M_UR50D" \
${sequence_out_file} \
${full_protein_esm_embedding_dir} \
--repr_layers 33 \
--include "per_tok" \
--truncation_seq_length 4096`
state=$command
# map pocket esm embedding
command=`python ${SurfDockdir}/model/datasets/get_pocket_embedding.py \
--protein_pocket_csv ${protein_pocket_csv} \
--embeddings_dir ${full_protein_esm_embedding_dir} \
--pocket_emb_save_dir ${pocket_emb_save_dir}`
state=$command
# save pocket esm embedding to single file
command=`python ${SurfDockdir}/model/datasets/esm_pocket_embeddings_to_pt.py \
--esm_embeddings_path ${pocket_emb_save_dir} \
--output_path ${pocket_emb_save_to_single_file}`
state=$command
#------------------------------------------------------------------------------------------------#
#------------------------ Step4 : Start Sampling Ligand Confromers ----------------------------#
#------------------------------------------------------------------------------------------------#
echo '------------------------ Step3 : Start Sampling Ligand Confromers ----------------------------'
diffusion_model_dir=${model_temp}/weight/docking
confidence_model_base_dir=${model_temp}/weight/posepredict
protein_embedding=${pocket_emb_save_to_single_file}
test_data_csv=${out_csv_file}
mdn_dist_threshold_test=3.0
version=6
dist_arrays=(3)
for i in ${dist_arrays[@]}
do
mdn_dist_threshold_test=${i}
command=`accelerate launch \
--multi_gpu \
--main_process_port ${main_process_port} \
--num_processes ${NUM_GPUS} \
${SurfDockdir}/scripts/inference_accelerate.py \
--data_csv ${test_data_csv} \
--model_dir ${diffusion_model_dir} \
--ckpt best_ema_inference_epoch_model.pt \
--confidence_model_dir ${confidence_model_base_dir} \
--confidence_ckpt best_model.pt \
--save_docking_result \
--mdn_dist_threshold_test ${mdn_dist_threshold_test} \
--esm_embeddings_path ${protein_embedding} \
--run_name ${confidence_model_base_dir}_test_dist_${mdn_dist_threshold_test} \
--project ${project_name} \
--out_dir ${docking_out_dir} \
--batch_size 40 \
--batch_size_molecule 1 \
--samples_per_complex 40 \
--save_docking_result_number 40 \
--head_index 0 \
--tail_index 10000 \
--inference_mode evaluate \
--wandb_dir ${temp}/docking_result/test_workdir`
state=$command
done
cat << 'EOF'
____ _ _ _ _ _ _ ____ _ _____ ____ _ _ ____ _ _ ____ _ ____ _ _
____ __ ____ _ ____ _ _ ____ _
/ ___| _ _ _ __ / _| _ \ ___ ___| | __ / ___| __ _ _ __ ___ _ __ | (_)_ __ __ _ | _ \ ___ _ __ ___| |
\___ \| | | | '__| |_| | | |/ _ \ / __| |/ / \___ \ / _` | '_ ` _ \| '_ \| | | '_ \ / _` | | | | |/ _ \| '_ \ / _ \ |
___) | |_| | | | _| |_| | (_) | (__| < ___) | (_| | | | | | | |_) | | | | | | (_| | | |_| | (_) | | | | __/_|
|____/ \__,_|_| |_| |____/ \___/ \___|_|\_\ |____/ \__,_|_| |_| |_| .__/|_|_|_| |_|\__, | |____/ \___/|_| |_|\___(_)
|_| |___/
____ _ _ _ _ _ _ ____ _ _____ ____ _ _ ____ _ _ ____ _ ____ _ _
EOF
|