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

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

https://adventar.org/calendars/9092

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

## Introduction

LocalColabFoldは[ColabFold](https://github.com/sokrypton/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記事](https://qiita.com/Ag_smith/items/fca48002fbdcb15145c0)の方も呼んでください。森脇のGitHubは[こちら](https://github.com/YoshitakaMo/localcolabfold)。

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

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

## MMseqs2のインストール

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

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

```bash
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をダウンロードします。上記シェルスクリプトでは以下の部分に対応します。

```bash
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の上でやらないと大変時間がかかります。ファイルがすべて用意されると以下のようになるはずです。

```shell
/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

```bash
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
```

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

```shell
/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

```bash
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
```

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

```shell
/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

```bash
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`ファイルを取り出しているだけです。

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

が作成されます。

### PDB MMCIFデータベース

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

```bash
# 日本からの場合は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
```

```bash
# 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 8 \
  ${INPUTFILE} \
  ${DATABASE_PATH} \
  ${OUTPUTDIR}
```

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

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

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

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

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

::::note
また、`${INPUTFILE}`には**複数のアミノ酸配列が記述された1つのFASTAファイル**、または**複数のfastaファイルを置いたディレクトリを指定することも可能です**。こちらを使う方が大規模にMSAを取得したいときに非常に有用で、同じアミノ酸配列が中に含まれていた場合はそれについて再度計算しなくなったり、計算前のオーバーヘッドがなくなったりとさらに高速化されます
::::

## 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`ディレクトリ）を指定します。


```bash
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`マーカーが出力されます。