266 lines
10 KiB
Python
266 lines
10 KiB
Python
"""
|
||
plot_trajectory.py
|
||
------------------
|
||
绘制轨迹数据的x-t, y-t, z-t, vx-t, vy-t, vz-t函数图像，支持多原子同时绘制。
|
||
|
||
用法：
|
||
python plot_trajectory.py # 绘制所有原子
|
||
python plot_trajectory.py output/trajectory.txt # 指定文件
|
||
python plot_trajectory.py output/trajectory.txt --atom-id 1 # 仅绘制原子1
|
||
|
||
参数：
|
||
trajectory_file: 轨迹数据文件，支持结构化 .txt 格式（默认：output/trajectory.txt）
|
||
--atom-id: 可选，指定只绘制单个原子的轨迹
|
||
"""
|
||
|
||
import numpy as np
|
||
import matplotlib.pyplot as plt
|
||
import sys
|
||
import os
|
||
|
||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||
import compute
|
||
|
||
def select_atom_series(data, atom_id=None):
|
||
"""Select one atom from trajectory arrays when the file stores many atoms."""
|
||
traj_x = data['traj_x']
|
||
traj_y = data['traj_y']
|
||
traj_z = data['traj_z']
|
||
traj_vx = data['traj_vx']
|
||
traj_vy = data['traj_vy']
|
||
traj_vz = data['traj_vz']
|
||
|
||
if traj_x.ndim == 1:
|
||
selected_id = int(data["plot_atom_id"]) if "plot_atom_id" in data else 1
|
||
return traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz, selected_id
|
||
|
||
atom_ids = data["atom_ids"] if "atom_ids" in data else np.arange(traj_x.shape[1]) + 1
|
||
if atom_id is None:
|
||
row = int(data["plot_atom_row"]) if "plot_atom_row" in data else 0
|
||
else:
|
||
matches = np.where(atom_ids == int(atom_id))[0]
|
||
if len(matches) == 0:
|
||
raise ValueError(f"原子序号 {atom_id} 不在轨迹文件中")
|
||
row = int(matches[0])
|
||
selected_id = int(atom_ids[row])
|
||
return (
|
||
traj_x[:, row], traj_y[:, row], traj_z[:, row],
|
||
traj_vx[:, row], traj_vy[:, row], traj_vz[:, row],
|
||
selected_id,
|
||
)
|
||
|
||
|
||
def load_trajectory_txt(file_path, atom_id=None):
|
||
"""加载结构化 .txt 格式的轨迹数据"""
|
||
data = compute.load_text_data(file_path)
|
||
traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz, selected_id = select_atom_series(data, atom_id)
|
||
NT = int(data['NT'])
|
||
DT = float(data['DT'])
|
||
return traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz, NT, DT, selected_id
|
||
|
||
|
||
def load_full_trajectory(file_path):
|
||
"""加载完整轨迹数据（所有原子），返回 2D 数组及原始数据字典。"""
|
||
data = compute.load_text_data(file_path)
|
||
traj_x = data['traj_x']
|
||
traj_y = data['traj_y']
|
||
traj_z = data['traj_z']
|
||
traj_vx = data['traj_vx']
|
||
traj_vy = data['traj_vy']
|
||
traj_vz = data['traj_vz']
|
||
NT = int(data['NT'])
|
||
DT = float(data['DT'])
|
||
n_atoms = traj_x.shape[1] if traj_x.ndim > 1 else 1
|
||
atom_ids = data.get("atom_ids", np.arange(n_atoms) + 1)
|
||
return traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz, NT, DT, atom_ids, data
|
||
|
||
def create_plots(traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz, NT, DT, output_dir='.', atom_ids=None, extra_data=None):
|
||
"""创建六个子图的图表，支持多原子绘制与能量曲线。"""
|
||
# 配置中文字体支持
|
||
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'WenQuanYi Micro Hei', 'DejaVu Sans']
|
||
plt.rcParams['axes.unicode_minus'] = False
|
||
|
||
# 生成时间数组
|
||
time = np.arange(NT) * DT
|
||
|
||
# 判断单原子（1D）或多原子（2D）
|
||
is_multi = traj_x.ndim == 2
|
||
n_atoms = traj_x.shape[1] if is_multi else 1
|
||
|
||
# 是否绘制能量子图（需要 multi-atom + extra_data 中有 G）
|
||
has_energy = is_multi and extra_data is not None and "G" in extra_data
|
||
n_rows = 3 if has_energy else 2
|
||
n_cols = 3
|
||
figsize = (15, 13) if has_energy else (15, 9)
|
||
|
||
# 创建图形和子图
|
||
fig, axes = plt.subplots(n_rows, n_cols, figsize=figsize)
|
||
title = '轨迹与能量分析' if has_energy else '轨迹分析'
|
||
if atom_ids is not None and not is_multi:
|
||
title += f' - 原子 {int(atom_ids)}'
|
||
fig.suptitle(title, fontsize=16)
|
||
|
||
plot_configs = [
|
||
(axes[0, 0], traj_x, 'x'),
|
||
(axes[0, 1], traj_y, 'y'),
|
||
(axes[0, 2], traj_z, 'z'),
|
||
(axes[1, 0], traj_vx, 'vx'),
|
||
(axes[1, 1], traj_vy, 'vy'),
|
||
(axes[1, 2], traj_vz, 'vz'),
|
||
]
|
||
|
||
if is_multi:
|
||
colors = plt.cm.tab10(np.linspace(0, 1, n_atoms))
|
||
if atom_ids is None:
|
||
atom_ids = np.arange(n_atoms) + 1
|
||
|
||
for ax, data_arr, label in plot_configs:
|
||
if is_multi:
|
||
for i in range(n_atoms):
|
||
aid = int(atom_ids[i])
|
||
ax.plot(time, data_arr[:, i], color=colors[i], linewidth=1.5, label=f"原子 {aid}")
|
||
ax.legend()
|
||
else:
|
||
ax.plot(time, data_arr, 'b-', linewidth=1.5)
|
||
ax.set_title(f'{label} - 时间')
|
||
ax.set_xlabel('时间 (s)')
|
||
ax.set_ylabel(label)
|
||
ax.grid(True, alpha=0.3)
|
||
|
||
# ── 能量子图 ─────────────────────────────────────
|
||
if has_energy:
|
||
masses = np.array(extra_data["atom_masses"]) # (n_atoms,)
|
||
G_vec = np.array(extra_data["G"]) # [gx, gy, gz]
|
||
|
||
# 动能 Ek = ½ m v²
|
||
ek = 0.5 * masses[np.newaxis, :] * (traj_vx**2 + traj_vy**2 + traj_vz**2)
|
||
# 重力势能 Ug = -m G·r
|
||
ug = -masses[np.newaxis, :] * (G_vec[0] * traj_x + G_vec[1] * traj_y + G_vec[2] * traj_z)
|
||
# 弹性势能 Us = ½ k (d - d₀)²
|
||
us = np.zeros_like(ek)
|
||
bond_pairs = extra_data.get("bond_pairs")
|
||
bond_stiffness = extra_data.get("bond_stiffness")
|
||
bond_rest_lengths = extra_data.get("bond_rest_lengths")
|
||
if bond_pairs is not None and len(bond_pairs) > 0:
|
||
for b_idx in range(len(bond_pairs)):
|
||
i, j = bond_pairs[b_idx]
|
||
dx = traj_x[:, j] - traj_x[:, i]
|
||
dy = traj_y[:, j] - traj_y[:, i]
|
||
dz = traj_z[:, j] - traj_z[:, i]
|
||
dist = np.sqrt(dx**2 + dy**2 + dz**2)
|
||
stretch = dist - bond_rest_lengths[b_idx]
|
||
us[:, i] += 0.5 * bond_stiffness[b_idx] * stretch**2
|
||
us[:, j] += 0.5 * bond_stiffness[b_idx] * stretch**2
|
||
|
||
e_total = ek + ug + us # (NT, n_atoms)
|
||
ek_sys = np.sum(ek, axis=1)
|
||
ug_sys = np.sum(ug, axis=1)
|
||
us_sys = np.sum(us, axis=1)
|
||
e_sys = ek_sys + ug_sys + us_sys
|
||
|
||
# 第 3 行左：各原子总能量
|
||
ax_e = axes[2, 0]
|
||
for i in range(n_atoms):
|
||
aid = int(atom_ids[i])
|
||
ax_e.plot(time, e_total[:, i], color=colors[i], linewidth=1.5, label=f"原子 {aid}")
|
||
ax_e.set_title("各原子总能量")
|
||
ax_e.set_xlabel("时间 (s)")
|
||
ax_e.set_ylabel("能量")
|
||
ax_e.grid(True, alpha=0.3)
|
||
ax_e.legend(loc="upper right")
|
||
|
||
# 第 3 行右：系统总能量
|
||
ax_sys = axes[2, 1]
|
||
ax_sys.plot(time, ek_sys, 'b-', linewidth=1.5, label="系统动能")
|
||
ax_sys.plot(time, ug_sys, 'g-', linewidth=1.5, label="系统重力势能")
|
||
if bond_pairs is not None and len(bond_pairs) > 0:
|
||
ax_sys.plot(time, us_sys, color='orange', linewidth=1.5, label="系统弹性势能")
|
||
ax_sys.plot(time, e_sys, 'r--', linewidth=1.5, label="系统总能量")
|
||
ax_sys.set_title("系统总能量")
|
||
ax_sys.set_xlabel("时间 (s)")
|
||
ax_sys.set_ylabel("能量")
|
||
ax_sys.grid(True, alpha=0.3)
|
||
ax_sys.legend(loc="upper right")
|
||
|
||
# 隐藏第 3 行第 3 列空子图
|
||
axes[2, 2].set_visible(False)
|
||
|
||
# 调整布局
|
||
plt.tight_layout(rect=[0, 0.03, 1, 0.95])
|
||
|
||
# 保存图像
|
||
output_path = os.path.join(output_dir, 'trajectory_plots.png')
|
||
plt.savefig(output_path, dpi=300, bbox_inches='tight')
|
||
print(f"图表已保存至: {output_path}")
|
||
|
||
# 显示图像（如果环境支持）
|
||
plt.show()
|
||
|
||
return output_path
|
||
|
||
def main():
|
||
import argparse
|
||
|
||
parser = argparse.ArgumentParser(description="绘制轨迹数据图表")
|
||
parser.add_argument("trajectory_file", nargs="?",
|
||
help="轨迹数据文件路径（默认: output/trajectory.txt）")
|
||
parser.add_argument("--atom-id", type=int, default=None,
|
||
help="指定只绘制某个原子的轨迹（默认绘制所有原子）")
|
||
args = parser.parse_args()
|
||
|
||
# 默认文件
|
||
script_dir = os.path.dirname(os.path.abspath(__file__))
|
||
default_file = os.path.join(compute.get_output_dir(script_dir), "trajectory.txt")
|
||
input_file = args.trajectory_file or default_file
|
||
|
||
# 检查文件是否存在
|
||
if not os.path.exists(input_file):
|
||
print(f"错误: 文件 '{input_file}' 不存在")
|
||
print(f"请先运行 compute.py 生成轨迹数据")
|
||
sys.exit(1)
|
||
|
||
file_ext = os.path.splitext(input_file)[1].lower()
|
||
|
||
try:
|
||
if args.atom_id is not None:
|
||
# 单原子模式（旧版兼容）
|
||
traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz, NT, DT, selected_id = \
|
||
load_trajectory_txt(input_file, args.atom_id)
|
||
ids_for_plot = selected_id
|
||
extra_data = None
|
||
n_atoms_str = f" 绘图原子序号: {selected_id}"
|
||
else:
|
||
# 全原子模式
|
||
traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz, NT, DT, atom_ids, raw_data = \
|
||
load_full_trajectory(input_file)
|
||
ids_for_plot = atom_ids
|
||
extra_data = raw_data
|
||
n_atoms = traj_x.shape[1] if traj_x.ndim > 1 else 1
|
||
n_atoms_str = f" 绘图原子数: {n_atoms}"
|
||
except Exception as e:
|
||
print(f"加载文件时出错: {e}")
|
||
sys.exit(1)
|
||
|
||
print(f"加载轨迹数据: NT={NT}, DT={DT}")
|
||
if args.atom_id is not None:
|
||
print(n_atoms_str)
|
||
else:
|
||
print(n_atoms_str)
|
||
print(f" 位置范围: x [{traj_x.min():.3f}, {traj_x.max():.3f}], "
|
||
f"y [{traj_y.min():.3f}, {traj_y.max():.3f}], "
|
||
f"z [{traj_z.min():.3f}, {traj_z.max():.3f}]")
|
||
print(f" 速度范围: vx [{traj_vx.min():.3f}, {traj_vx.max():.3f}], "
|
||
f"vy [{traj_vy.min():.3f}, {traj_vy.max():.3f}], "
|
||
f"vz [{traj_vz.min():.3f}, {traj_vz.max():.3f}]")
|
||
|
||
# 创建图表
|
||
output_path = create_plots(
|
||
traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz,
|
||
NT, DT, compute.get_output_dir(script_dir), ids_for_plot, extra_data,
|
||
)
|
||
|
||
print("绘图完成！")
|
||
|
||
if __name__ == "__main__":
|
||
main()
|