File size: 4,883 Bytes
c2767f4 | 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 | #!/bin/bash
set -euo pipefail
export LD_LIBRARY_PATH="$CONDA_PREFIX/lib/:$LD_LIBRARY_PATH"
export LD_LIBRARY_PATH="$CONDA_PREFIX/lib/python3.11/site-packages/fastpt/torch/lib:$LD_LIBRARY_PATH"
export LD_LIBRARY_PATH=${ROCM_PATH}/opencl/lib:$LD_LIBRARY_PATH
SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
EXAMPLE_DIR=$(cd "${SCRIPT_DIR}/.." && pwd)
REPO_ROOT=$(cd "${SCRIPT_DIR}/../../../.." && pwd)
# source "${REPO_ROOT}/env.sh"
if [[ -n "${ROCM_PATH:-}" && -f "${ROCM_PATH}/cuda/env.sh" ]]; then
source "${ROCM_PATH}/cuda/env.sh"
fi
export PYTHONPATH="${REPO_ROOT}/src:${REPO_ROOT}:${PYTHONPATH:-}"
export HIP_VISIBLE_DEVICES="${HIP_VISIBLE_DEVICES:-0}"
export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-${HIP_VISIBLE_DEVICES}}"
export OMP_NUM_THREADS="${OMP_NUM_THREADS:-4}"
export DIFFDOCK_RADIUS_ON_CPU="${RADIUS_ON_CPU:-false}"
DIFFDOCK_DATA_ROOT="${DIFFDOCK_DATA_ROOT:-${ONESCIENCE_DATASETS_DIR}/diffdock}"
export TORCH_HOME="${TORCH_HOME:-${DIFFDOCK_DATA_ROOT}/torch_home}"
SCORE_MODEL_DIR="${SCORE_MODEL_DIR:-${DIFFDOCK_DATA_ROOT}/score_model}"
SCORE_CKPT="${SCORE_CKPT:-best_ema_inference_epoch_model.pt}"
CONFIDENCE_MODEL_DIR="${CONFIDENCE_MODEL_DIR:-${DIFFDOCK_DATA_ROOT}/confidence_model}"
CONFIDENCE_CKPT="${CONFIDENCE_CKPT:-best_model_epoch75.pt}"
ENABLE_CONFIDENCE="${ENABLE_CONFIDENCE:-true}"
OLD_CONFIDENCE_MODEL="${OLD_CONFIDENCE_MODEL:-true}"
OUT_DIR="${OUT_DIR:-${EXAMPLE_DIR}/outputs/scnet_inference}"
CONFIG_PATH="${CONFIG_PATH:-${OUT_DIR}/inference_config.yml}"
mkdir -p "${OUT_DIR}"
OUT_DIR=$(cd "${OUT_DIR}" && pwd)
CONFIG_DIR=$(dirname "${CONFIG_PATH}")
CONFIG_NAME=$(basename "${CONFIG_PATH}")
mkdir -p "${CONFIG_DIR}"
CONFIG_PATH=$(cd "${CONFIG_DIR}" && pwd)/"${CONFIG_NAME}"
DEFAULT_SHARED_CSV="${DIFFDOCK_DATA_ROOT}/datasets/inferdata/protein_ligand_example.csv"
PROTEIN_LIGAND_CSV="${PROTEIN_LIGAND_CSV:-}"
if [[ -z "${PROTEIN_LIGAND_CSV}" && -f "${DEFAULT_SHARED_CSV}" ]]; then
PROTEIN_LIGAND_CSV="${DEFAULT_SHARED_CSV}"
fi
COMPLEX_NAME="${COMPLEX_NAME:-6o5u_test}"
PROTEIN_PATH="${PROTEIN_PATH:-${EXAMPLE_DIR}/data/6o5u_protein_processed.pdb}"
PROTEIN_SEQUENCE="${PROTEIN_SEQUENCE:-}"
LIGAND_DESCRIPTION="${LIGAND_DESCRIPTION:-${EXAMPLE_DIR}/data/6o5u_ligand.sdf}"
DEVICE="${DEVICE:-auto}"
SAMPLES_PER_COMPLEX="${SAMPLES_PER_COMPLEX:-10}"
BATCH_SIZE="${BATCH_SIZE:-10}"
INFERENCE_STEPS="${INFERENCE_STEPS:-20}"
ACTUAL_STEPS="${ACTUAL_STEPS:-}"
NO_RANDOM="${NO_RANDOM:-false}"
NO_FINAL_STEP_NOISE="${NO_FINAL_STEP_NOISE:-true}"
CROP_BEYOND="${CROP_BEYOND:-}"
yaml_value() {
if [[ -z "${1:-}" || "${1}" == "null" ]]; then
printf "null"
else
local value
value=$(printf "%s" "$1" | sed "s/'/''/g")
printf "'%s'" "$value"
fi
}
yaml_bool() {
if [[ "${1,,}" == "true" ]]; then
printf "true"
else
printf "false"
fi
}
if [[ "${ENABLE_CONFIDENCE,,}" == "true" ]]; then
CONFIDENCE_MODEL_VALUE=$(yaml_value "${CONFIDENCE_MODEL_DIR}")
else
CONFIDENCE_MODEL_VALUE="null"
fi
cat > "${CONFIG_PATH}" <<EOF
runtime:
device: $(yaml_value "${DEVICE}")
loglevel: INFO
out_dir: $(yaml_value "${OUT_DIR}")
model:
model_dir: $(yaml_value "${SCORE_MODEL_DIR}")
ckpt: $(yaml_value "${SCORE_CKPT}")
old_score_model: false
confidence:
confidence_model_dir: ${CONFIDENCE_MODEL_VALUE}
confidence_ckpt: $(yaml_value "${CONFIDENCE_CKPT}")
old_confidence_model: $(yaml_bool "${OLD_CONFIDENCE_MODEL}")
input:
protein_ligand_csv: $(yaml_value "${PROTEIN_LIGAND_CSV}")
complex_name: $(yaml_value "${COMPLEX_NAME}")
protein_path: $(yaml_value "${PROTEIN_PATH}")
protein_sequence: $(yaml_value "${PROTEIN_SEQUENCE}")
ligand_description: $(yaml_value "${LIGAND_DESCRIPTION}")
lm_embeddings: null
crop_beyond: $(yaml_value "${CROP_BEYOND}")
sampling:
samples_per_complex: ${SAMPLES_PER_COMPLEX}
batch_size: ${BATCH_SIZE}
inference_steps: ${INFERENCE_STEPS}
actual_steps: $(yaml_value "${ACTUAL_STEPS}")
sigma_schedule: expbeta
inf_sched_alpha: 1.0
inf_sched_beta: 1.0
no_random: $(yaml_bool "${NO_RANDOM}")
no_final_step_noise: $(yaml_bool "${NO_FINAL_STEP_NOISE}")
ode: false
choose_residue: false
initial_noise_std_proportion: 1.0
temp_sampling_tr: 1.0
temp_psi_tr: 0.0
temp_sigma_data_tr: 0.5
temp_sampling_rot: 1.0
temp_psi_rot: 0.0
temp_sigma_data_rot: 0.5
temp_sampling_tor: 1.0
temp_psi_tor: 0.0
temp_sigma_data_tor: 0.5
EOF
echo "DiffDock inference config: ${CONFIG_PATH}"
echo "Output directory: ${OUT_DIR}"
echo "Score model: ${SCORE_MODEL_DIR}/${SCORE_CKPT}"
echo "Confidence rerank: ${ENABLE_CONFIDENCE}"
echo "TORCH_HOME: ${TORCH_HOME}"
echo "Radius on CPU: ${DIFFDOCK_RADIUS_ON_CPU}"
if [[ -n "${PROTEIN_LIGAND_CSV}" ]]; then
echo "Input CSV: ${PROTEIN_LIGAND_CSV}"
else
echo "Single input: ${PROTEIN_PATH} + ${LIGAND_DESCRIPTION}"
fi
cd "${REPO_ROOT}"
python "${SCRIPT_DIR}/sample_diffdock.py" --config "${CONFIG_PATH}"
|