备注：基于LAMMPS的data+strj文件+Python MDAnalysis

代码示例：

import MDAnalysis as mda
from MDAnalysis.analysis import msd
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

# 1. 加载轨迹文件
# 如果你有 data 文件，最好加上它作为拓扑，否则 MDA 只能从 dump 文件中猜测信息
# u = mda.Universe("traj.lammpstrj", format="LAMMPSDUMP")
u = mda.Universe("annealing.data", "run2.lammpstrj", format="LAMMPSDUMP")

# 2. 挑选目标分子
# 在 MDA 中，LAMMPS 的 Molecule ID 通常被映射为 'resid'
# 假设你想计算 Molecule ID 为 1 到 10 的分子
target_atoms = u.select_atoms("resid 1:330")

# 检查是否选对了原子
print(f"成功选择原子数量: {len(target_atoms)}")

# 3. 坐标“去卷绕” (Unwrapping) - 非常重要！
# MSD 计算必须基于连续坐标。如果你的 dump 文件是原子的原始坐标 (x, y, z)，
# 且原子跨越了周期性边界，直接计算会出错。
# 如果你 dump 的是 xu, yu, zu，这一步可以跳过。
# 如果是 x, y, z，且有 image flags，可以执行：
# u.atoms.unwrap(compound='fragments')

# 4. 使用内置工具计算 MSD
# fft=True 使用快速傅里叶变换，速度极快
msd_analysis = msd.EinsteinMSD(target_atoms, select='all', fft=True)
msd_analysis.run()

# 5. 提取结果
msd_values = msd_analysis.results.timeseries
frames = np.arange(len(msd_values))
dt = 10  # TODO: 这里的 dt 是你 dump 的时间间隔（lammps步长 * dump频率），步长单位和后面作图时候的一致，为ps
times = frames * dt

# 6. 绘图
plt.plot(times, msd_values)
plt.xlabel("Time (ps)")
plt.ylabel(r"MSD ($\AA^2$)")
plt.title("MSD of Target Molecules")
plt.show()

# 7. 计算扩散系数 D
# 根据爱因斯坦关系式：MSD = 6Dt (在 3D 情况下)
# 拟合线性部分（通常取轨迹的中间段）
start_frame = int(len(msd_values) * 0.1)
end_frame = int(len(msd_values) * 0.9)
slope, intercept = np.polyfit(times[start_frame:end_frame], msd_values[start_frame:end_frame], 1)
diffusion_coefficient = slope / 6.0
print(f"扩散系数 D = {diffusion_coefficient}")

# 8. 导出数据到 TXT 
output_file = "msd_results.txt"
data_to_save = np.column_stack((times, msd_values))
header = f"Time(ps)    MSD(A^2)\nDiffusion Coefficient D = {diffusion_coefficient:.6e} A^2/ps"

np.savetxt(output_file, data_to_save, fmt="%.6f", delimiter="\t", header=header)
print(f"结果已成功导出至: {output_file}")

# 9. 导出数据到 Excel
output_excel = "msd_results.xlsx"

# 创建一个数据框 (DataFrame)
df = pd.DataFrame({
    'Time (ps)': times,
    'MSD (A^2)': msd_values
})

# 导出到 Excel
# index=False 表示不保存行索引（即左侧的 0, 1, 2...）
df.to_excel(output_excel, index=False, sheet_name='MSD_Data')

print(f"结果已成功导出至 Excel: {output_excel}")