aenetとASEを用いてANNポテンシャルによるMDシミュレーションをやってみる #Python

aenet (Atomic Energy Network) はANN (人工ニューラルネットワーク) 原子間相互作用ポテンシャルを扱うためのパッケージです．DFT計算をもとに新規ANNポテンシャルを構築したり，学習済みポテンシャルを用いてエネルギーや力を予測したりすることができます．

aenetは原子シミュレーション用のPythonライブラリであるASE (Atomic Simulation Environment) のインターフェースを提供しているため，ANNポテンシャルを用いた構造最適化や分子動力学シミュレーション (MD) を行うことができます．

本記事では学習済みのANNポテンシャルを用いて，TiO₂のMDシミュレーションを行ってみます．

インストール

aenetのPythonインターフェースをビルドします．

プロジェクトの各ディレクトリ (lib, src, python3) でmakeを実行します．srcディレクトリには複数のMakefile (ifort vs. gfortran, mpi vs. serialなど) が用意されています．必要に応じてMakefileを選択してください．

各ディレクトリにREADMEが用意されているので，こちらも参照してください．

$ git clone https://github.com/atomisticnet/aenet.git
$ git checkout v2.0.4 # 最新バージョン
$ cd aenet/lib
$ make # libをコンパイル
$ cd ../src
$ make -f makefiles/Makefile.gfortran_serial lib # srcをコンパイル. ターゲットは必ずlibにする．
$ cd ../python3
$ python3 setup.py install --user # Pythonインターフェースをコンパイル

なお，libやsrcをビルドする際のgccのバージョンとpythonインターフェースをビルドする際のgccのバージョンが一致していない場合，ビルドは完了しても正しく動作しません．私はこれでハマりました．

学習済みデータ

公式サイトからTiO₂の学習済みデータ (aenet-example-02-TiO2-Chebyshev.tar.bz2) がダウンロード可能です．本記事でもこの学習済みデータを用います．

以下ではaenet-example-02-TiO2-Chebyshev/03-predict/set001ディレクトリ内の学習済みnnファイルを用いて計算を行います．

簡単な例

もっとも簡単な例としてルチル型TiO₂のエネルギーを計算してみます．

from ase.spacegroup import crystal
from aenet.ase_calculator import ANNCalculator

TiO2_rutile =crystal(['Ti', 'O'], basis=[(0, 0, 0), (0.3, 0.3, 0.0)],
                     spacegroup=136, cellpar=[4.6, 4.6, 2.95, 90, 90, 90])
calc = ANNCalculator({'Ti':'Ti.15t-15t.nn', 'O':'O.15t-15t.nn'})
TiO2_rutile.set_calculator(calc)
e = TiO2_rutile.get_potential_energy()
print(f'energy: {e}')

とても簡単にANNポテンシャルのエネルギーが計算できました．

ベンチマーク

ノートPCで計算時間を計測してみます．原子数の異なるスーパーセルを複数個作成して，それぞれエネルギー計算に要した時間をプロットしました．

MDシミュレーション

最後にMDシミュレーションを行ってみます．
ASEのCalculatorをANNCalculatorに置き換えるだけです．

from ase import units
from ase.spacegroup import crystal
from ase.md.nvtberendsen import NVTBerendsen
from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
from aenet.ase_calculator import ANNCalculator

dt = 2 * units.fs
temp = 300
nsteps = 400
taut = 20 * units.fs

TiO2_rutile =crystal(['Ti', 'O'], basis=[(0, 0, 0), (0.3, 0.3, 0.0)],
                     spacegroup=136, cellpar=[4.6, 4.6, 2.95, 90, 90, 90])
calc = ANNCalculator({'Ti':'Ti.15t-15t.nn', 'O':'O.15t-15t.nn'})
TiO2_rutile.set_calculator(calc)
MaxwellBoltzmannDistribution(TiO2_rutile, temp * units.kB)
dyn = NVTBerendsen(TiO2_rutile, dt, temp, taut=taut, trajectory='md.traj')
def myprint():
    print(f'time={dyn.get_time() / units.fs: 5.0f} fs ' + \
          f'T={TiO2_rutile.get_temperature(): 3.0f} K')
dyn.attach(myprint, interval=20)
dyn.run(nsteps)

参考

aenet
- ホームページ http://ann.atomistic.net/
- GitHub https://github.com/atomisticnet/aenet
- N. Artrith and A. Urban, Comput. Mater. Sci. 114 (2016) 135-150. http://dx.doi.org/10.1016/j.commatsci.2015.11.047
Qiita
- PIMDとAENETを使って機械学習分子シミュレーションをやってみよう https://qiita.com/cometscome_phys/items/660a84119b5d4acce903