🍅

DiffDock実行コンテナの作成 | Singularity (on WSL2 & with GPU)

2022/10/07に公開約8,000字

コードブロック における 🍅 は、プロンプト を表している
$SINGULARITY_IMAGE~/apps/singularity_image です (私の環境)

はじめに

前々回・前回に引き続き WSL2 にインストールした Singularity で遊んでみる

前々回:Singularity (on WSL2 & with GPU) はじめました! | 環境構築とOmegaFold実行コンテナ作成
前回:Singularity (on WSL2 & with GPU) はじめました! | GNINA実行コンテナ作成

今回は DiffDock という、わりと最近報告された分子ドッキングプログラムを実行するコンテナイメージ ( .sif ) を作成し、簡単な動作テストを実施する

Predicting the binding structure of a small molecule ligand to a protein -- a task known as molecular docking -- is critical to drug design. Recent deep learning methods that treat docking as a regression problem have decreased runtime compared to traditional search-based methods but have yet to offer substantial improvements in accuracy. We instead frame molecular docking as a generative modeling problem and develop DiffDock, a diffusion generative model over the non-Euclidean manifold of ligand poses. To do so, we map this manifold to the product space of the degrees of freedom (translational, rotational, and torsional) involved in docking and develop an efficient diffusion process on this space. Empirically, DiffDock obtains a 38% top-1 success rate (RMSD<2 Å) on PDBBind, significantly outperforming the previous state-of-the-art of traditional docking (23%) and deep learning (20%) methods. Moreover, DiffDock has fast inference times and provides confidence estimates with high selective accuracy.

参照:arXiv | DiffDock: Diffusion Steps, Twists, and Turns for Molecular Docking
参照:Twitter | @HannesStaerkによる紹介

以下、本記事を執筆した際の環境。Singularity (on WSL2 & with GPU) 環境の構築については前々回の記事を参照

specifications
CPU Ryzen 7 3700X (AMD)
グラフィックボード GeForce RTX 2080 Ti VENTUS GP OC (MSI)
メモリ DDR4-3200 16 GB×2 (Crucial)
version
Windows 10 21H2
NVIDIA グラフィクス ドライバー 516.94
WSL2 5.10.102.1
Ubuntu on WLS2 20.04.5

DiffDockを実行するコンテナイメージの作成

.defファイルの作成

DiffDock を実行するコンテナイメージ ( .sif ) をビルドするための .def ファイルを下記のように作成する。

現状 DiffDock の実行には ESM が必要らしい。

We will soon also provide weights of a trained model without ESM2 embeddings such that this step is not necessary.

参照:GitHub | DiffDock: Diffusion Steps, Twists, and Turns for Molecular Docking
参照:GitHub | Evolutionary Scale Modeling

ESM の初回実行時にはモデルパラメーターがダウンロードされるが、コンテナイメージ作成後にコンテナイメージ内にダウンロードするのは面倒なので、コンテナイメージをビルドする際にあらかじめダウンロードしておく

$SINGULARITY_IMAGE/def_file/DiffDock.def
BootStrap: docker
From: nvidia/cuda:11.6.2-cudnn8-devel-ubuntu20.04

%post
  # localtime関係のWARNINGに対処
  touch /etc/localtime
  
  # DiffDockに必要なパッケージのインストール (冗長)
  apt update && apt upgrade -y
  apt install -y build-essential python3 python3-venv git curl
  
  # localeにen_US.UTF-8を追加する (WSL2だけで使用するなら不要)
  apt install -y locales
  locale-gen en_US.UTF-8

  # GitHubから/usr/local/appsにDiffDockを複製する
  mkdir -p /usr/local/apps
  cd /usr/local/apps
  git clone https://github.com/gcorso/DiffDock.git
  DIFFDOCK_HOME="/usr/local/apps/DiffDock"
  
  # ${DIFFDOCK_HOME}にPython仮想環境.venvを作成しアクティベート
  cd ${DIFFDOCK_HOME}
  python3 -m venv .venv
  . ./.venv/bin/activate
  
  # Python仮想環境内にpipでDiffDockの依存パッケージをインストール
  python3 -m pip install --upgrade pip
  python3 -m pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu116
  python3 -m pip install torch-scatter -f https://data.pyg.org/whl/torch-1.12.1+cu116.html
  python3 -m pip install torch-sparse -f https://data.pyg.org/whl/torch-1.12.1+cu116.html
  python3 -m pip install torch-geometric
  python3 -m pip install torch-cluster -f https://data.pyg.org/whl/torch-1.12.1+cu116.html
  python3 -m pip install torch-spline-conv -f https://data.pyg.org/whl/torch-1.12.1+cu116.html
  python3 -m pip install pyyaml scipy networkx biopython rdkit-pypi e3nn spyrmsd pandas biopandas
  
  # GitHubから/usr/local/apps/DiffDockにesmを複製しインストール
  cd ${DIFFDOCK_HOME}
  git clone https://github.com/facebookresearch/esm
  cd ${DIFFDOCK_HOME}/esm
  python3 -m pip install -e .
  
  # ESMで使用するモデルをあらかじめダウンロード
  mkdir -p ${DIFFDOCK_HOME}/esm/model_weights/.cache/torch/hub/checkpoints
  cd ${DIFFDOCK_HOME}/esm/model_weights/.cache/torch/hub/checkpoints
  curl -L -O https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t33_650M_UR50D.pt
  curl -L -O https://dl.fbaipublicfiles.com/fair-esm/regression/esm2_t33_650M_UR50D-contact-regression.pt

%environment
  # DiffDockおよびESM実行のための環境設定
  export DIFFDOCK_HOME="/usr/local/apps/DiffDock"
  . ${DIFFDOCK_HOME}/.venv/bin/activate
  export PYTHONPATH="${DIFFDOCK_HOME}"
  export PYTHONPATH="$PYTHONPATH:${DIFFDOCK_HOME}/esm"

%runscript
  # コンテナ実行時のスクリプトを定義
  # "$@"は追加の引数を利用するための設定
  python3 -m inference \
  --model_dir ${DIFFDOCK_HOME}/workdir/paper_score_model \
  --ckpt best_ema_inference_epoch_model.pt \
  --confidence_model_dir ${DIFFDOCK_HOME}/workdir/paper_confidence_model \
  --confidence_ckpt best_model_epoch75.pt "$@"

.defファイルからコンテナイメージをビルドする

作成した .def ファイルからコンテナイメージ ( .sif ) をビルドする。

🍅 cd $SINGULARITY_IMAGE
🍅 singularity build --fakeroot DiffDock.sif ./def_file/DiffDock.def

DiffDockを実行してみる

D drive ( WSL2 において /mnt/d ) 上で DiffDock を実行してみる。
/mnt/d/diffdock/test/7vu6 というディレクトリを作成し、その中に以下のファイルを準備する。

蛋白質インプットファイル:
/mnt/d/diffdock/test/7vu6/input/7vu6_A_noHET.pdb
PDB code: 7VU6 (with chain A without HETATM)

リガンドインプットファイル:
/mnt/d/diffdock/test/7vu6/input/ensitrelvir_CID_162533924.sdf
PubChem CID: 162533924

DiffDock 実行スクリプト:
/mnt/d/diffdock/test/7vu6/diffdock_run.sh

diffdock_run.sh
#!/usr/bin/bash

# DiffDock.sifにおけるDiffDockの場所
DIFFDOCK_DIR="/usr/local/apps/DiffDock"
# diffdock_run.shを置いた場所を以降の起点とする
INFERENCE_DIR="$(cd $(dirname $0) && pwd)"

# 蛋白質およびリガンドのインプットファイルの場所 (都度修正)
protein_path="${INFERENCE_DIR}/input/7vu6_A_noHET.pdb"
ligand_path="${INFERENCE_DIR}/input/ensitrelvir_CID_162533924.sdf"

# ESM2で使用する.fastaの準備 (インプットのprotein.pdbから作成)
mkdir -p ${INFERENCE_DIR}/esm2_output
singularity exec --nv --nvccli --bind /mnt/d:/mnt/d ${SINGULARITY_IMAGE}/DiffDock.sif \
  python3 ${DIFFDOCK_DIR}/datasets/esm_embedding_preparation.py \
    --protein_path ${protein_path} \
    --out_file ${INFERENCE_DIR}/esm2_output/prepared_for_esm.fasta 

# ESM2によるembedding for proteinの作成
singularity exec --nv --nvccli --bind /mnt/d:/mnt/d ${SINGULARITY_IMAGE}/DiffDock.sif \
  python3 ${DIFFDOCK_DIR}/esm/scripts/extract.py \
    ${DIFFDOCK_DIR}/esm/model_weights/.cache/torch/hub/checkpoints/esm2_t33_650M_UR50D.pt \
    ${INFERENCE_DIR}/esm2_output/prepared_for_esm.fasta \
    ${INFERENCE_DIR}/esm2_output \
    --repr_layers 33 --include per_tok

# DiffDockの実行
singularity run --nv --nvccli --bind /mnt/d:/mnt/d ${SINGULARITY_IMAGE}/DiffDock.sif \
  --protein_path ${protein_path} \
  --ligand ${ligand_path} \
  --out_dir ${INFERENCE_DIR}/results \
  --esm_embeddings_path ${INFERENCE_DIR}/esm2_output \
  --inference_steps 20 \
  --save_visualisation \
  --samples_per_complex 40 \
  --batch_size 10 \
  --cache_path ${INFERENCE_DIR}/cache/cache

シェルスクリプトを実行する。

🍅 cd /mnt/d/diffdock/test/7vu6
🍅 chmod +x ./diffdock_run.sh
🍅 ./diffdock_run.sh

👇 DiffDock によるdocking simulation結果 ( SARS-CoV-2 3CL Protease vs. ensitrelvir )
Blue🐳: Crystal structure (PDB code: 7VU6)
Orange🍊: DiffDock prediction (top score)

Discussion

ログインするとコメントできます