一键处理AIMD获得MSD并绘图脚本
多粒子运动的均方位移mean squre displacement (MSD)的定义为:
VASP的MD计算轨迹文件输出到XDATCAR文件. 通过python代码将XDATCAR 输出到xyz文件,然后根据选定的具体原子来计算其轨迹相对于初始位置一定步长下的演化获得其MSD
方法参考:
https://zhuanlan.zhihu.com/p/542642528
https://zhuanlan.zhihu.com/p/499032498
执行命令以及过程
python x-xyz-msd-2.py XDATCAR --step 50 Cu
选定元素所有原子MSD
部分原子不同方向MSD
代码全文
#!/usr/bin/env python3
"""
将VASP的XDATCAR文件转换为展开的xyz格式并计算均方位移(MSD)
增强功能:支持多原子选择、元素选择、Δt步长和绘图
作者: lipai@mail.ustc.edu.cn
"""import numpy as np
import os
import sys
import argparse
import matplotlib.pyplot as plt
from copy import deepcopydef parse_arguments():"""解析命令行参数"""parser = argparse.ArgumentParser(description='将XDATCAR转换为xyz格式并计算MSD')parser.add_argument('xdatcar_file', help='VASP的XDATCAR文件')parser.add_argument('atoms', nargs='+', help='原子索引(从0开始)或元素符号,-1表示所有原子')parser.add_argument('--step', type=int, default=1,help='Δt步长 (默认: 1)')parser.add_argument('--no-plot', action='store_true',help='不生成MSD图像')parser.add_argument('--output-dir', default='.',help='输出目录 (默认: 当前目录)')return parser.parse_args()def xdatcar_to_xyz(xdatcar_file='XDATCAR', xyz_file='XDATCAR.xyz'):"""将XDATCAR文件转换为展开的xyz格式Args:xdatcar_file: 输入的XDATCAR文件名xyz_file: 输出的xyz文件名"""try:xdatcar = open(xdatcar_file, 'r')xyz = open(xyz_file, 'w')except FileNotFoundError:print(f"错误: 找不到文件 {xdatcar_file}")sys.exit(1)# 读取系统信息system = xdatcar.readline().strip()scale = float(xdatcar.readline().rstrip('\n'))print(f"缩放因子: {scale}")# 读取晶格向量a1 = np.array([float(s)*scale for s in xdatcar.readline().rstrip('\n').split()])a2 = np.array([float(s)*scale for s in xdatcar.readline().rstrip('\n').split()])a3 = np.array([float(s)*scale for s in xdatcar.readline().rstrip('\n').split()])# 构建注释行comment = f'Lattice="{a1[0]} {a1[1]} {a1[2]} {a2[0]} {a2[1]} {a2[2]} {a3[0]} {a3[1]} {a3[2]}" 'comment += 'Properties=species:S:1:pos:R:3 pbc="T T T"'# 读取元素信息element_names = xdatcar.readline().rstrip('\n').split()element_numbers = xdatcar.readline().rstrip('\n').split()# 创建原子名称列表Natom = 0Ntype = len(element_names)Nname = []element_start_indices = {} # 记录每种元素的起始索引for t in range(Ntype):element_start_indices[element_names[t]] = NatomNatom += int(element_numbers[t])for i in range(int(element_numbers[t])):Nname.append(element_names[t])print(f"原子类型数: {Ntype}, 总原子数: {Natom}")print(f"元素组成: {dict(zip(element_names, element_numbers))}")# 处理轨迹frame_count = 0f_prev = Noneall_coords = [] # 存储所有帧的坐标while True:line = xdatcar.readline()if len(line) == 0:breakframe_coords = np.zeros((Natom, 3)) # 当前帧的直角坐标f_next = np.zeros((Natom, 3)) # 当前帧的分数坐标xyz.write(f"{Natom}\n{comment}\n")for atom in range(Natom):p = xdatcar.readline().rstrip('\n').split()f_next[atom, :] = np.array([float(s) for s in p])# 轨迹展开if f_prev is not None:f_next[atom, :] -= np.around(f_next[atom, :] - f_prev[atom, :])# 转换为直角坐标c_coords = f_next[atom, 0]*a1 + f_next[atom, 1]*a2 + f_next[atom, 2]*a3frame_coords[atom, :] = c_coordsxyz.write(f"{Nname[atom]} {c_coords[0]} {c_coords[1]} {c_coords[2]}\n")all_coords.append(frame_coords)f_prev = deepcopy(f_next)frame_count += 1xdatcar.close()xyz.close()print(f"处理完成: 共 {frame_count} 帧轨迹")return np.array(all_coords), Nname, element_start_indices, element_names, element_numbersdef get_atom_indices_from_input(atom_inputs, Nname, element_start_indices, element_names, element_numbers):"""根据用户输入获取原子索引列表Args:atom_inputs: 用户输入的原子标识列表Nname: 原子名称列表element_start_indices: 元素起始索引字典element_names: 元素名称列表element_numbers: 元素数量列表"""atom_indices = []for atom_input in atom_inputs:# 尝试解析为整数try:atom_idx = int(atom_input)if atom_idx == -1: # 所有原子atom_indices = list(range(len(Nname)))breakelif 0 <= atom_idx < len(Nname):atom_indices.append(atom_idx)else:print(f"警告: 原子索引 {atom_idx} 超出范围 (0-{len(Nname)-1})")except ValueError:# 不是整数,尝试作为元素符号处理if atom_input in element_names:start_idx = element_start_indices[atom_input]count = int(element_numbers[element_names.index(atom_input)])atom_indices.extend(range(start_idx, start_idx + count))print(f"选择元素 {atom_input} 的所有原子: 索引 {start_idx} 到 {start_idx+count-1}")else:print(f"警告: 未知元素符号 '{atom_input}'")# 去重并排序atom_indices = sorted(list(set(atom_indices)))return atom_indicesdef calculate_msd(xyz_coords, atom_indices, step=1, output_dir=".", no_plot=False):"""计算均方位移(MSD)Args:xyz_coords: 轨迹坐标数组 [帧数, 原子数, 3]atom_indices: 要计算MSD的原子索引列表step: Δt步长output_dir: 输出目录no_plot: 是否不生成图像"""FrameNum, AtomNum, _ = xyz_coords.shape# 为所有原子创建输出目录if len(atom_indices) > 1:msd_dir = os.path.join(output_dir, "msd_all")os.makedirs(msd_dir, exist_ok=True)else:msd_dir = output_dir# 存储所有原子的MSD数据用于绘图all_msd_data = {}for atom_idx in atom_indices:if atom_idx >= AtomNum:print(f"错误: 原子索引 {atom_idx} 大于总原子数 {AtomNum}")continueprint(f"计算原子 {atom_idx} 的MSD...", end="")# 初始化MSD数组max_del_t = (FrameNum - 1) // stepmsd = np.zeros(max_del_t)msd_i = np.zeros((max_del_t, 3)) # 三个方向的分量# 计算MSDfor i, del_t in enumerate(range(1, FrameNum, step)):if i >= max_del_t:breakFd = FrameNum - del_t # 可用的时间起点数for n in range(Fd):displacement = xyz_coords[n + del_t, atom_idx, :] - xyz_coords[n, atom_idx, :]msd_i[i, :] += displacement**2 / Fdmsd[i] = np.sum(msd_i[i, :])print("完成")# 存储MSD数据all_msd_data[atom_idx] = {'msd': msd,'msd_i': msd_i,'time_steps': np.arange(1, FrameNum, step)[:len(msd)]}# 输出结果if len(atom_indices) > 1:filename = os.path.join(msd_dir, f"{atom_idx}")else:filename = f"{atom_idx}"# 保存总MSDwith open(f"{filename}.msd", "w") as f:for i in range(len(msd)):f.write(f"{all_msd_data[atom_idx]['time_steps'][i]} {msd[i]}\n")# 保存三个方向的MSD分量with open(f"{filename}.msd_i", "w") as f:for i in range(len(msd)):f.write(f"{all_msd_data[atom_idx]['time_steps'][i]} {msd_i[i, 0]} {msd_i[i, 1]} {msd_i[i, 2]}\n")# 绘制MSD图像if not no_plot:plot_msd(all_msd_data, atom_indices, output_dir)def plot_msd(all_msd_data, atom_indices, output_dir):"""绘制MSD图像Args:all_msd_data: 所有原子的MSD数据atom_indices: 原子索引列表output_dir: 输出目录"""# 设置绘图样式plt.style.use('default')colors = plt.cm.tab10(np.linspace(0, 1, len(atom_indices)))# 创建总MSD图fig1, ax1 = plt.subplots(figsize=(10, 6))for i, atom_idx in enumerate(atom_indices):if atom_idx not in all_msd_data:continuedata = all_msd_data[atom_idx]ax1.plot(data['time_steps'], data['msd'], label=f'Atom {atom_idx}', color=colors[i], linewidth=2)ax1.set_xlabel('Time Step (Δt)', fontsize=12)ax1.set_ylabel('MSD (Ų)', fontsize=12)ax1.set_title('Mean Square Displacement', fontsize=14)ax1.legend(fontsize=10)ax1.grid(True, alpha=0.3)# 保存总MSD图plt.tight_layout()plt.savefig(os.path.join(output_dir, 'msd_total.png'), dpi=300, bbox_inches='tight')plt.close(fig1)# 为每个原子创建详细MSD图(如果原子数量不多)if len(atom_indices) <= 10:for atom_idx in atom_indices:if atom_idx not in all_msd_data:continuedata = all_msd_data[atom_idx]fig2, ax2 = plt.subplots(figsize=(10, 6))ax2.plot(data['time_steps'], data['msd'], label='Total MSD', color='black', linewidth=3)ax2.plot(data['time_steps'], data['msd_i'][:, 0], label='X direction', linestyle='--', alpha=0.8)ax2.plot(data['time_steps'], data['msd_i'][:, 1], label='Y direction', linestyle='--', alpha=0.8)ax2.plot(data['time_steps'], data['msd_i'][:, 2], label='Z direction', linestyle='--', alpha=0.8)ax2.set_xlabel('Time Step (Δt)', fontsize=12)ax2.set_ylabel('MSD (Ų)', fontsize=12)ax2.set_title(f'MSD for Atom {atom_idx}', fontsize=14)ax2.legend(fontsize=10)ax2.grid(True, alpha=0.3)plt.tight_layout()if len(atom_indices) > 1:filename = os.path.join(output_dir, "msd_all", f"msd_atom_{atom_idx}.png")else:filename = os.path.join(output_dir, f"msd_atom_{atom_idx}.png")plt.savefig(filename, dpi=300, bbox_inches='tight')plt.close(fig2)print(f"MSD图像已保存到 {output_dir}")def main():"""主函数"""args = parse_arguments()# 第一步: 转换XDATCAR为xyz格式print("第一步: 转换XDATCAR为xyz格式")xyz_coords, atom_names, element_start_indices, element_names, element_numbers = xdatcar_to_xyz(args.xdatcar_file)# 第二步: 解析原子选择print("\n第二步: 解析原子选择")atom_indices = get_atom_indices_from_input(args.atoms, atom_names, element_start_indices, element_names, element_numbers)if not atom_indices:print("错误: 没有找到有效的原子选择")sys.exit(1)print(f"选择的原子索引: {atom_indices}")# 第三步: 计算MSDprint(f"\n第三步: 计算均方位移(MSD), Δt步长 = {args.step}")calculate_msd(xyz_coords, atom_indices, args.step, args.output_dir, args.no_plot)print("\n所有计算完成!")if __name__ == "__main__":main()