Open Catalyst Project (OCP)とは?

Meta AIとCarnegie Mellon Universityが共同で進めているプロジェクトで、
化学反応のための触媒の性質を機械学習モデルで予測することを目的にしています。

厳密にどう切り分けているかは筆者もよくわかっていませんが、

(1) 触媒化学のための密度汎関数理論(DFT)計算のデータセット
(2) それを機械学習で予測するコンテスト
(3) およびそのデータセットを学習した機械学習原子間ポテンシャル
の総称になっているようです。(ちょっとわかりにくいです。)

この記事では (3) について、どうやって使い始めたら良いのかを紹介します。
最近では fairchem という名前のリポジトリになったようですね。

OCPでは様々な機械学習モデル (SchNet, Gemnet, Equiformer_v2など)を使って物質の構造データからポテンシャルエネルギーや力などを推定できます。

すでに学習済みのモデルを使えるので、最初に学習データを用意することなく使い始められます。

この記事で説明すること

上記のように、材料科学や触媒化学などの分野で便利そうなOCPですが、インストールしてちゃんと動かすまでが大変で、しかも解説記事などが少なかったので書くことにしました。具体的には以下のようなコードを動かし、学習済みモデルでの構造最適化計算ができるところまでの設定方法になります。

(以前のOCPレポジトリ時代はすごく大変だった記憶がありますが、fairchemになったものは割とインストールしやすくなった気がします)

今回は上記リポジトリの Quick Start のコード
( Cuのfcc(100)面上でのCO吸着の構造最適化計算 )
を動かしてみます。

最終的に動かしたいテスト計算

from ase.build import fcc100, add_adsorbate, molecule
from ase.optimize import LBFGS
from fairchem.core import OCPCalculator

# Set up your system as an ASE atoms object
slab = fcc100('Cu', (3, 3, 3), vacuum=8)
adsorbate = molecule("CO")
add_adsorbate(slab, adsorbate, 2.0, 'bridge')

calc = OCPCalculator(
    model_name="EquiformerV2-31M-S2EF-OC20-All+MD",
    local_cache="pretrained_models",
    cpu=False,
)
slab.calc = calc

# Set up LBFGS dynamics object
dyn = LBFGS(slab)
dyn.run(0.05, 100)

テスト環境

CPU: AMD Threadripper 3970X (32-core, SMT: disable)
GPU: Nvidia RTX-3060Ti
OS: Ubuntu 22.04 LTS

Python: 3.9.2 (installed using pyenv)

この記事では、計算でGPUが使えるように環境を整備します。
かなり速さが違うと思います。
厳密な計測ではありませんが、この環境で試すと1秒(GPU)と9分(CPU)でした。(後述）

cuda 11.8 のインストール

この後インストールするPytorchと関連パッケージとの兼ね合いからバージョンを選択します。今回は11.8が良さそうなのでそれを使います。

マイナーリビジョン(11.XのXの部分)は合ってなくても経験上いけると思います。
以下公式サイトに従って次のようにインストールします。

> wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-ubuntu2204.pinsudo mv cuda-ubuntu2204.pin /etc/apt/preferences.d/cuda-repository-pin-600
> wget https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda-repo-ubuntu2204-11-8-local_11.8.0-520.61.05-1_amd64.deb
> sudo dpkg -i cuda-repo-ubuntu2204-11-8-local_11.8.0-520.61.05-1_amd64.deb
> sudo cp /var/cuda-repo-ubuntu2204-11-8-local/cuda-*-keyring.gpg /usr/share/keyrings/
> sudo apt update
> sudo apt -y install cuda-11-8

私の環境では、/usr/local 以下に cuda-11.8 フォルダができました。
.bashrc に次の環境変数を設定します。

.bashrc

export PATH="/usr/local/cuda-11.8/bin:$PATH"
export LD_LIBRARY_PATH="/usr/local/cuda-11.8/lib64:$LD_LIBRARY_PATH"

source ~/.bashrc　として変更を反映します。

Python 仮想環境

pyenvを使いました。インストール方法などは割愛します。以下のgithubリポジトリから指示に従ってインストールします。

インストール後に、fairchem計算用のディレクトリを作成し、そこに仮想環境を構築します。
pyenv versions コマンドで見ると、3.9.2が既にインストールされていたのでこれをベースに仮想環境を作ります。python3の細かいバージョンはそんなに古いやつじゃなければ適当で良いです。

> mkdir fairchem
> cd fairchem
> pyenv virtualenv 3.9.2 fairchem
> pyenv local fairchem
(fairchem)>

プロンプトに名付けた仮想環境名が出れば無事設定完了です。

Pytorch のインストール

Pytorchの公式ドキュメントの通りに行います。

> python3.9 -m pip install --upgrade pip'
> pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

torch-geometric その他のパッケージも必要です。

# pytorchのバージョンを確認
> python -c "import torch; print(torch.__version__)"
2.3.1+cu118

# pytorchで呼んでるcudaのバージョンを確認
> python -c "import torch; print(torch.version.cuda)"
11.8

# torch_geometricをインストール
> pip install torch_geometric

#　他の必要パッケージをインストール
> pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.3.0+cu118.html

fairchem.core のインストール

いよいよfairchemをインストール

> pip install fairchem-core

Quick Startの実行

コード(再掲)

from ase.build import fcc100, add_adsorbate, molecule
from ase.optimize import LBFGS
from fairchem.core import OCPCalculator

# Set up your system as an ASE atoms object
slab = fcc100('Cu', (3, 3, 3), vacuum=8)
adsorbate = molecule("CO")
add_adsorbate(slab, adsorbate, 2.0, 'bridge')

calc = OCPCalculator(
    model_name="EquiformerV2-31M-S2EF-OC20-All+MD",
    local_cache="pretrained_models",
    cpu=False,
)
slab.calc = calc

# Set up LBFGS dynamics object
dyn = LBFGS(slab)
dyn.run(0.05, 100)

実行結果

WARNING:root:Detected old config, converting to new format. Consider updating to avoid potential incompatibilities.
WARNING:root:Skipping scheduler setup. No training set found.
WARNING:root:No seed has been set in modelcheckpoint or OCPCalculator! Results may not be reproducible on re-run

      Step     Time          Energy          fmax
LBFGS:    0 11:51:40        3.447266       11.352313
LBFGS:    1 11:51:40        0.809082        6.038936
LBFGS:    2 11:51:40        0.690918        7.082108
LBFGS:    3 11:51:40        0.350098        3.341267
LBFGS:    4 11:51:40        0.163940        3.049078
LBFGS:    5 11:51:40        0.062439        2.273317
LBFGS:    6 11:51:40       -0.161255        1.114096
LBFGS:    7 11:51:41       -0.262939        0.936718
LBFGS:    8 11:51:41       -0.376465        0.566177
LBFGS:    9 11:51:41       -0.457031        0.480961
LBFGS:   10 11:51:41       -0.475098        0.354243
LBFGS:   11 11:51:41       -0.481201        0.346403
LBFGS:   12 11:51:41       -0.492920        0.462854
LBFGS:   13 11:51:41       -0.503906        0.334754
LBFGS:   14 11:51:41       -0.511719        0.160774
LBFGS:   15 11:51:41       -0.519043        0.190647
LBFGS:   16 11:51:41       -0.525879        0.229972
LBFGS:   17 11:51:41       -0.531250        0.250488
LBFGS:   18 11:51:41       -0.540039        0.137720
LBFGS:   19 11:51:41       -0.544434        0.026457

True

なんか警告が出ますね。一体何が原因なのかわかる方いたらご教授ください。

[参考] CPUでやった場合
OCPCalculatorの引数をcpu=True にしてCPU計算をしてみました。
GPUでは数秒だった計算が数分かかりました。
他の数値計算も同時に行っていたりしていたので、本来はもう少しCPUでも速そうですが、GPUを使った方がやはり速そうです。

      Step     Time          Energy         fmax
LBFGS:    0 10:10:27        3.435445       11.3571
LBFGS:    1 10:10:56        0.804873        6.0475
LBFGS:    2 10:11:25        0.689559        7.0846
LBFGS:    3 10:11:54        0.347147        3.3434
LBFGS:    4 10:12:23        0.163328        3.0373
LBFGS:    5 10:12:52        0.062281        2.2672
LBFGS:    6 10:13:21       -0.162037        1.1142
LBFGS:    7 10:13:51       -0.262403        0.9378
LBFGS:    8 10:14:20       -0.378077        0.5592
LBFGS:    9 10:14:50       -0.456981        0.4777
LBFGS:   10 10:15:19       -0.474649        0.3526
LBFGS:   11 10:15:48       -0.481642        0.3457
LBFGS:   12 10:16:17       -0.493434        0.4654
LBFGS:   13 10:16:47       -0.503320        0.3323
LBFGS:   14 10:17:16       -0.512020        0.1605
LBFGS:   15 10:17:45       -0.518738        0.1905
LBFGS:   16 10:18:14       -0.525454        0.2341
LBFGS:   17 10:18:43       -0.530845        0.2512
LBFGS:   18 10:19:12       -0.539970        0.1350
LBFGS:   19 10:19:42       -0.543461        0.0270

True

この後のチュートリアルなど

この動画が実際の研究事例を詳しく説明していて面白かったです。

Case study として以下の論文が挙げられており、tutorialとしてgithubに同じ事例の計算例が収められています。

ここで使われている ocpmodels や ocdata.core といったパッケージはそれぞれ、fairchem.core や fairchem.data.oc.core に移されているので、若干修正が必要になってきますが、ちょっと直せば動きました。実際の論文の図とほぼ同じものが出力できるのでやってみると良いでしょう。

筆者も使い始めたばかりなので、これからもっと色々試せればいいなと思っています。