QEのpw.xによるMD計算結果をDeePMD-kitで使用する Part.1 データ準備編

機械学習力場のツールであるDeePMD-kitでは学習データに第一原理MDの計算を使用できる…はずなのですが、QEの出力をいい感じに読み込んでくれないし、やり方を調べてもあまり出てこない。ということで、PythonのASEを中間に挟む方法でいい感じに読み込めましたのでやり方を解説します。

なお、Dockerを使ったDeePMD-kitの環境構築はこちらで解説しています。

QEのpw.xによるMD計算

「DeePMD-kitで作成した力場で結晶の密度を再現する」という目的で進めていきます。計算速度が速い適当なサンプルということで、炭素のダイヤモンド結晶を使用します。モデルはASEで作ります。

from ase.build import bulk
from ase.visualize import view
from ase.units import mol

diamond = bulk("C", "diamond", cubic=True)

n=2
unitcell = diamond.repeat((n, n, n))
unitcell.write("diamond.in", format='espresso-in')

density = (np.sum(unitcell.get_masses())/mol) / (unitcell.get_volume()*1e-24)
print(density, "[g/cm3]")
view(unitcell)

これで作られるダイヤモンド結晶は原子数が64個、密度は3.51g/cm³です。

そしてdiamond.inを開いて編集します。編集部分は赤字のところです。

&CONTROL
   calculation='vc-md'
   dt=20
   nstep=200
   pseudo_dir='../'
/
&SYSTEM
   ntyp             = 1
   nat              = 64
   ibrav            = 0
   nosym=.true.
/
&ELECTRONS
/
&IONS
   ion_temperature='initial'
   tempw=300.0
/
&CELL
   cell_dynamics='pr'
   press=0.001
/

ATOMIC_SPECIES
C 12.011 C.pbe-n-kjpaw_psl.1.0.0.UPF

K_POINTS gamma

CELL_PARAMETERS angstrom
7.14000000000000 0.00000000000000 0.00000000000000
0.00000000000000 7.14000000000000 0.00000000000000
0.00000000000000 0.00000000000000 7.14000000000000

ATOMIC_POSITIONS angstrom
C 0.0000000000 0.0000000000 0.0000000000 
以下略

精度を何も考慮していない入力ファイルですので、真似する際はいろいろ書き加えてください。

そして計算します。

mpiexec -n 8 pw.x -in diamond.in |tee diamond.out

「登場する元素が軽い」「k_pointsがgamma」という理由で、64個という原子数の割には計算速度が速いはずです。またEstimated total dynamical RAMは578MBでした。
計算の標準出力はdiamond.outというファイル名で出てきます。

なお、cp.xのMD計算も使用できます。その際に気を付けるポイントは以下の記事にまとめました。

ASEでファイル変換

DeePMD-kitはQEのファイルに対応していますが、pw.xのファイルはSCFの計算結果しか読めず、MDの計算結果は1フレームしか読んでくれません。（cp.xのほうだとトラジェクトリで読んでくれます 関連記事）
DeePMD-kitはASEのtrajectoryファイルにも対応しており、QEのMD計算結果をASEのtrajectoryに変換し、それをDeePMD-kitに読み込ませることができます。今回はこの方法で進めます。

ASEでのデータ変換は簡単です。出力ファイル名を diamond.traj とします。

from ase.io import read, write

qe_out    = "diamond.out"
traj_name = "diamond.traj"
traj      = read(qe_out, format="espresso-out", index=":")
write(traj_name, traj,  format="traj")

注意点として、QEの時系列データの読み込みは index=”:” の指定が必要です。indexの指定方法の詳細は公式のドキュメントをご確認ください。

なお、このindexの指定方法によっては学習データの間引きが可能です。（偶数stepのみ使用する、10step毎に使用する等）

dpdataで変換

出力した diamond.traj を dpdata で変換します。
dockerの中で変換をする場合は
docker cp ./diamond.traj deepmd_kit:/root/diamond/
のような形で起動中のコンテナにコピーします。

その後、以下のPythonコードを実行します。（公式チュートリアルとほぼ同じです）

import dpdata

data = dpdata.MultiSystems.from_file(file_name="./diamond.traj", fmt="ase/structure")
data_training, data_validation, dict= data.train_test_split(40)
data_training.to_deepmd_npy('training_data')
data_validation.to_deepmd_npy('validation_data')

出力は training_data/C64の中に set.000 type.raw type_map.raw が生成されます。set.000の中には box.npy coord.npy energy.npy force.npy virial.npy があります。validation_dataについても同様です。
Python実行時にdataの中身を確認すると以下のようになります。

>>> print(data.systems)
{'C64': Data Summary
Labeled System
-------------------
Frame Numbers      : 200
Atom Numbers       : 64
Including Virials  : Yes
Element List       :
-------------------
C
64}
>>> print(data.systems["C64"].data)
{'atom_numbs': [64], 'atom_names': ['C'], 'atom_types': array([0, 0, 0, 0,,
 'cells':省略
 'coords':省略
 'nopbc': False
 'energies':省略
 'forces':省略 
 'virials':省略

input.jsonの作成

次にdp trainで使用するinput.jsonを作ります。公式チュートリアルを少し書き換えただけですので、正確性や効率は保証しません。

{
    "_comment": " model parameters",
    "model": {
        "type_map":     ["C"],
        "descriptor" :{
            "type":             "se_e2_a",
            "sel":              [158],
            "rcut_smth":        0.50,
            "rcut":             6.00,
            "neuron":           [10, 20, 40],
            "resnet_dt":        false,
            "axis_neuron":      4,
            "seed":             1,
            "_comment":         " that's all"
        },
        "fitting_net" : {
            "neuron":           [100, 100, 100],
            "resnet_dt":        true,
            "seed":             1,
            "_comment":         " that's all"
        },
        "_comment":     " that's all"
    },

    "learning_rate" :{
        "type":         "exp",
        "decay_steps":  5000,
        "start_lr":     0.001,
        "stop_lr":      3.51e-8,
        "_comment":     "that's all"
    },

    "loss" :{
        "type":         "ener",
        "start_pref_e": 0.02,
        "limit_pref_e": 1,
        "start_pref_f": 1000,
        "limit_pref_f": 1,
        "start_pref_v": 0,
        "limit_pref_v": 0,
        "_comment":     " that's all"
    },

    "training" : {
        "training_data": {
            "systems":          ["./training_data/C64"],
            "batch_size":       "auto",
            "_comment":         "that's all"
        },
        "validation_data":{
            "systems":          ["./validation_data/C64"],
            "batch_size":       "auto",
            "numb_btch":        1,
            "_comment":         "that's all"
        },
        "numb_steps":   1000000,
        "seed":         10,
        "disp_file":    "lcurve.out",
        "disp_freq":    1000,
        "save_freq":    10000,
        "_comment":     "that's all"
    },

    "_comment":         "that's all"
}

selの値 [158] は以下を参考に決めました。

機械学習の実行

ここからはいよいよ機械学習の処理を進めていきます。