File size: 10,939 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 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 | #!/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}
#------------------------------------------------------------------------------------------------#
#------------------------------------ Step1 : Setup Params --------------------------------------#
#------------------------------------------------------------------------------------------------#
export precomputed_arrays="${temp}/precomputed/precomputed_arrays"
## Please set the GPU devices you want to use
gpu_string="7"
echo "Using GPU devices: ${gpu_string}"
IFS=',' read -ra gpu_array <<< "$gpu_string"
NUM_GPUS=${#gpu_array[@]}
export CUDA_VISIBLE_DEVICES=${gpu_string}
## Please set the main Parameters
main_process_port=2957${gpu_array[-1]}
## Please set the project name
project_name='SurfDock_Screen_samples_skip_target_processed'
# /home/caoduanhua/NM_submit_code/SurfDock
# 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}/Screen_result/processed_data/${project_name}/test_samples_8A_surface
## Please set the path to the input data
data_dir=${SurfDockdir}/model/data/Screen_sample_dirs/test_samples
## Please set the path to the output csv file
out_csv_dir=${temp}/Screen_result/processed_data/${project_name}/input_csv_files/
out_csv_file=${out_csv_dir}/test_samples.csv
## Please set the path to the esmbedding file
esmbedding_dir=${temp}/Screen_result/processed_data/${project_name}/test_samples_esmbedding
## Please set the path to the Screen ligand library file
Screen_lib_path=${SurfDockdir}/model/data/Screen_sample_dirs/test_samples/1a0q/1a0q_ligand_for_Screen.sdf
## Please set the path to the docking result directory
docking_out_dir=${temp}/Screen_result/docking_result/${project_name}
#------------------------------------------------------------------------------------------------#
# -----------------------Step1 : Processed Target Structure -------------------------------------#
#----------------(Set target_have_processed as true if you have done with your pipeline)---------#
#------------------------------------------------------------------------------------------------#
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}`
state=$command
fi
#------------------------------------------------------------------------------------------------#
#----------------------------- Step2 : Compute Target Surface -----------------------------------#
#------------------------------------------------------------------------------------------------#
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} \
`
state=$command
#------------------------------------------------------------------------------------------------#
#-------------------------------- Step3 : 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} \
--Screen_ligand_library_file ${Screen_lib_path} \
`
state=$command
#------------------------------------------------------------------------------------------------#
#-------------------------------- Step4 : 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
#------------------------------------------------------------------------------------------------#
#------------------------ Step5 : 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}
cd ${SurfDockdir}/scripts/bash_scripts/test_scripts
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 400 \
--batch_size_molecule 10 \
--samples_per_complex 40 \
--save_docking_result_number 40 \
--head_index 0 \
--tail_index 10000 \
--inference_mode Screen \
--wandb_dir ${temp}/docking_result/test_workdir`
state=$command
done
#------------------------------------------------------------------------------------------------#
#---------------- Step6 : Start Rescoring the Pose For Screening -----------------#
#------------------------------------------------------------------------------------------------#
echo '---------------- Step4 : Start Rescoring the Pose For Screening -----------------'
# project_name='SurfDock_Screen_samples/repeat_zero'
# surface_out_dir=${SurfDockdir}/model/data/Screen_sample_dirs/${project_name}/test_samples_8A_surface
# data_dir=${SurfDockdir}/model/data/Screen_sample_dirs/test_samples
out_csv_file=${out_csv_dir}/score_inplace.csv
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} \
--Screen_ligand_library_file ${Screen_lib_path} \
--is_docking_result_dir \
--docking_result_dir ${docking_out_dir} \
`
state=$command
confidence_model_base_dir=${model_temp}/weight/screen
test_data_csv=${out_csv_file}
version=6
dist_arrays=(3)
for i in ${dist_arrays[@]}
do
mdn_dist_threshold_test=${i}
echo mdn_dist_threshold_test : ${mdn_dist_threshold_test}
command=`accelerate launch \
--multi_gpu \
--main_process_port ${main_process_port} \
--num_processes 1 \
${SurfDockdir}/scripts/evaluate_score_in_place.py \
--data_csv ${test_data_csv} \
--confidence_model_dir ${confidence_model_base_dir} \
--confidence_ckpt best_model.pt \
--model_version version6 \
--mdn_dist_threshold_test ${mdn_dist_threshold_test} \
--esm_embeddings_path ${protein_embedding} \
--run_name ${project_name}_test_dist_${mdn_dist_threshold_test} \
--project ${project_name} \
--out_dir ${docking_out_dir} \
--batch_size 40 \
--wandb_dir ${temp}/wandb/test_workdir`
state=$command
done
cat << 'EOF'
____ _ _ _ _ _ _ ____ _ _____ ____ _ _ ____ _ _ ____ _ ____ _ _
____ __ ____ _ ____ _ _ ____ _
/ ___| _ _ _ __ / _| _ \ ___ ___| | __ / ___| __ _ _ __ ___ _ __ | (_)_ __ __ _ | _ \ ___ _ __ ___| |
\___ \| | | | '__| |_| | | |/ _ \ / __| |/ / \___ \ / _` | '_ ` _ \| '_ \| | | '_ \ / _` | | | | |/ _ \| '_ \ / _ \ |
___) | |_| | | | _| |_| | (_) | (__| < ___) | (_| | | | | | | |_) | | | | | | (_| | | |_| | (_) | | | | __/_|
|____/ \__,_|_| |_| |____/ \___/ \___|_|\_\ |____/ \__,_|_| |_| |_| .__/|_|_|_| |_|\__, | |____/ \___/|_| |_|\___(_)
|_| |___/
____ _ _ _ _ _ _ ____ _ _____ ____ _ _ ____ _ _ ____ _ ____ _ _
EOF
|