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


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

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

本記事では学習済みのANNポテンシャルを用いて,TiO2の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のバージョンが一致していない場合,ビルドは完了しても正しく動作しません.私はこれでハマりました.

学習済みデータ

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

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

簡単な例

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

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)

参考