ESMFoldとは

2021年に登場したAlphaFold2より60倍くらい速く構造予測ができると噂のソフトウェア。この高速さを売りにして、2ヶ月くらいでMGnify90メタゲノムデータベースに登録されている6億のメタゲノムタンパク質（環境中に存在する身元不明のタンパク質）の構造をすべて予測した結果が ESM Metagenomics Atlas にて公開されている。（ちなみにAlphaFold2を使った予測構造はUniProtデータベースに登録されている2億のタンパク質に対して行われており、AlphaFold Protein Structure Databaseにてその構造が公開されている。）

論文には、約半数のタンパク質が高い信頼度（pLDDT）で予測できたと書かれている。

ESMFoldインストールログ

最近のLinux環境ならば以下の手順でインストールできるが、CUDA 11.6がインストールされており、そのnvccへのPATHが通っていることが必要。
condaを使って環境を$INSTALLDIRに作り、その中にconda系でpytorch系をインストールし（最新のPyTorch 2.0系でも動作した）、その後pipで不足しているパッケージをインストールすることにした。

###
# インストール先を指定
INSTALLDIR=/home/apps/ESMFold
###
mkdir -p ${INSTALLDIR}
cd ${INSTALLDIR}
wget -q -P . https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
bash ./Miniconda3-latest-Linux-x86_64.sh -b -p ${INSTALLDIR}/conda
rm Miniconda3-latest-Linux-x86_64.sh
. "${INSTALLDIR}/conda/etc/profile.d/conda.sh"
export PATH="${INSTALLDIR}/conda/condabin:${PATH}"
conda create -p ${INSTALLDIR}/esmfold-conda python=3.9 -y
conda activate ${INSTALLDIR}/esmfold-conda
conda update -n base conda -y

# 依存パッケージのインストール。cuda 11.6とpytorch 1.12.0を使用することにした（新しすぎると動かない可能性）
conda install pytorch==1.12.0 torchvision==0.13.0 torchaudio==0.12.0 cudatoolkit=11.6 -c pytorch -c conda-forge
conda install -y -c conda-forge setuptools openmm==7.5.1 pdbfixer==1.7 einops fairscale
conda install -y -c bioconda hmmer==3.3.2 hhsuite==3.3.0

# Install ESMFold
python3.9 -m pip install git+https://github.com/facebookresearch/esm.git
python3.9 -m pip install "fair-esm[esmfold]"
# 必ず先にPyTorchとCUDA 11.8を使える環境にしておかないと失敗する
python3.9 -m pip install 'openfold @ git+https://github.com/aqlaboratory/openfold.git@4b41059694619831a7db195b7e0988fc4ff3a307'

# Install dllogger
python3.9 -m pip install 'dllogger @ git+https://github.com/NVIDIA/dllogger.git'
# Modelingに必要なパッケージ
python3.9 -m pip install pytorch_lightning==1.5.10 wandb==0.12.21
python3.9 -m pip install omegaconf hydra-core pandas pytest matplotlib biotite mypy flake8 jupyter ipykernel

これによって、$INSTALLDIRのパス上にesmfoldの動作環境がインストールされる。${INSTALLDIR}/esmfold-conda/bin以下にはesm-foldというコマンドがインストールされ、これを使うことでターミナルから配列ファイルを引数にとって予測させることができる。

これを便利に使うために、${INSTALLDIR}/binというディレクトリを作り、その中にesmfoldというbashスクリプトラッパーを書くとよいかもしれない。

###
# インストール先を指定
INSTALLDIR=/home/apps/ESMFold
###
mkdir -p ${INSTALLDIR}/bin

# 'esmfold'というbashスクリプトを${INSTALLDIR}/bin以下に作成する
cat << 'EOS' > ${INSTALLDIR}/bin/esmfold
#!/bin/bash

INSTALLDIR=/home/apps/ESMFold
. "${INSTALLDIR}/conda/etc/profile.d/conda.sh"
export PATH="${INSTALLDIR}/conda/condabin:${PATH}"
conda activate ${INSTALLDIR}/esmfold-conda

# Environment Modulesでcuda 11.8を呼び出すための設定
module load cuda/11.8

# 共用計算機用に、環境変数TORCH_HOMEを指定しておくことで、
# 学習済みパラメータファイルを指定したディレクトリに保存させる
export TORCH_HOME=${INSTALLDIR}

# ${INSTALLDIR}/esmfold-conda/bin/esm-fold を呼び出し、引数を取れるようにする
esm-fold $@
EOS

# 実行権限付与
chmod +x ${INSTALLDIR}/bin/esmfold

使い方の例

GPUを1つ使う。

run_esmfold.sh

#!/bin/sh

INSTALLDIR=/home/apps/ESMFold
export PATH=${INSTALLDIR}/bin:${PATH}

esmfold \
  -i foo.fasta \
  -o . \
  --chunk-size 128 \
  --num-recycles 4

-iはインプットのFASTAファイル。Multimerを予測したい場合は

foo.fasta

>hoge
AAAAAA:BBBBBB

のように、:で配列を区切る。

-o <directory_name>はアウトプットファイルを置くためのディレクトリを指定
--chunk-size 128は配列長が大きい場合（700くらい？）に使うと良い。そうでない場合は不要。
--num-recycles 4はリサイクル回数の指定。4回がデフォルト。

詳しいオプションはGitHub参照。

pLDDTスコアは埋め込まれているが、PAEは出てこないっぽい。原理的にはできるはずだし
https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/ESMFold.ipynb
によれば表示させることはできるようだ。

GPU起動のオーバーヘッドがあるので起動に時間が掛かるが、計算が始まると2〜3分で終わる。930残基くらいの複合体予測も5〜10分くらいで終わる（RTX3090を使用）。

感想

~~類縁配列が非常に少ないorphanタンパク質やde novo タンパク質の構造予測にはAlphaFold2よりも良い性能を示すようだ。~~
こっそり書かれているFigure S8によれば、orphanタンパク質に対してもAlphaFold2の構造予測精度を上回らないことが示されていました。de novoタンパク質については不明。
天然に存在する2億以上のうちの多く（98％以上）のタンパク質についてはAlphaFold2の高信頼度(pLDDT)には及ばず、予測が失敗することが多々あります。全体で見れば、類縁配列が豊富なタンパク質は比較的高いpLDDTで予測されることが多いが、そうでないものについては失敗しがち。

AlphaFold2の計算速度より60倍以上速いことがウリだが、(Local)ColabFoldがAlphaFold2と同精度を保ちつつ20〜30倍ほど高速化を達成しているので、複合体予測も基本的にはこっちでいいかな……って気がしてますね……。