LocalColabFoldを完全ローカル環境で動作させる #bioinformatics

創薬 (dry) Advent Calendar 2023の12月8日（金）のテーマ

超大容量RAM搭載サーバーを持っている、または名古屋大学のスパコン不老の上でLocalColabFoldを計算中のインターネット通信なしで使う方法の紹介。

Introduction

LocalColabFoldはColabFoldが持つ非常に高速な構造予測と高機能なオプションが売りのAlphaFold2の派生品です。また、ColabFoldの機能性を自前のマシン環境で無制限に使うというのがLocalColabFoldのコンセプトなので、ColabFoldができることはLocalColabFoldでもすべて実現できます。
特にColabFoldを用いた複合体予測はそのMSAの計算パートがさらに高速化され、AlphaFold2を用いて単純に複合体予測させるよりも大変速く予測できることで重宝します。多くのタンパク質は小分子や他のタンパク質と結合してその機能を発揮することが多いこともあるため、ColabFoldのこの特長を用いて、タンパク質–タンパク質複合体の全探索を行うことに大きく貢献します。

一方で、ColabFoldはインターネットを経由して入力とするアミノ酸配列情報をMSA生成のサーバーに送信しなければならないため、製薬企業さんからは「便利なんだけど、アミノ酸配列情報を外部に出してはいけないという制約があるから使うことができない」と言う話をよく聞きました。

しかし、2023年10月にColabFoldを完全ローカル化させることができるようになりましたので、計算環境さえ揃っていれば、このColabFoldの制約を気にすることなく使えるようになりました。まだ変な使い方をしようとするとちょいちょいバグがあるんですけど（※筆者が対応を放置している）、まぁたいていはなんとかなります。

この記事では、そういった超ハイスペックなマシンを持っている製薬企業のために、計算環境を設定する方法を記します。

動作に必要な計算環境

AlphaFold2もColabFoldも、入力とするアミノ酸配列に対してMSAを計算させるパートと得られたMSAを入力の特徴量としてタンパク質立体構造を予測させるパートに分かれています。AlphaFold2はそのパイプライン仕様上融合した形になっていますが、ColabFoldはそれらの計算環境を分けることができます。具体的には、MSAの計算パートは大容量RAMとSSDを持つ計算ノードで行わせ、そこから得られたMSAファイル（.a3mファイル）を使って、GPUがある計算ノードで予測するというように切り分けることができます。

MSAを作成させるマシン環境

2TBのSSD（これ専用に外部SSDを追加した方が良い）
768GB以上のRAM
8〜16コア以上のCPUコア数

MSA作成部分では実は1つの配列に対する要求CPUコア数は少なく、2〜4コア割り当てれば十分な気がします。コア数が多いマシンほど、異なるアミノ酸インプット配列に対するMSAを並列して同時に作成させることができます。

構造予測させるマシン環境

VRAMが大きいGPU。例：RTX3090やRTX4090。

こちらでは大容量SSDやRAMを用意させる必要はありませんが、原理上はMSA計算マシンと1台にまとめることも可能です（私はやっていないですけど）。RTX3090や4090ならば合計1800〜2000残基くらいまでは予測できます。もし複数台用意できるのであればこちらをたくさん用意しましょう。

LocalColabFoldのインストール

インストールは以前のQiita記事の方も呼んでください。森脇のGitHubはこちら。

こちらでインストールされるcolabfold_searchがMSA作成計算のためのコマンドです。これにより、入力の配列からMSAファイル（a3mファイルフォーマット）を生成します。colabfold_batchが構造予測部分を担います。colabfold_batchは通常インターネットを経由してMSAをWebサーバーから取得しようとしますが、インプットにa3mファイルフォーマットを指定した場合はインターネット接続なしで（ローカルで）計算を行います。

この記事では前者のcolabfold_searchを使えるようにすることが主題です。

MMseqs2のインストール

colabfold_searchは事実上MMseqs2を使った配列検索のためのPythonで書かれたラッパースクリプトで、つまりはMMseqs2がインストールされていないと動作しません。なのでそれをインストールします。

インストール方法は以下の通り

git clone https://github.com/soedinglab/MMseqs2.git
cd MMseqs2/
git checkout 71dd32ec43e3ac4dabf111bbc4b124f1c66a85f1
mkdir build ; cd build
# コンパイルにはある程度新しいGCCコンパイラが必要。GCC 9以上？
cmake .. -DCMAKE_INSTALL_PREFIX=/path/to/your/mmseqs2/for_colabfold -DCMAKE_BUILD_TYPE=Release
make -j8 install

コミットハッシュが71dd32ec43e3ac4dabf111bbc4b124f1c66a85f1のものを使うのがポイントです。これを使うとウェブサーバーとほぼ同じような結果が得られるようになります（※完全ではないですが）。
オプションでMPI対応にもできますが、デフォルトでOpenMPのマルチスレッド処理をしてくれるので、近年のマルチコアCPUであればそれを付ける必要はありません。

これを使って以下のデータベースをセットアップします。

配列データベースセットアップ

ColabFoldをローカルで動かすためにはAlphaFold2の時と同様に、自前で配列データベースと構造データベースをセットアップする必要がありますので、その準備をします。

一括でセットアップしてくれるスクリプトはhttps://github.com/sokrypton/ColabFold/blob/main/setup_databases.sh にあるのですが、日本からだとダウンロードがかなり遅いので気長にやる必要があると思います。セットアップ時はaria2コマンドを使えるようにしておくと便利かもしれません。

以下でインストールするデータベースはSSDの/mnt/databasesに作成し、その上で実行するものとします。以降この/mnt/databasesは各ユーザーのPATHに適宜置き換えてください。また、全部で配列データベースのファイルサイズは最終的には合計1.2TBほどになりますが、処理の都合上途中でそれ以上に容量を使うことになる（2TBくらい？）ので、あらかじめ余裕を持っておきましょう。

配列データベース1：UNIREF30DB

ColabFoldチームが作成したクラスタリング済みのUNIREF30DBをダウンロードします。上記シェルスクリプトでは以下の部分に対応します。

UNIREF30DB="uniref30_2302"
if [ ! -f UNIREF30_READY ]; then
  downloadFile "https://wwwuser.gwdg.de/~compbiol/colabfold/${UNIREF30DB}.tar.gz" "${UNIREF30DB}.tar.gz"
  tar xzvf "${UNIREF30DB}.tar.gz"
  mmseqs tsv2exprofiledb "${UNIREF30DB}" "${UNIREF30DB}_db"
  if [ -z "$MMSEQS_NO_INDEX" ]; then
    mmseqs createindex "${UNIREF30DB}_db" tmp1 --remove-tmp-files 1
  fi
  if [ -e ${UNIREF30DB}_db_mapping ]; then
    ln -sf ${UNIREF30DB}_db_mapping ${UNIREF30DB}_db.idx_mapping
  fi
  if [ -e ${UNIREF30DB}_db_taxonomy ]; then
    ln -sf ${UNIREF30DB}_db_taxonomy ${UNIREF30DB}_db.idx_taxonomy
  fi
  touch UNIREF30_READY
fi

つまり、 https://wwwuser.gwdg.de/~compbiol/colabfold/uniref30_2302.tar.gz から該当圧縮ファイルをダウンロードしてきて解凍し、mmseqs tsv2exprofiledbとmmseqs createindexコマンドを使って配列データベースをセットアップする形になっています。またln -sfでシンボリックリンクを貼っています。

ファイルのダウンロードとmmseqsの実行は相当時間がかかりますし、SSDの上でやらないと大変時間がかかります。ファイルがすべて用意されると以下のようになるはずです。

/mnt/databases
├── uniref30_2302_db
├── uniref30_2302_db_aln
├── uniref30_2302_db_aln.dbtype
├── uniref30_2302_db_aln.index
├── uniref30_2302_db.dbtype
├── uniref30_2302_db_h
├── uniref30_2302_db_h.dbtype
├── uniref30_2302_db_h.index
├── uniref30_2302_db.idx
├── uniref30_2302_db.idx.dbtype
├── uniref30_2302_db.idx.index
├── uniref30_2302_db.idx_mapping -> uniref30_2302_db_mapping
├── uniref30_2302_db.idx_taxonomy -> uniref30_2302_db_taxonomy
├── uniref30_2302_db.index
├── uniref30_2302_db_mapping
├── uniref30_2302_db_seq
├── uniref30_2302_db_seq.dbtype
├── uniref30_2302_db_seq_h -> uniref30_2302_db_h
├── uniref30_2302_db_seq_h.dbtype -> uniref30_2302_db_h.dbtype
├── uniref30_2302_db_seq_h.index -> uniref30_2302_db_h.index
├── uniref30_2302_db_seq.index
├── uniref30_2302_db_taxonomy
└── UNIREF30_READY

->がついているものはシンボリックリンクファイルです。

配列データベース2：COLABDB

MMSEQS_NO_INDEX=${MMSEQS_NO_INDEX:-}
if [ ! -f COLABDB_READY ]; then
  downloadFile "https://wwwuser.gwdg.de/~compbiol/colabfold/colabfold_envdb_202108.tar.gz" "colabfold_envdb_202108.tar.gz"
  tar xzvf "colabfold_envdb_202108.tar.gz"
  mmseqs tsv2exprofiledb "colabfold_envdb_202108" "colabfold_envdb_202108_db"
  # TODO: split memory value for createindex?
  if [ -z "$MMSEQS_NO_INDEX" ]; then
    mmseqs createindex "colabfold_envdb_202108_db" tmp2 --remove-tmp-files 1
  fi
  touch COLABDB_READY
fi

の処理部分に対応します。同様にして配列データベースができると、

/mnt/databases
├── COLABDB_READY
├── colabfold_envdb_202108_db
├── colabfold_envdb_202108_db_aln
├── colabfold_envdb_202108_db_aln.dbtype
├── colabfold_envdb_202108_db_aln.index
├── colabfold_envdb_202108_db.dbtype
├── colabfold_envdb_202108_db_h
├── colabfold_envdb_202108_db_h.dbtype
├── colabfold_envdb_202108_db_h.index
├── colabfold_envdb_202108_db.idx
├── colabfold_envdb_202108_db.idx.dbtype
├── colabfold_envdb_202108_db.idx.index
├── colabfold_envdb_202108_db.index
├── colabfold_envdb_202108_db_seq
├── colabfold_envdb_202108_db_seq.dbtype
├── colabfold_envdb_202108_db_seq_h -> colabfold_envdb_202108_db_h
├── colabfold_envdb_202108_db_seq_h.dbtype -> colabfold_envdb_202108_db_h.dbtype
├── colabfold_envdb_202108_db_seq_h.index -> colabfold_envdb_202108_db_h.index
└── colabfold_envdb_202108_db_seq.index

のようになります。
->がついているものはシンボリックリンクファイルです。

配列データベース3：PDB

if [ ! -f PDB_READY ]; then
  downloadFile "https://wwwuser.gwdg.de/~compbiol/colabfold/pdb100_230517.fasta.gz" "pdb100_230517.fasta.gz"
  mmseqs createdb pdb100_230517.fasta.gz pdb100_230517
  if [ -z "$MMSEQS_NO_INDEX" ]; then
    mmseqs createindex pdb100_230517 tmp3 --remove-tmp-files 1
  fi
  touch PDB_READY
fi

の処理部分に相当します。同様にして配列データベースができると、

/mnt/databases
├── pdb100_230517
├── pdb100_230517.dbtype
├── pdb100_230517_h
├── pdb100_230517_h.dbtype
├── pdb100_230517_h.index
├── pdb100_230517.idx
├── pdb100_230517.idx.dbtype
├── pdb100_230517.idx.index
├── pdb100_230517.index
├── pdb100_230517.lookup
├── pdb100_230517.source
├── pdb100_a3m.ffdata
├── pdb100_a3m.ffindex
└── PDB100_READY

が作成されます。

配列データベース4：PDB100

if [ ! -f PDB100_READY ]; then
  downloadFile "https://wwwuser.gwdg.de/~compbiol/data/hhsuite/databases/hhsuite_dbs/pdb100_foldseek_230517.tar.gz" "pdb100_foldseek_230517.tar.gz"
  tar xzvf pdb100_foldseek_230517.tar.gz pdb100_a3m.ffdata pdb100_a3m.ffindex
  touch PDB100_READY
fi

の処理部分に相当します。やっていることはpdb100_foldseek_230517.tar.gzの圧縮ファイルからpdb100_a3m.ffdataとpdb100_a3m.ffindexファイルを取り出しているだけです。

/mnt/databases
├── pdb100_a3m.ffdata
├── pdb100_a3m.ffindex
└── PDB_READY

が作成されます。

PDB MMCIFデータベース

こちらは配列ではなくProtein Data Bankの構造ファイルそのものです。以下の処理に相当します。

# 日本からの場合はPDB_SERVERをPDBjのものに変更したほうがよい。
# PDB_SERVER="${2:-"rsync.wwpdb.org::ftp"}"
PDB_SERVER="${2:-"ftp.pdbj.org::ftp"}"
# PDB_PORT="${3:-"33444"}"
if [ ! -f PDB_MMCIF_READY ]; then
  mkdir -p pdb/divided
  mkdir -p pdb/obsolete
  rsync -rlpt -v -z --delete ${PDB_SERVER}/data/structures/divided/mmCIF/ pdb/divided
  rsync -rlpt -v -z --delete ${PDB_SERVER}/data/structures/obsolete/mmCIF/ pdb/obsolete
  touch PDB_MMCIF_READY
fi

もしAlphaFold2をローカルで動かす準備をされていた事がある場合、download_pdb_mmcif.shスクリプトで生成されたpdb_mmcif/mmcif_filesディレクトリをそのまま利用することができます。一見上のディレクトリ構成と異なるファイル構成になりますが、両対応しています。なので、このステップは飛ばしてください。

こちらは高速なファイル読み書きが必要ないので、SSD上に置く必要はありません。言い換えれば、もしMSA計算マシンに他の別のネットワークドライブをマウントしている状態であれば、そちらのドライブ上に置いていてもColabFoldの計算速度にほとんど影響しません。

以上で必要なデータベースの作成は終了です。

colabfold_searchコマンドでローカルでMSAを高速生成する

以上のデータベースが揃ったら、LocalColabFoldのcolabfold_searchコマンドを用いることで高速にMSAを生成することができます。例として、以下のFASTAファイルras_raf.fastaを用いてRAS-RAF複合体のMSAファイル生成をやってみましょう。ColabFoldでは:を挟むことで複合体を予測することができます（※本家AlphaFold2の入力形式と異なることに注意）。もちろん単量体を予測させることもできます。

ras_raf.fasta

>tr|RAS_RAF
MTEYKLVVVGAGGVGKSALTIQLIQNHFVDEYDPTIEDSYRKQVVIDGETCLLDILDTAG
QEEYSAMRDQYMRTGEGFLCVFAINNTKSFEDIHQYREQIKRVKDSDDVPMVLVGNKCDL
AARTVESRQAQDLARSYGIPYIETSAKTRQGVEDAFYTLVREIRQHKLRKLNPPDESGPG
CMSCKCVLS:
PSKTSNTIRVFLPNKQRTVVNVRNGMSLHDCLMKALKVRGLQPECCAVFRLLHEHKGKKA
RLDWNTDAASLIGEELQVDFL

# PATHの指定
MMSEQS_PATH="/path/to/your/mmseqs2/for_colabfold"
DATABASE_PATH="/mnt/databases"
# 入力ファイル・出力ディレクトリの指定
INPUTFILE="ras_raf.fasta"
OUTPUTDIR="ras_raf"

colabfold_search \
  --use-env 1 \
  --use-templates 1 \
  --db-load-mode 2 \
  --db2 pdb100_230517 \
  --mmseqs ${MMSEQS_PATH}/bin/mmseqs \
  --threads 4 \
  ${INPUTFILE} \
  ${DATABASE_PATH} \
  ${OUTPUTDIR}

5分くらいで、結果としてRASとRAFについてのMSAファイルtr_RAS_RAF.a3mと、PDBでの配列ヒット結果ファイルtr_RAS_RAF_pdb100_230517.m8が得られます。

以下これについての補足

--db-load-mode 2としていることで、RAMに展開されている配列データベースインデックスファイル（拡張子.idx）の中身を高速に検索できるようです。また、あるサーバーにおいて常時この*.idxファイルをメモリに展開した状態を維持したい場合は、vmtouchというツールがあり、これを使って

cd /mnt/databases
vmtouch -f -w -t -l -d -m 1000G *.idx

を実行ししばらく処理終了まで待つと、*.idxファイルがRAMに展開された状態になります。このvmtouchプロセスが生きている間、効果が持続するようです。

名古屋大学のスパコン不老では、Type3ノードにログインしたあとでのみDATABASE_PATH="/nvme1/z44243z/colabfold_env"に上記データベースが存在するので試してみてください（MMseqs2は自前のディレクトリ上にインストールしておく必要があります）。

並列処理用スクリプト

上記の通り、colabfold_searchの実態であるmmseqs2はそれほどコア数を必要としないので、GNU parallelを使うと同時並行してMSAを計算させることが可能です。大変効率良く計算できるようになるので、これを使ってみましょう。GNU parallelの使い方は https://bicycle1885.hatenablog.com/entry/2014/08/10/143612 などが参考になります。

GNU parallelのソースからのインストールはここでは省略しますが、そんなに難しくはないです。RedHat系OSの場合は

$ sudo dnf -y install parallel

Debian系であれば

$ sudo apt -y install parallel

でもコンパイル済みのものがインストール可能です。

あとは

NUM_PARALLEL=8
parallel -j ${NUM_PARALLEL} --noswap \
    'input=`basename {} .fasta` ; mkdir -p ${input} ; \
    if [ ! -s ${input}/${input}.a3m ]; then \
        colabfold_search   --use-env 1 \
                           --use-templates 1 \
                           --db-load-mode 2 \
                           --db2 pdb100_230517 \
                           --mmseqs /path/to/mmseqs \
                           --threads 4 \
                           ${input}.fasta /mnt/databases ${input} > ${input}/${input}.log ; \
    fi' ::: `find . -name "*.fasta"`

のような形で、parallelでcolabfold_searchを並列化させることができます。上の例の場合はNUM_PARALLEL=8と設定されているので、そのディレクトリ中にあるfastaファイルすべて（find . -name "*.fasta"で取得している）に対して、1ノード上で常時8個の同時処理を実現しています。

colabfold_batchでローカルで構造を予測する

colabfold_batchは、指定されたINPUTFILEに対してその立体構造を予測するためのコマンドです。LocalColabFoldがインターネット接続を試みようとするかしないかは、以下の条件で決まります。

INPUTFILEに拡張子.fastaが指定された場合、MSA取得をインターネットを通じて行おうとする。
INPUTFILEにMSAファイルであることを示す拡張子.a3mが指定された場合、インターネットを通じたMSA取得は行われない
--templatesが指定された場合、--pdb-hit-fileと--local-pdb-pathが指定されていればインターネットを通じてテンプレート構造の検索と取得を行わない。指定されていなければ、インターネットを通じて取得を試みる。
--templatesが指定され、かつ--custom-template-path /path/to/dirが指定された場合、その/path/to/dir内に存在するPDBまたはmmcifファイルのみをテンプレート構造の候補としてみなし、インターネット接続は行わない。

という仕様なので、先のcolabfold_searchで2つのファイルが得られたら、colabfold_batchのINPUTFILEに先程のtr_RAS_RAF.a3mファイルを、PDBHITFILEにtr_RAS_RAF_pdb100_230517.m8を、そしてLOCALPDBPATHにAlphaFold2のときに作成したpdb_mmcif/mmcif_filesディレクトリ、またはdividedディレクトリがある場合はその親ディレクトリ（/path/to/pdb/divided/db/2dba.cif.gzだった場合は/path/to/pdbディレクトリ）を指定します。

INPUTFILE="tr_RAS_RAF.a3m"
PDBHITFILE="tr_RAS_RAF_pdb100_230517.m8"
LOCALPDBPATH="/path/to/pdb_mmcif/mmcif_files"
RANDOMSEED=0

colabfold_batch \
  --amber \
  --templates \
  --use-gpu-relax \
  --pdb-hit-file ${PDBHITFILE} \
  --local-pdb-path ${LOCALPDBPATH} \
  --random-seed ${RANDOMSEED} \
  ${INPUTFILE} \
  ras_raf

5〜6分くらいで、構造予測が終了すると思います。正しく終了すると0バイトのtr_RAS_RAF.done.txtマーカーが出力されます。