[自由研究に]ML-Potential_CHGNet,

Localで動かせる機械学習ポテンシャル、力場の紹介も兼ねたメモです。安価なPCでも動かせるため、中高生の方にぜひ研究者気分を少し味わってみて欲しいです。

紹介するのは少し前のものになりますが、下記のCederらのCHGNetです。

CHGNet as a pretrained universal neural network potential for charge-informed atomistic modelling - Nature Machine Intelligence

しばらくMatbenchで一番上に記載されていました。

Matbench leaderboead

機械学習(ML)ポテンシャルを使った分子動力学(MD)シミュレーションは低スペックのPCでも手軽に動かせるものが多いものの、性能的にはまだまだ発展途上だと思います。

一方で低スペックのPCで動かせることは非常に重要で、人の理解は頭よりも体感が重要なことが殆どであるらしく、安価なPCで実際に自分で操作ができることは集団の認知・理解の底上げの大きな助けになります。

0. 環境

OS：Windows
CPU：Intel(R) Core i9-13900KF
RAM：128GB
GPU：RTX 4090

1. cif読込み

Materials Project(MP)からNa3V2(PO4)3のcifデータをダウンロードします。
論文ではNa4V2(PO4)3からNaをランダムに半分削除していますが、Na4V2(PO4)3がMPになかったので、Na3V2(PO4)3のNaを1/3削除します。

from pymatgen.core import Structure

filename = "Na3V2(PO4)3"
structure = Structure.from_file(f"{path}/{filename}.cif")
supercell = structure.make_supercell([2, 2, 2], in_place=False)

import random
n_Na = int(supercell.composition["Na+"])
remove_ids = random.sample(list(range(n_Na)), n_Na // 3)

supercell.remove_sites(remove_ids)
print(supercell.composition)

Na+64 V3+64 P5+96 O2-384

2. LBFGSLineSearchで構造最適化

Atomic Simulation Environment(ASE)のLBFGSLineSearchを使って最適化します。ポテンシャルは上述のCHGNetを使います。ちなみにパラメータ数は412, 525でした。

from chgnet.model.dynamics import CHGNetCalculator
from ase.optimize import LBFGSLineSearch
from pymatgen.io.ase import AseAtomsAdaptor

pot = CHGNetCalculator(
    potential=chgnet,
    properties='energy',
    compute_stress=True,
    compute_hessian=False,
    stress_weight=0.01,
    )

outpath = r".\MD\231123_NVP"
atoms_0 = AseAtomsAdaptor().get_atoms(supercell)
atoms = atoms_0.copy()
atoms.calc = pot
qn=LBFGSLineSearch(atoms, trajectory=f"{outpath}/{supercell.composition}.traj")
qn.run(fmax=0.05)

3. 磁気モーメント

論文同様、Vの磁気モーメントを可視化してみます。

# 最適化前
Unrelaxed_supercell = AseAtomsAdaptor.get_structure(atoms_0).make_supercell([2, 2, 2], in_place=False)
s_object = Unrelaxed_supercell.as_dataframe().Species
s = s_object.astype(str).str.contains('V')
s[s].index
m = chgnet.predict_structure(Unrelaxed_supercell)["m"][s[s].index]

# 最適化後
Relaxed_supercell = AseAtomsAdaptor.get_structure(traj[-1]).make_supercell([2, 2, 2], in_place=False)
m2 = chgnet.predict_structure(Relaxed_supercell)["m"][s[s].index]

# 可視化
df_mag = pd.DataFrame({"Unrelaxed":m, "CHGNet relaxed":m2})
fig = df_mag.hist(
    nbins=200,
    sharex=True,
    sharey=True,
    backend="plotly",
    barmode="overlay",
    layout={"title": "Magmom distribution"},
    opacity=0.7,
    range_x=[0, 5],
    template="plotly_white",
)
fig.layout.legend.update(title="", x=1, y=1, xanchor="right", yanchor="top")
fig.layout.xaxis.title = "Magnetic moment"
fig.show()

青が最適化前、赤が最適化後の磁気モーメント(CHGNetによる予測)分布になっています。
CHGNetのポテンシャルを使った最適化で全体的にV3+-likeだったものからV4+-likeなものが増えている様子が見て取れます。

いくつかのV原子の最適化中の磁気モーメントを見てみます。

mag_list = []
for st in tqdm(traj):
    Relaxed_supercell = AseAtomsAdaptor.get_structure(st).make_supercell([2, 2, 2], in_place=False)
    mag = chgnet.predict_structure(Relaxed_supercell)["m"][s[s].index]
    mag_list.append(mag)
mag_map = np.array(mag_list)

import plotly.express as px
fig = px.line(mag_map[:, ::80])
fig.show()

横軸が最適化中の経過、縦軸がCHGNetで予測される磁気モーメントです。
当初、すべてのVの磁気モーメントは2付近で、その中の一部が1.2前後にすぐ落ち着いている様子がわかり摂動のような現象は生じていないように見えます。

4. まとめ

CHGNetは磁気モーメントを予測可能な初めてのML-ポテンシャルで、パラメータ数も比較的少なく安価なPCでそこそこ早く計算が可能です。
精度はさておき、磁気モーメントを描写できるのもまた楽しいですね。

どこかでその妥当性についても投稿出来たらいいなぁと思ってます。