【DP系列教程-8】DeePMD-kit拟合Dipole / Polarizability
众所周知，DeePMD-kit 在拟合势能面上有良好的表现。然而许多小伙伴不知道的是，DeePMD-kit 其实提供了多种 fitting_net 的模式。今天给大家介绍如何使用 DeePMD-kit 拟合Tensor，即Dipole[1]以及Polarizability[2]（下称 polar）。在数值上，dipole / polar 就是一个三维向量 / 三阶矩阵（并满足一定的物理约束）。一个完整的体系会有一个全局的 dipole / polar，而我们也可以人为分解为每一个原子上都有一个 atomic dipole / polar。为了保证拟合的精度，DeePMD-kit 采用 atomic tensor 的模式。如果需要拟合全局值，只需要将其对原子求和即可。
数据准备
Label 文件
有些系统可能提供 atomic label，而有些则只针对每一个体系（frame）提供了一个全局的 label。幸运的是，DeePMD-kit 同时支持对 global label、atomic label 的拟合，也支持对多个 label 类型不同的体系同时拟合。类似于ener模式，我们需要提供coord.raw，type.raw，box.raw。此外，我们还需要根据拟合的 tensor 类型和该系统的 label 类型，提供dipole.raw或是atomic_dipole.raw（类似的，提供polarizability.raw或atomic_polarizability.raw），这里是否添加前缀atomic_取决于该系统的 label 是否对每一个原子都有 label。如有需要，参见https://deepmd.readthedocs.io/en/master/use-deepmd-kit.html#prepare-data。
Json 修改
准备好数据后，我们需要修改 json 文件。首先需要修改fitting_net。根据拟合的类型不同，将type字段修改为dipole或polar（注意：global_polar模式已经融入polar模式中）。
"fitting_net": {
"type": "dipole",（或是"polar"）
"sel_type": [0],
"neuron": [100, 100, 100],
"resnet_dt": true,
"seed": 1,
"_comment": " that's all"
},
这里的sel_type指的是需要针对哪些原子类型拟合 tensor。第二处修改为 json 中的loss；
"loss" : {
"type": "tensor",
"pref": 1.0,
"pref_atomic": 1.0,
"_comment": "that's all"
},
在loss段中，type被统一填写为tensor。pref与pref_atomic分别表示针对一个 global label 的体系和一个 atomic label 的体系时，拟合他们的 label 所对应的权重。如果你希望平等的对待这些体系，可以将他们都设为1.0；如果你完全不希望使用体系中的 global label，可以将pref设为0.0。
训练方式及原理
数据集就位之后就可以直接训练：
dp train input.json
这里简单介绍一下 tensor 模式和 ener 模式下网络的区别。首先，ener 模式采用的 loss 为 energy、force、virial 的有监督拟合，即
loss = weight['ener'] * tf.reduce_mean(tf.square(hat_ener - ener))
+ weight['force'] * tf.reduce_mean(tf.square(hat_force - force))
+ weight['virial'] * tf.reduce_mean(tf.square(hat_virial - virial))
这里 hat 表示是网络输出的结果。与之相对，tensor 模式下损失函数只考虑对 tensor label 的拟合，即
loss = pref * tf.reduce_mean(tf.square(polar - hat_polar)) # global label的部分
+ pref_atomic * tf.reduce_mean(tf.square(atomic_polar - hat_atomic_polar)) / normalized_term # local label的部分
这里hat_polar和hat_atomic_label为用户提供的 label。此时网络的输出不再是标量，而是一个natoms x 3 (dipole)或natoms x (3 x 3) (polar)的张量，其中natoms表示有该 tensor 的原子个数（例如水分子中的氧的个数）。那么，DeePMD-kit 是如何做到这一点的呢？熟悉 ener 模式的小伙伴们都知道，DeePMD-kit 由一个 embedding network 和一个 fitting network 组成。在 ener 模式中，embedding network 负责将原始数据(coordinate)编码成一系列特征并输入到 fitting network，后者从这些特征中学习出 energy function 的表示。在 tensor 模式中，embedding network 是保持不变的，有所变化的是 fitting network：DeePMD-kit 使用了一个改进版本的全连接网络，使得在保持输入不变的情况下调整了输出的维度，从而改变的网络拟合的对象。
模型测试
完成训练之后，拟合模型就可以被用于预测各种体系的 tensor 了。这里简单介绍一下如何使用。首先，在训练的目录下（即有"model.ckpt.index"、"checkpoint"等文件的目录下运行）
dp freeze -o tensor.pb
生成 tensorflow 定义的二进制模型文件。接下来，可以直接运行
dp -m tensor.pb -s /path/to/system
来测试模型在某个 global label system 下的预测结果：
# number of test data : 100
Dipole L2err : 1.085408e-02 eV/A
# -----------------------------------------------
如果希望预测某个 atomic label system 的结果，可以运行
dp -m tensor.pb -s /path/to/system --atomic
除了使用 shell，也可以在 python 代码中导入模型用于预测。以 dipole 模式为例：
from deepmd import DeepDipole # use it to load your trained model
import numpy as np
import dpdata
nframe = 30 # feel free to alter this parameter
dd = dpdata.System('/path/to/system', fmt =
'deepmd/npy').sub_system(range(nframe))
dipole = DeepDipole('dipole.pb')
coord = dd['coords'].reshape([nframe,-1])
cell = dd['cells'].reshape([nframe,-1])
atype = dd['atom_types']
natom = len(atype)
result = dipole.eval(coord, cell, atype)
print(result.shape)
# (30, 64, 3)
这样我们就得到了模型输出的结果。
参考资料
[1]Dipole: https://journals.aps.org/prb/abstract/10.1103/PhysRevB.102.041121
[2]Polarizability: https://pubs.rsc.org/no/content/articlelanding/2020/cp/d0cp01893g/unauth#!divAbstract
你点的每一次“在看”，都能让更多人了解我们