这是【未来已来】系列学习笔记，为将来会发生的事情而学习，做好准备。目前神经网络势函数发展的技术路线非常多，尤其是基于神经网络势函数的大模型，会是未来AI for Science领域最重要的发展方向，笔者认为最有希望取得决定性突破的技术路线是消息传递等变图神经网络（equivariant message-passing neural networks）。这个学习笔记不保证内容的正确性，也不会过多介绍背景，仅仅是一个学习记录，大概会写十几篇笔记，从算法到代码实现，一点点来记录学习过程。
在消息传递等变图神经网络-学习笔记01中，讨论了消息传递的基本原理，后面就是怎么构建具体的message了。
在这个框架里，看到有个输入是rij，这就是每个神经网络势函数都需要的近邻原子相对位置的输入了。就是两个原子坐标的矢量差。怎么去构建这个近邻原子列表呢？大概步骤可以理解简单弄成两步：（1）ASE读入轨迹文件。（2）构建近邻原子列表
（1）ASE读入轨迹
from ase.io import read, write, Trajectoryfrom ase.calculators.singlepoint import SinglePointCalculatorimages = read('./OUTCAR', ':')write('dft_structures.traj', images)
这样就把一个OUTCAR文件转化成了ase atom格式的文件，并且输出成'dft_structures.traj'了。
然后
from ase.visualize import viewlen(images)view(images[0], viewer='x3d')
可以直接看某一帧的结构，比如[0]代表第一帧。len(images)可以看轨迹有几帧。
（2）构建近邻原子列表
用到asap3库，asap3 是一个用于分子动力学（Molecular Dynamics, MD）的加速库，它与 ASE（Atomic Simulation Environment）集成，用于高效地进行大规模分子动力学模拟。它的主要特点是通过优化算法和并行化显著提高计算效率。asap3 是一个 C++ 编写的 Python 包，因此需要编译和安装：
pip install asap3然后如果想要提取某一帧的结构的某一个原子近邻原子
import asap3# 读取 OUTCAR 文件中的所有结构images = read('./OUTCAR', index=':')atoms = images[-1] # 以数据集中最后一个结构为例nl = asap3.FullNeighborList(rCut=5.0, atoms=atoms) # cutoff设置为5.0 Å# 以下代码表示计算结构中第0个原子的neighbor list，其中indices, n_diff, n_dist_square分别为近邻原子的索引编号，近邻原子与中心原子矢量差，和两个原子距离的平方indices, n_diff, n_dist_square = nl.get_neighbors(0)print(indices, n_diff, n_dist_square)
输出为：
[ 24 91 95 119 121 123 130 142 178 190 118 120 126 177 192 193 26 2832 146 6 31 144 145 8 99 134 96 102 181 2 4 12 40 150 17 39 132 133 137] [[-3.82247000e+00 -1.83000000e-03 3.19000000e-03][-3.79000000e-03 -4.41843000e+00 -1.55882000e+00][ 1.91655000e+00 -3.31751000e+00 3.90000000e-04][-1.90946000e+00 -3.31592000e+00 5.12000000e-03]...[ 1.92005988e+00 1.10343219e+00 4.06873000e+00][ 3.50988100e-03 2.20623219e+00 8.24140000e-01][ 3.82805988e+00 2.20960219e+00 8.30160000e-01]] [14.61129043 21.95245782 14.67903665 14.64138915 15.14990033 7.324156845.66848635 20.25897198 20.30801197 20.29931947 22.23313254 15.500980637.60302293 20.21891577 21.26288625 5.56352712 15.11826928 21.9388155614.62038297 20.28148146 23.96776973 7.59028124 21.3849368 20.1520842814.68822916 7.34242627 20.34948657 15.49248464 22.24238362 5.5771948415.17431308 7.31589907 22.00279968 14.64607485 20.30897047 15.466315757.56348837 22.14459534 21.45875636 5.54667953 20.22554991]
实际上我们要把所有帧和每一帧的所有原子近邻列表都拿到，构建函数：
def get_neighborlist(cutoff, atoms):nl = asap3.FullNeighborList(cutoff, atoms)pair_i_idx = []pair_j_idx = []n_diff = []for i in range(len(atoms)):diff, _ = nl.get_neighbors(i)pair_i_idx += [i] * len(indices) # local index of pair i# local index of pair jn_diff.append(diff)= np.concatenate(pair_j_idx)pairs = np.stack((pair_i_idx, pair_j_idx), axis=1)n_diff = np.concatenate(n_diff)return pairs, n_diff= get_neighborlist(5.0, atoms) # atoms为某一帧结构n_diff)
1. def get_neighborlist(cutoff, atoms):
这个函数接受两个参数：
cutoff: 截断半径，决定邻居原子的最大距离，比如5 angstrom。atoms: 一个ASE原子对象，包含了原子的位置和类型。
2. nl = asap3.FullNeighborList(cutoff, atoms)
创建一个
FullNeighborList对象nl，使用指定的cutoff截断半径来计算atoms结构中的原子邻居关系。
3. pair_i_idx = [] 和 pair_j_idx = []
初始化两个空列表
pair_i_idx和pair_j_idx，分别用于存储原子对的索引。pair_i_idx用来记录每一对原子的第一个原子（i）。pair_j_idx用来记录每一对原子的第二个原子（j）。
4. n_diff = []
初始化一个空列表
n_diff用于存储邻居原子之间的位置差（位移矢量）。
5. for i in range(len(atoms)):
遍历每个原子（从索引
i为 0 到len(atoms)-1）。
6. indices, diff, _ = nl.get_neighbors(i)
使用
FullNeighborList对象nl的get_neighbors(i)方法获取原子i的所有邻居信息。indices是原子i的邻居原子的索引列表。diff是一个列表，包含邻居原子与原子i的位置差（位移矢量）。_表示不需要的返回值（通常是距离平方，这里未使用）。
7. pair_i_idx += [i] * len(indices)
对于原子
i的每一个邻居，将i添加到pair_i_idx列表中。[i] * len(indices)会创建一个长度为len(indices)的列表，所有元素都是i。
8. pair_j_idx.append(indices)
将原子
i的邻居索引indices添加到pair_j_idx列表中。每次都会把一个邻居的索引列表附加到pair_j_idx。
9. n_diff.append(diff)
将邻居原子与
i的位置差（位移矢量）diff添加到n_diff列表中。
10. pair_j_idx = np.concatenate(pair_j_idx)
将
pair_j_idx列表中的所有邻居原子索引按顺序连接成一个一维数组。原本每个元素是一个索引列表，concatenate将它们展开成一个平坦的数组。
11. pairs = np.stack((pair_i_idx, pair_j_idx), axis=1)
使用
np.stack将pair_i_idx和pair_j_idx按列（axis=1）合并成一个二维数组pairs，其中每行表示一对原子(i, j)。
12. n_diff = np.concatenate(n_diff)
将
n_diff列表中的所有位置差（位移矢量）连接成一个一维数组。
13. return pairs, n_diff
返回
pairs（原子对的索引）和n_diff（邻居原子之间的位移矢量）。
14. pairs, n_diff = get_neighborlist(5.0, atoms)
调用
get_neighborlist函数，传入截断半径 5.0 Å 和原子结构atoms。返回结果pairs包含所有原子对的索引，n_diff包含对应的位移矢量。
得到的输出为：
[[ 0 24][ 0 91][ 0 95]...[197 175][197 194][197 195]] [[-3.82247e+00 -1.83000e-03 3.19000e-03][-3.79000e-03 -4.41843e+00 -1.55882e+00][ 1.91655e+00 -3.31751e+00 3.90000e-04]...[ 2.25584e+00 8.06630e-01 -2.98133e+00][-1.75045e+00 2.71777e+00 4.45400e-02][-1.99210e+00 2.51283e+00 1.87348e+00]]
这样我们便获得了 n x 2 的索引矩阵pairs和 n x 3 的n_diff，其中n为所有pair的数目。此外我们还可以把 能量、原子数、元素组成、原子坐标、晶胞参数、原子受力得到，这些也是训练输入所必须的。
print(atoms.get_potential_energy())print(atoms.get_global_number_of_atoms())print(atoms.numbers)print(atoms.positions)print(atoms.cell[:])atoms.get_forces(apply_constraint=False)
输出为：
-1579.86779664198[8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 88 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 88 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 88 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 88 8 8 8 8 8 8 8 8 8 58 58 58 58 58 58 58 58 58 58 58 58 58 5858 58 58 58 58 58 58 58 58 58 58 58 58 58 58 58 58 58 58 58 58 58 58 5858 58 58 58 58 58 58 58 58 58 58 58 58 58 58 58 58 58 58 58 58 58 58 5858 58 79 6 6 28][][][]...[][][][]][][][]]array([[-7.600000e-04, -5.831300e-02, -1.352000e-02],[],[],[],[],...
为了创建一个 AseDataReader 类，旨在读取 ASE 数据并将其转化为 PyTorch 张量以便于训练，我们需要对类进行以下设计：
初始化：在初始化方法中，接受一个可选的
cutoff参数（默认为 5.0），并初始化邻居列表对象asap3.FullNeighborList。此对象需要在每次调用时使用不同的atoms对象。__call__方法：使得类的实例可以像函数一样被调用。在__call__方法中，接收atoms对象，计算邻居列表，并提取其他相关信息（如能量、原子类型、位置、力等）。数据转化：将
ASE中的数据（如能量、原子位置、力等）转换为PyTorch张量，并计算邻居对信息。
以下是 AseDataReader 类的完整实现：
import torchimport numpy as npimport asap3from ase.io import readclass AseDataReader:def __init__(self, cutoff=5.0):"""初始化数据读取器，设置邻居截断距离（cutoff）。参数：cutoff (float): 邻居的截断距离，默认值为 5.0 Å。"""self.cutoff = cutoffself.neighbor_list = Nonedef get_neighborlist(self, atoms):"""计算并返回每个原子的邻居列表及邻居对的位移矢量。参数：atoms (ase.Atoms): 一个 ASE Atoms 对象，包含当前的原子结构。返回：pairs (torch.Tensor): 一个二维 PyTorch 张量，表示所有邻居对的索引。n_diff (torch.Tensor): 一个二维 PyTorch 张量，表示所有邻居对之间的位移矢量。"""nl = asap3.FullNeighborList(self.cutoff, atoms)pair_i_idx = [] # 用于存储原子 i 的索引pair_j_idx = [] # 用于存储邻居 j 的索引n_diff = [] # 用于存储邻居对之间的位移矢量for i in range(len(atoms)):indices, diff, _ = nl.get_neighbors(i)pair_i_idx += [i] * len(indices) # 当前原子的索引 ipair_j_idx.append(indices) # 当前原子的邻居的索引 jn_diff.append(diff) # 邻居对之间的位移矢量pair_j_idx = np.concatenate(pair_j_idx)pairs = np.vstack((pair_i_idx, pair_j_idx)).T # 将原子 i 和邻居 j 的索引堆叠成对n_diff = np.concatenate(n_diff) # 邻居对之间的位移矢量# 转换为 PyTorch 张量并返回pairs = torch.tensor(pairs, dtype=torch.long)n_diff = torch.tensor(n_diff, dtype=torch.float32)return pairs, n_diffdef __call__(self, atoms):"""接收 ASE Atoms 对象并提取相关信息，并计算邻居对信息。参数：atoms (ase.Atoms): 一个 ASE Atoms 对象，表示某一帧的原子结构。返回：data (dict): 包含能量、原子数、类型、位置、力和邻居对数据的字典。"""# 更新邻居列表对象self.neighbor_list = asap3.FullNeighborList(self.cutoff, atoms)# 计算邻居列表和位移矢量pairs, n_diff = self.get_neighborlist(atoms)# 获取总能量energy = torch.tensor(atoms.get_potential_energy(), dtype=torch.float32)# 获取原子数num_atoms = atoms.get_global_number_of_atoms()# 获取原子类型（elems）elems = torch.tensor(atoms.numbers, dtype=torch.long)# 获取原子位置positions = torch.tensor(atoms.positions, dtype=torch.float32)# 获取原子力forces = torch.tensor(atoms.get_forces(apply_constraint=False), dtype=torch.float32)image_idx = torch.zeros((atoms_data['num_atoms'],), dtype=torch.long)# 返回数据字典return {'energy': energy,'num_atoms': num_atoms,'elems': elems,'positions': positions,'forces': forces,'pairs': pairs,'n_diff': n_diff,'image_idx' : image_idx}# 示例：使用 AseDataReader 来处理 ASE 数据data_reader = AseDataReader(cutoff=5.0)# 读取 ASE 文件（如 OUTCAR）images = read('./OUTCAR', index=':')print(data_reader(images[-1]))'''# 处理每一帧的原子数据for atoms in images:data = data_reader(atoms)print(data) # 打印处理后的数据字典'''
