自己学習ハイブリッドモンテカルロ法を試してみる

自己学習ハイブリッドモンテカルロ法は機械学習分子動力学法におけるポテンシャル作成を自動化しつつ計算精度が元のシミュレーションと同等になるという手法です。どのように使うか、ということについて、紹介したいと思います。英語の原論文はこちら(出版版とプレプリント版)。紹介スライドは例えばこちらにあります。

何をするか

SiO2の原子間相互作用をニューラルネットワークで模擬します。具体的には、第一原理計算ソフトウェアであるQuantum Espressoで計算されたエネルギーを模擬するニューラルネットワークを構築し、自動微分によって力を計算して、MDを動かすことができるようになります。ニューラルネットワークを学習する際に、自己学習ハイブリッドモンテカルロ法を用います。

背景

分子動力学計算(MD)を精度良く行うためには、力を精度良く求める必要があります。また、エネルギーの微分が力です。そのため、エネルギーを精度良く求めれば精度の良いMDが実行できます。そこで、第一原理計算によってエネルギーを計算して、その微分を計算することでMDを実行する、第一原理分子動力学法という手法が使われてきました。しかし、この手法は時間ステップの刻み(1フェムト秒程度)ごとに第一原理計算を行うため、非常に計算が重いです。そこで、このエネルギーの計算を行う部分をニューラルネットワークに置き換えるという、機械学習分子動力学法（MLMD）が登場しました。この手法は非常に速く、さまざまな場所で使われています。
MLMDは非常に良い手法ですが、問題の一つとして、ニューラルネットワークの訓練をどうやるか、というところがあります。

使うソフトウェア

• Docker+PIMD+aenet+Quantum Espresso

こちらのPIMDとAENETを使って機械学習分子シミュレーションをやってみよう: 2021年Docker版を参考にしてください。手順としては、Dockerをインストールした後

docker pull cometscome/pimd_aenet_qe:pimd_2_6_0

でイメージをpullしておいてください。

また、PIMDのexamplesディレクトリの中身を使うため、こちらから2.6.0をダウンロードし、解凍しておいてください。

• VESTA

構造の可視化とデータのフォーマットの変換用にVESTAを使います。インストールしておきましょう。

フォルダ

以後は、SiO2というディレクトリを作成し、その中で作業を行います。Dockerを走らせる際にはこのディレクトリを共有フォルダとして使用しますので、例えば、SiO2ディレクトリに移動した状態で

docker run -v ./:/SiO2 --rm -it cometscome/pimd_aenet_qe:pimd_2_6_0

のように実行することになります。なお、WindowsやMacでは発生しませんがLinuxの場合パーミッションがらみでエラーがでることがあります。その場合は、

 docker run -v ./:/SiO2 -u "$(id -u $USER):$(id -g $USER)" \
    -v /etc/passwd:/etc/passwd:ro \
    -v /etc/group:/etc/group:ro --rm -it cometscome/pimd_aenet_qe:pimd_2_6_0

のようにしてください（参考）。また、この場合には、

source /opt/intel/oneapi/setvars.sh

をコンテナの実行時に行うことでintel MPIを有効にする必要があります。また、Docker内でのMPIを実行する場合、Linuxの場合にはどうやらホスト側でもintel MPIが入っている（有効化されている）必要があるようです。

SiO2ディレクトリに移動後

mkdir pseudo
cd pseudo
wget --no-check-certificate http://pseudopotentials.quantum-espresso.org/upf_files/Si.pz-vbc.UPF 
wget --no-check-certificate https://pseudopotentials.quantum-espresso.org/upf_files/O.pz-kjpaw.UPF

として、今回使う擬ポテンシャルをダウンロードしておいてください。もしwgetがない、と言われた場合には、macではbrewなどで入れるか、curlでダウンロードしてください。

対象とする物質

SiO2とします。この物質のSiとOに関するニューラルネットワークを作ってみましょう。SiO2の構造は色々ありますが、ここではα-quartzとします。

構造ファイルの入手

まず、結晶構造を入手します。これには、Materials Projectが便利です。α-quartzはこちらです。構造データはcifファイルというフォーマットを使います。Materials Projectのユーザー登録（無料）をして、

図のように、アイコンの一番下のファイルのエクスポートを選び、CIFファイルを選んでダウンロードしてください。これでSiO2.cifというファイルを入手できます。このファイルをSiO2ディレクトリに入れておきます。

構造の確認とデータフォーマット変換

このファイルをVESTAで開きます。その後、Export Dataを選び、VASPファイルを選択、「Cartesian coordinates」を選んでください。このファイルをPIMDに読み出せるようにstructure.datに変換するのですが、このPIMDのフォーマットはxyzフォーマットです。毎回変換するのも面倒ですので、Juliaのコードを作りました。Julia言語はPythonのように気楽にかけてFortranやcのように高速に書ける言語です（弊著も参考にしてみてください）。

Juliaは、Windowsであれば

winget install julia -s msstore

MacあるいはLinuxであれば

curl -fsSL https://install.julialang.org | sh

で簡単にインストールできます。

henkan.jl

function test(filename)
    bohr = 0.529177210903
    data = readlines(filename)
    num = countlines(filename)

    u = split(data[2])
    factor = parse(Float64,u[1])
    u1 = parse.(Float64,split(data[3]))*factor/bohr
    u2 = parse.(Float64,split(data[4]))*factor/bohr
    u3 = parse.(Float64,split(data[5]))*factor/bohr
    println("input.dat for PIMD")
    println("""
    <iboundary>
    1
    $(u1[1])\t$(u2[1])\t$(u3[1])
    $(u1[2])\t$(u2[2])\t$(u3[2])
    $(u1[3])\t$(u2[3])\t$(u3[3])
    """)

    u = split(data[6])
    atoms = []
    for i=1:length(u)
        push!(atoms,u[i])
    end
    u = split(data[7])
    numatoms = []
    println(u)
    println(length(u))
    for i=1:length(u)
        push!(numatoms,parse(Int64,u[i]))        
    end
    println(numatoms)
    totalatoms = sum(numatoms)
    fp = open("structure.dat","w")
    println(fp,totalatoms)
    println(fp,"ANGSTROM")
    count = 8

    for i=9:9+totalatoms-1
        
        for is = 1:length(numatoms)
            for iatom = 1:numatoms[is]
                count += 1
                u = split(data[count])
                r = parse.(Float64,u)
                println(fp,atoms[is],"\t",data[count],"\t",is)
            end
        end
        
        if count == num
            close(fp)
            exit()
        end
    end


    
end

filename = ARGS[1]
test(filename)

こちらのファイルをSiO2ディレクトリにhenkan.jlという名前で保存し、

julia henkan.jl SiO2.vasp

と実行してください。出力結果は、

input.dat for PIMD
<iboundary>
1
9.287939852914283	-4.6439699264571415	0.0
0.0	8.043591861479895	0.0
0.0	0.0	10.263671611315054

SubString{String}["3", "6"]
2
Any[3, 6]

となります。ここに出てきたiboundaryの部分はinput.datで使うことになります。
これでstructure.datというファイルが生成されます。また、SiO2.vaspを開きますと、

Si3 O6
1.0
        4.9149661064         0.0000000000         0.0000000000
       -2.4574830532         4.2564855069         0.0000000000
        0.0000000000         0.0000000000         5.4313011169

のようにありますが、これがユニットセルの情報です。henkan.jlでは、input.datのフォーマットがvaspフォーマットの転置になっているので転置し、ユニットがボーア長なのでオングストロームからボーアに直しています。

第一原理MDの実行

次に、ニューラルネットワークを作成するための最初のデータを生成します。SiO2ディレクトリに新しくMDディレクトリを作成します。次に、input.datを作成します。ゼロから作るよりもexamplesディレクトリからコピーした方が楽ですので、解凍したPIMDディレクトリから、examples/SiO2/qe_md/input.datをSiO2ディレクトリにコピーしてください。また、examples/SiO2/qe_md/input_default.datもコピーしておいてください。

input.datは

<method>
PIMD

<nbead>
1


<iprint_rdf>
10

<ensemble>
NVT

<bath_type>
MNHC

<temperature>
300.0


<ipotential>
QE

<nstep>
100

<dt>
1.0


<iboundary>
  1
  9.232239523 -4.170626664  0.000000000
  0.000000000  7.223738307  0.000000000
  0.000000000  0.000000000  9.132346975


<qe_output>
1 10

<qe_input_file_name>
qe.dat

<iprint_xsf_aenet>
1
<istep_train_aenet>
-1

のようになっています。はステップ数で、今は100ステップ行うことになっています。<iboundary>の部分を先ほどhenkan.jlで出力されたものに置き換えます。また、学習用の結晶構造とエネルギーのデータを書き出すために<istep_train_aenet>を-1以外に書き換えます。例えば最後にトレーニングをするのであれば、100にしておきましょう。また、訓練をする際の最小の構造数は<minxsf_train_aenet>で定義されていまして、デフォルトでは300になっています。デフォルト値がどうなっているかを見るには、input_default.datを見てください。ここでは、この値を100としておきます。

<qe_input_file_name>はQuantum Espressoの設定ファイルが何かを指定するものでして、今はqe.datになっています。examples/SiO2/qe_md/qe.datをコピーしてください。Quantum Espressoの計算の設定をいじりたい場合にはここをいじってください。今はexamplesと基本的に同じということにします。ただし、気をつけなければならない点は、qe.datの中で

pseudo_dir = '../../../lib/qe/pseudo/',

というところがありますが、ここは擬ポテンシャルの保存場所を示しているところです。ここではSiO2ディレクトリにpseudoディレクトリを作成し、

pseudo_dir = '../pseudo/',

としておきます。

結果、

<method>
PIMD

<nbead>
1


<iprint_rdf>
10

<ensemble>
NVT

<bath_type>
MNHC

<temperature>
300.0


<ipotential>
QE

<nstep>
100

<dt>
1.0


<iboundary>
1
9.287939852914283	-4.6439699264571415	0.0
0.0	8.043591861479895	0.0
0.0	0.0	10.263671611315054


<qe_output>
1 10

<qe_input_file_name>
qe.dat

<iprint_xsf_aenet>
1

<istep_train_aenet>
100

<minxsf_train_aenet>
100

となります。MDディレクトリの中身は

pimder@4c27654ebd0e:/SiO2/MD$ ls
input.dat  input_default.dat  qe.dat  structure.dat

となっているはずです。ここまでデータ収集をする分には問題ありません。100ステップ後に訓練する場合には、aenetに関するパラメータを設定する必要があります。

訓練用パラメータの設定

訓練用のパラメータはexamples/SiO2/qe_slhmcを参考にしてください。まず、必要なファイルはSiとOのstpファイルです。qe_slhmcのディレクトリから二つの.stpをMDディレクトリにコピーしておきます。stpファイルは原子の周辺環境をどのように記述するか、というものです。例えば、

DESCR
  N. Artrith and A. Urban, Comput. Mater. Sci. 114 (2016) 135-150.
  N. Artrith, A. Urban, and G. Ceder, Phys. Rev. B 96 (2017) 014112.
END DESCR

ATOM O

ENV 2
Si
O

RMIN 0.55d0

BASIS type=Chebyshev
radial_Rc = 8.0  radial_N = 16 angular_Rc = 6.5  angular_N = 4

のようなものです。ここで、ENVには今考えている原子の数と種類を書きます。また、RMINは考慮すべき原子の最小距離です。また、BASIS type=Chebyshevはチェビシェフ多項式基底を使うことを指示しています。radial_Rc = 8.0 radial_N = 16 angular_Rc = 6.5 angular_N = 4の部分は、radial_Rcは二体の情報に関する動径方向の長さカットオフで、この半径より遠い原子は考慮に入れません。anbular_Rcは三体の情報に関する長さカットオフです。radial_Nは二体の情報をチェビシェフ多項式展開した時の最大次数です。angular_Nは３個の原子のなす角度方向のチェビシェフ多項式展開の最大次数です。

次に、input.datに追記します。追記するのは、

<generate_aenet>
OUTPUT SiO2.train
TYPES
2
Si 0.0   | eV
O   0.0  | eV
SETUPS
Si  Si.fingerprint.stp
O    O.fingerprint.stp

です。ここでstpファイルを指定しています。また、Si 0.0 | eVの部分はSiが孤立していた場合のエネルギーを指定していますが、学習データの原子数が同じ場合には、ここは0にしておいて問題ありません。訓練データの中に原子数の異なるものが入っている場合には、ここを使っている第一原理計算による孤立原子のエネルギーを入れておかないと、訓練がうまくいきません。

また、訓練で作るニューラルネットワークの情報は

<train_aenet>
TRAININGSET SiO2.train
TESTPERCENT 10
ITERATIONS  5000
MAXENERGY 1.0
TIMING
METHOD
bfgs
NETWORKS
! atom   network         hidden
! types  file-name       layers  nodes:activation
Si     Si.15-15.ann    2      15:tanh 15:tanh
O       O.15-15.ann    2      15:tanh 15:tanh

に記述します。詳しくはaenetのマニュアルを見てください。ここでは、隠れ層2層でユニット数が15のニューラルネットワークを設定しています。
これらをinput.datに追記してください。

ユニットセルサイズについて

ユニットセルサイズは、通常、ニューラルネットワークで考慮する周辺原子環境のカットオフよりも大きい方がいいです。つまり、周辺原子環境を見る球がユニットセル内に一つは入る程度のサイズがないと、周期境界条件のアーティファクトを引っ掛ける可能性があり、より大きな系への適用が難しい可能性があります。ここでは試すためにサイズは小さくしています。

リスタートについて

PIMDを実行する場合、.iniファイルがある場合はリスタートができます。一方、まっさらな状態からリスタートしたい場合には.iniと.outを消せば良いです

データ収集

さて、準備はできましたので、あとはMDディレクトリでPIMDを実行すればいいということになります。SiO2ディレクトリに移動し、

docker run -v ./:/SiO2 --rm -it cometscome/pimd_aenet_qe:pimd_2_6_0

を実行するとPIMD+aenet+QEの環境に入ることができます。その後、

cd /SiO2/MD

でMDのディレクトリに移動し、

mpirun -np 2 pimd.mpi.x

などとすれば、2並列でPIMDが動きます。実行すると、structuresディレクトリがMDディレクトリの直下に生成されます。そして、その中は

structures % ls
00000001.xsf	00000003.xsf	00000005.xsf	00000007.xsf	00000009.xsf	00000011.xsf	00000013.xsf	00000015.xsf	00000017.xsf
00000002.xsf	00000004.xsf	00000006.xsf	00000008.xsf	00000010.xsf	00000012.xsf	00000014.xsf	00000016.xsf	00000018.xsf

のようになっています（ここでは18構造まで入っています）。並列数は2だと遅すぎるなと思った場合は、適宜ご自分の計算機に合わせて選ぶとよいでしょう。訓練の途中経過はtrain.scrに書き出されていますので、そちらを見るとlossがどうなっているかなどがわかるかと思います。計算時間はtrain.timeに書き出されています。

計算した結果が trained_networksにできます。ここでは訓練の各ステップでのニューラルネットワークファイルannファイルが格納されているはずです。ここで、番号の付いていないO.15-15.annとSi.15-15.annが最終的なファイルです。

SLHMCの実行

次に、作ったannファイルを用いてSLHMCを行ってさらにニューラルネットワークを改善します。SiO2のディレクトリに新しくSLHMCディレクトリを作ります。この時、

cp -r MD SLHMC

とMDをコピーします。そして、SLHMCディレクトリに移動し、

rm *.out
rm *.ini
cp ../MD/aenet.ini ./
cp trained_networks/O.15-15.ann ./
cp trained_networks/Si.15-15.ann ./

とaenet.ini以外のiniファイルを削除します。aenet.iniファイルはstructuresに格納するデータの連番情報を保持していますので、これを残しておくと次の計算のファイル名が続きになります。
次に、input.datをSLHMCをするように変更します。examples/SiO2/qe_slhmcにあるinput.datを参考にして変更します。

具体的には、

<method>
PIHMC

<nbead>
1


<iprint_rdf>
10

<ensemble>
NVT

<bath_type>
MNHC

<temperature>
300.0


<ipotential>
DUAL

<dual_potential>
QE  AENET

<nstep>
1000

<dt>
0.25


<iboundary>
1
9.287939852914283	-4.6439699264571415	0.0
0.0	8.043591861479895	0.0
0.0	0.0	10.263671611315054


<qe_output>
1 10

<qe_input_file_name>
qe.dat

<iprint_xsf_aenet>
1

<minxsf_train_aenet>
100

<istep_train_aenet>
100

<istep_adjust_hmc>
100

<istep_hmc>
2

<ntype_aenet>
2
Si Si.15-15.ann
O   O.15-15.ann

<generate_aenet>
OUTPUT SiO2.train
TYPES
2
Si 0.0   | eV
O   0.0  | eV
SETUPS
Si  Si.fingerprint.stp
O    O.fingerprint.stp

<train_aenet>
TRAININGSET SiO2.train
TESTPERCENT 10
ITERATIONS  2000
MAXENERGY 1.0
TIMING
METHOD
bfgs
NETWORKS
! atom   network         hidden
! types  file-name       layers  nodes:activation
Si     Si.15-15.ann    2      15:tanh 15:tanh
O       O.15-15.ann    2      15:tanh 15:tanh

とします。MDディレクトリのinput.datと比較すると何を追加したかわかるかと思います。いくつかポイントとしましては、aenetによるMDを実行するためにポテンシャルの指定を

<ntype_aenet>
2
Si Si.15-15.ann
O   O.15-15.ann

で指定していること、

<ipotential>
DUAL

<dual_potential>
QE  AENET

でQuantum Espressoとaenetを使うと指定していること、でしょうか。<dt>は今度はaenetによるMDの時間刻みになりますので、0.25と小さめを指定しています。

あとはこれを実行して、しばらく待ってみてください。aenetによるMDの長さはアクセプト率に応じて自動調整するようになっています。どのように調整されたかは、hmc.outに出力されています。
今回の場合、例えば、hmc.outは

================================
    step   moves   last%  total%
--------------------------------
     100       2 0.85000 0.85000
     200       4 0.79000 0.82000
     300       8 0.80000 0.81333
     400      16 0.56000 0.75000
     500      32 0.34000 0.66800
     600      64 0.17000 0.58500
     700      64 0.51000 0.57429
     800     128 0.49000 0.56375
     900     128 0.69000 0.57778
    1000     128 0.74000 0.59400

のようになっています。ここで、movesはMDステップです。右の二つはアクセプト率です。最後は128ステップに１度QEをやっているので、これは0.25fs*128=32fsごとに1度QEをやっています。

機械学習MDのテスト

できあがった機械学習ポテンシャルを用いて、機械学習MDをやってみましょう。ディレクトリとして新しくMLMDを作り、そこにファイルを入れます。
MLMDの場合はinput.datは

<method>
PIMD

<nbead>
1

<iprint_rdf>
10

<ensemble>
NVT

<bath_type>
MNHC

<temperature>
300.0

<ipotential>
AENET

<nstep>
10000

<dt>
1.0

<iboundary>
1
9.287939852914283	-4.6439699264571415	0.0
0.0	8.043591861479895	0.0
0.0	0.0	10.263671611315054


<ntype_aenet>
2
Si Si.15-15.ann
O   O.15-15.ann

となります。必要なファイルはannファイルとstructure.datとinput.datとinput_default.datです。実行してみると高速に動くことがわかると思います。