MLIP バックエンド¶
計算バックエンドの選び方と、バックエンドごとのインストール、モデル名、精度、再現性の設定、Hessian の計算方式をまとめたページです。デフォルトのバックエンドは UMA(Meta の Universal Models for Atoms)で、-b/--backend で ORB、MACE、AIMNet2 も選べます。4 つとも機械学習原子間ポテンシャル(MLIP)です。
バックエンドごとの特性¶
バックエンドは、計算を行うどのコマンドでも -b/--backend で選びます。
# UMA(デフォルト)
pdb2reaction opt -i input.pdb -q 0
# ORB
pdb2reaction opt -i input.pdb -q 0 -b orb
# MACE
pdb2reaction opt -i input.pdb -q 0 -b mace
# AIMNet2
pdb2reaction opt -i input.pdb -q 0 -b aimnet2
バックエンド |
インストール |
モデル名 |
|
解析 Hessian |
複数ワーカー |
|---|---|---|---|---|---|
|
同梱( |
|
|
あり(autograd) |
あり |
|
|
|
|
あり(autograd) |
なし |
|
専用の conda 環境に pdb2reaction を入れてから |
|
|
あり |
なし |
|
|
|
|
あり |
なし |
--backend-model NAME は、選んだ --backend のモデルを替えます(例:--backend uma --backend-model uma-m-1p1)。
実行時には、読み込むバックエンドとモデルが [backend] Preparing MLIP model (<backend> / <model>)... の形で表示されます。デフォルトのモデルでは括弧の中が UMA / UMA-S-1.2 (OMol)、ORB / ORB-v3-conservative-OMol、MACE / MACE-OMOL-0、AIMNet2 / aimnet2 になります。
精度(precision)¶
--precision fp32|fp64 は、どのバックエンドでも MLIP の推論の浮動小数点精度を決めます。値は UMA の precision、ORB の precision、MACE の default_dtype に渡されます。AIMNet2 には精度の設定がなく、fp32 を指定しても何も変わりません。
--precision を指定しないときは、バックエンドごとのデフォルト値を使います。
バックエンド |
デフォルト値 |
理由 |
|---|---|---|
|
fp32 |
上流の fairchem の基準の設定です。 |
|
fp64 |
ORB の |
|
fp64 |
MACE は上流で |
|
fp32 |
精度の設定がありません。 |
どの値を選ぶかは目的で決めます。
目的 |
推奨 |
理由 |
|---|---|---|
通常の計算 |
指定しない |
上のデフォルト値(UMA・AIMNet2 は fp32、ORB・MACE は fp64)のままにします。 |
速さを優先するスクリーニング |
必要なときだけ |
ORB・MACE の精度が下がります(使用上の注意点)。 |
最終の TS と Hessian |
指定しない。UMA で n_imag ≥ 2 のときは |
精度によらず、 |
fp64 は次のように指定します。
pdb2reaction tsopt -i ts.pdb -q 0 --precision fp64 ...
pdb2reaction freq -i opt.pdb -q 0 --precision fp64 ...
pdb2reaction irc -i ts.pdb -q 0 --precision fp64 ...
YAML では次のように書きます。
calc:
precision: fp64
決定論的実行と再現性¶
--deterministic を付けると、同じソフトウェアと GPU で、同じ入力から同じ結果が得られます。付けないと、GPU では同じ入力の 2 回の計算でも最後の桁が違うことがあります。
--deterministic は、PyTorch の決定論的アルゴリズム(torch.use_deterministic_algorithms)を有効にし、GPU で決定論的に動く版の無い PyTorch の演算 1 つを置き換えます。DFT の計算、自作の ASE calculator、PyTorch の外の GPU のコードは制御しません。
pdb2reaction opt -i input.pdb -q 0 --deterministic
pdb2reaction all -i r.pdb p.pdb -q -1 --tsopt --deterministic
プロセス全体に効きます。
allに付ければallが実行する MLIP の段すべてに効くので、段ごとに付ける必要はありません。計算は遅くなります。決定論的な GPU の演算は処理が遅いので、繰り返しの計算を一致させる必要があるときだけ使ってください。
実行する演算に PyTorch の決定論的な版が無いときは、再現しない結果を黙って出さずに、エラーで止まります。
バックエンド |
|
|---|---|
|
対応 |
|
PyTorch の決定論的モードは有効になります。入れた版で 2 回の計算が一致するか確かめてください |
|
非対応:エラーで止まります(使用上の注意点) |
|
自作の ASE calculator しだいで、このフラグでは保証できません |
ワーカーと Hessian の計算方式¶
--uma-workers N(デフォルト 1)は UMA の予測器を N 個並列に動かし、--uma-workers-per-node(デフォルト 1)はそのうち 1 ノードで動かす数を決めます。どちらのフラグも opt、tsopt、freq、irc、sp、all、path-opt、path-search、scan、scan2d、scan3d にあります。ORB、MACE、AIMNet2 はこれらを警告を出して無視します。ワーカーを増やすと速くなるかは HPC 実行例 › ウォールタイム見積り にあります。
Hessian の計算方式¶
--hessian-calc-mode で Hessian の計算方式を選びます。FiniteDifference(デフォルト)は力の中心差分を取り、Analytical は選んだデバイスで 2 階の自動微分を行います。UMA、ORB、MACE、AIMNet2、DFT は解析 Hessian を計算でき、自作の calculator は FiniteDifference だけに対応します。
xTB 溶媒補正¶
MLIP のバックエンドでは、--solvent NAME で xTB の溶媒和エネルギー E_xTB(solvent) - E_xTB(vacuum) と、それに対応する力と Hessian の差を MLIP のポテンシャル面に足します。主な用途は小分子の溶液中の計算です。--solvent-model で alpb(デフォルト)か cpcmx を選びます。--solvent-xtb-cmd には追加の引数を含めた xTB のコマンドを渡せます。xTB の SCC(自己無撞着電荷)の反復が収束しないときは、'xtb --etemp 1000' のように指定してください。xTB は pdb2reaction と一緒には入らないので、conda install -c conda-forge xtb で導入してください。
補正のたびに xTB を 2 回(溶媒中と真空中)実行するので、200〜300 原子くらいが実用の目安です(上限ではありません)。実際の系と計算機で計算時間を測ってください。この補正は名前の付いたバルクの溶媒を表すので、酵素クラスターに使うのはその環境を仮定する根拠があるときだけにしてください。溶液中の障壁とクラスターの障壁を比べるときは、反応する化学種、電荷、多重度、バックエンドとモデル、エネルギーの基準をそろえてください。
DFT バックエンド¶
sp、opt、tsopt、irc、freq、scan、scan2d、scan3d、path-opt、path-search、all は -b dft --func-basis FUNCTIONAL/BASIS --dft-engine gpu|cpu(デフォルトは wb97m-v/def2-svp と gpu)を受け付け、すべてのエネルギーと力を PySCF/GPU4PySCF で計算します。ポピュレーション解析付きの一点計算には、別の pdb2reaction dft コマンドがあります。低メモリモード、CPU のスレッド数とホスト RAM、SCF のチェックポイントは MLIP の TS を DFT で確かめる にあります。
カスタムバックエンド — 任意の ASE Calculator を使う(--calc-file)¶
組み込みの MLIP バックエンドのほかに、--calc-file で任意の ASE Calculator を実行時に使えます。pdb2reaction 本体を変える必要はありません。GFN-xTB(tblite か xtb-python 経由)、DFTB+、ORCA、Psi4 など、ASE に対応した計算エンジンをつなげます。受け渡しは標準の ASE Calculator の形(エネルギーは eV、力は eV/Å)です。
ASE Calculator を返す get_calculator 関数を持つ Python ファイルを書きます。
# my_calc.py(最小の例)
from ase.calculators.emt import EMT
def get_calculator(charge=0, spin=1, device="auto", **kwargs):
return EMT()
EMT() を、GFN-xTB の tblite.ase.TBLite(...)、DFTB+ の ASE calculator、ase.calculators.orca.ORCA(...) など、使いたいエンジンに替えてください。このファイルを各コマンドか all に渡すと custom バックエンドが選ばれ、--backend の指定より優先されます。
pdb2reaction sp -i model.xyz --calc-file my_calc.py -q 0 -m 1
pdb2reaction opt -i model.xyz --calc-file my_calc.py -q 0 -m 1
pdb2reaction tsopt -i ts.xyz --calc-file my_calc.py -q 0 -m 1
pdb2reaction freq -i ts.xyz --calc-file my_calc.py -q 0 -m 1
pdb2reaction all -i R.pdb P.pdb -c 'A:LIG' --calc-file my_calc.py -q 0 -m 1
関数が引数か
**kwargsで受け取れば、charge、spin、deviceが渡されるので、総電荷が要るエンジン(xTB など)も設定できます。spinは多重度で、mult・multiplicityの名前でも渡します。関数の名前は--calc-file-func-name NAMEで変えられ、その名前に Calculator のインスタンスを置いてもかまいません。Hessian は力の有限差分で求めるので、
freqとtsopt --opt-mode hessはどのエンジンでも動きます。--freeze-links・--freeze-atomsによる固定もふつうどおり効きます。all、sp、opt、tsopt、freq、irc、scan・scan2d・scan3d、path-opt、path-searchで使えます。allは、calculator を使うすべての段に同じ関数を渡します。独自の--backend名を持つ、インストールできるバックエンドにするときは 開発者向け を見てください。
Python API¶
クイックスタート¶
import numpy as np
from pdb2reaction.backends.uma import UMACalculator
# 例: 中性一重項の2原子系(GPUが利用可能ならGPU、なければCPU)
calc = UMACalculator(charge=0, spin=1, model="uma-s-1p2", device="auto")
# UMACalculator には Bohr 単位の座標(形状: [n_atoms, 3])を渡します
coords_bohr = np.array([
[0.0, 0.0, 0.0],
[2.2, 0.0, 0.0], # 約 1.16 Å
])
symbols = ["C", "O"]
# 注: これらのメソッドは dict を返すため、適切なキーで値を取り出します
energy_h = calc.get_energy(symbols, coords_bohr)["energy"] # float (Hartree)
forces_h_bohr = calc.get_forces(symbols, coords_bohr)["forces"] # ndarray (Hartree/Bohr)
hessian_h_bohr2 = calc.get_hessian(symbols, coords_bohr)["hessian"] # torch.Tensor (Hartree/Bohr²)
座標は Bohr で与えます。内部で Å に変換して UMA で計算し、エネルギーと微分を Hartree、Hartree/Bohr、Hartree/Bohr² に戻します。
device="auto"は、CUDA が使えれば GPU を、使えなければ CPU を選びます。pysisyphus(同梱の最適化ライブラリ)の geometry オブジェクトに付けるか、上のように直接呼び出します。
Calculator ファクトリ¶
backends モジュールには、MLIP の calculator をプログラムから作るファクトリがあります。
from pdb2reaction.backends import create_calculator, create_ase_calculator
関数 |
説明 |
|---|---|
|
pysisyphus 用の MLIP calculator を作ります。受け付ける kwargs はバックエンドごとに違い、選んだバックエンドが受け付けないキーは警告を出して捨てます。ただし UMA 用のキー( |
|
ASE 用の calculator を作ります。受け付ける kwargs はバックエンドごとに違い、使えないキーは警告なしに捨てられます。UMA・ORB・MACE の calculator は構造ごとの電荷とスピンを |
from pdb2reaction.backends import create_calculator
# 解析 Hessian 付きの UMA の calculator
calc = create_calculator(
backend="uma",
charge=0,
spin=1,
device="auto",
hessian_calc_mode="Analytical",
)
返される calculator は pysisyphus の calculator の形を持ちます。get_energy、get_forces、get_hessian は (atoms: List[str], coords: np.ndarray) を受け取り、coords は Bohr 単位です。戻り値は "energy"(Hartree)、"forces"(Hartree/Bohr)、"hessian"(Hartree/Bohr²)を持つ dict です。固定原子の力は 0 になり、Hessian は動ける原子のブロック(return_partial_hessian=True)か、固定原子の行と列を 0 にした全体の行列です。
設定の一覧¶
calculator の主な引数です。YAML の calc 節のキーと同じです。xTB の溶媒のキーを含む calc のすべてのキーは YAML 設定の一覧 › calc にあります。
オプション |
説明 |
デフォルト |
|---|---|---|
|
MLIP バックエンド |
|
|
総電荷。YAML に書いたときだけ使い、 |
なし |
|
スピン多重度(2S+1) |
|
|
選んだバックエンドのモデル( |
|
|
MLIP の数値精度( |
|
|
UMA のバッチに記録するタスクの名前 |
|
|
|
|
|
並列の UMA 予測器(UMA だけ。ORB・MACE・AIMNet2 は警告を出して無視します) |
|
|
UMA の近傍の作り方の上書き |
|
|
固定する原子の番号。Python API では 0 始まり(CLI と YAML は 1 始まり) |
None |
|
Hessian の計算方式( |
|
|
全体の行列でなく、動ける原子の Hessian だけを返す |
|
|
Hessian を float64 で組み立てて返す |
|
|
Hessian を |
|
|
Hessian の計算時間の内訳を表示 |
|
|
Hessian の計算中の CUDA VRAM の使用量を表示(UMA だけ) |
|
開発者向け¶
バックエンドディスパッチャのパターン¶
from pdb2reaction.backends import create_calculator, create_ase_calculator
calc = create_calculator(
backend="uma", # one of: "uma", "orb", "mace", "aimnet2", "auto"
charge=0, spin=1,
device="cuda", workers=1,
model="uma-s-1p2",
)
# calc is a pysisyphus-compatible MLIPCalculator.
# ASE-based stages (e.g. DMF path optimization) use the ASE factory:
ase_calc = create_ase_calculator(backend="uma", model="uma-s-1p2", device="cuda")
pysisyphus を使う構造と経路の段は create_calculator(...) を、DMF(Direct Max Flux)のように ASE を使う段は create_ase_calculator(...) を使います。backend="auto" は UMA、ORB、MACE、AIMNet2 の順に試し、最初にインポートできたものを使います。YAML の calc.backend: auto も同じで、-b は auto を受け付けません。
ファイルマップ¶
ファイル |
役割 |
|---|---|
|
|
|
|
|
UMA(Meta FAIR の fairchem-core):autograd の Hessian |
|
Orb(Orbital Materials):precision と compile_model |
|
MACE:default_dtype |
|
AIMNet2:電荷を入力に取るモデル |
|
PySCF/GPU4PySCF の DFT/HF の calculator。段の間で SCF の状態を引き継ぎ、同じ座標の結果を再利用し、必要なら SCF のチェックポイントを書きます |
組み込みのバックエンドを独自の --backend 名で足すときは、CONTRIBUTING のレシピ 3.2「Add an MLIP backend」に従ってください。
使用上の注意点¶
ORB と MACE の
--precision fp32はスクリーニング専用です。有限差分 Hessian のノイズが増えるので、結果を使う前に n_imag を確かめてください。AIMNet2 は
--precision fp64にも--deterministicにも対応せず、どちらもエラーで止まります。AIMNet2 はモデルへの入力を float32 にし、力を PyTorch の決定論的モードの外にある独自の CUDA のコードで計算するので、力はビット単位で再現しません(エネルギーは再現します)。繰り返しの計算を一致させたいときは、UMA、ORB、MACE のいずれかで--deterministicを付け、同じ環境で 2 回実行して比べてください。
UMA で
--uma-workersを 2 以上にすると、--hessian-calc-mode Analyticalとは併用できず、エラーで止まります。並列の予測器は autograd のモデルを持たないためです。解析 Hessian には--uma-workers 1を、複数のワーカーにはFiniteDifferenceを使ってください。モデルの精度と Hessian の精度は別の設定です。エネルギーと力は常に float64 で返り、Hessian もデフォルトでは float64 で組み立てます。
calc.hessian_double: falseにすると、モデル本来の dtype(ふつうは float32)で返します。--precision fp64のときは Hessian も常に float64 になり、設定ファイルのhessian_double: falseは警告を出して上書きされます。CI のジョブや Python API の
create_calculatorでは、環境変数PDB2REACTION_STRICT_DETERMINISTIC=1で--deterministicと同じモードになります。
関連ドキュメント¶
アーキテクチャ:ディレクトリの構成と依存の向き
HPC 実行例:PBS、Open MPI、Ray で
workers・workers_per_nodeを複数のノードに広げるテンプレートMLIP の TS を DFT で確かめる:DFT の設定、メモリ、チェックポイント
トラブルシューティング:計算が失敗したとき