diff --git a/plot_wave.py b/plot_wave.py index 9d3bf96..09dd2c9 100644 --- a/plot_wave.py +++ b/plot_wave.py @@ -261,13 +261,73 @@ def compute_per_atom_energy(x, y, z, vx, vy, vz, masses, return ek_atom, pe_atom, et_atom +def compute_energy_flux(x, y, z, vx, vy, vz, + bond_pairs, bond_stiffness, bond_rest_lengths): + """计算每帧每根键的能流密度(Hardy 公式)。 + + 对键 b 连接原子 i, j(j > i,方向 i→j): + + J_b = ½ · F_{b,i} · (v_i + v_j) + + 其中 F_{b,i} = k(d - r₀) * (r_j - r_i)/d 是键对原子 i 的弹力矢量, + 点乘取两端速度均值。 + + - J > 0:能量从 i 流向 j(沿键方向正流) + - J = 0:驻波,能量不流动 + - 沿链从左到右,J 的分布揭示能量传播方向 + + Returns: + flux: (n_frames, n_bonds) 每帧每键的能流(标量) + bond_xpos: (n_bonds,) 各键中点的初始 x 坐标(用于绘图横轴) + """ + if bond_pairs is None or len(bond_pairs) == 0: + n_frames = x.shape[0] + return np.zeros((n_frames, 0)), np.zeros(0) + + n_bonds = len(bond_pairs) + n_frames = x.shape[0] + flux = np.zeros((n_frames, n_bonds)) + + # 键中点初始 x 坐标(用于横轴定位) + bond_xpos = np.array([ + 0.5 * (x[0, bond_pairs[b, 0]] + x[0, bond_pairs[b, 1]]) + for b in range(n_bonds) + ]) + + for b in range(n_bonds): + i, j = int(bond_pairs[b, 0]), int(bond_pairs[b, 1]) + k = bond_stiffness[b] + r0 = bond_rest_lengths[b] + + # 键矢量与长度 + dx_ = x[:, j] - x[:, i] + dy_ = y[:, j] - y[:, i] + dz_ = z[:, j] - z[:, i] + d = np.sqrt(dx_**2 + dy_**2 + dz_**2) + d = np.maximum(d, 1e-12) + + # 弹力矢量(作用于原子 i,指向 j 方向) + fac = k * (d - r0) / d # 标量因子 + fx = fac * dx_ + fy = fac * dy_ + fz = fac * dz_ + + # 能流 = F_i · (v_i + v_j) / 2 + flux[:, b] = 0.5 * (fx * (vx[:, i] + vx[:, j]) + + fy * (vy[:, i] + vy[:, j]) + + fz * (vz[:, i] + vz[:, j])) + + return flux, bond_xpos + + def plot_wave(output_dir, save_gif=False, save_mp4=False, show=True): """主绘图函数:读取 display.txt 并生成波形+能量动画。 - 布局(3行×1列,纵向排列): + 布局(4行×1列,纵向排列): 行0:x/y/z 位移波形叠加在同一子图(vs 原子序号) 行1:每粒子动能、势能、总能叠加在同一子图 - 行2:系统总能量随时间变化 + 行2:键能流密度 J(Hardy 公式,vs 键中点位置) + 行3:系统总能量随时间变化 Args: output_dir: 输出目录(含 display.npz 或 display.txt) @@ -324,6 +384,11 @@ def plot_wave(output_dir, save_gif=False, save_mp4=False, show=True): bond_pairs, bond_stiffness, bond_rest_lengths, atom_ids, driver_info) + # ── 能流密度 ── + flux, bond_xpos = compute_energy_flux( + x, y, z, vx, vy, vz, + bond_pairs, bond_stiffness, bond_rest_lengths) + # ── y 轴范围 ── def get_ylim(arr): vmax = np.max(np.abs(arr)) @@ -351,15 +416,23 @@ def plot_wave(output_dir, save_gif=False, save_mp4=False, show=True): e_max = max(np.max(e_total), 0.01) * 1.3 p_max = max(np.max(np.abs(power)) * 1.3, 0.01) + # 能流 y 轴范围(对称,正负各半) + if flux.size > 0: + flux_vmax = np.max(np.abs(flux)) + flux_vmax = flux_vmax if flux_vmax > 1e-12 else 1.0 + flux_ylim = (-flux_vmax * 1.2, flux_vmax * 1.2) + else: + flux_ylim = (-1.0, 1.0) + atom_idx = np.arange(n_atoms) - # ── 图形布局:3 行 × 1 列,纵向排列 ── + # ── 图形布局:4 行 × 1 列,纵向排列 ── plt.rcParams['font.sans-serif'] = ['Microsoft YaHei', 'SimHei', 'DejaVu Sans'] plt.rcParams['axes.unicode_minus'] = False - fig, (ax_wave, ax_energy, ax_ep) = plt.subplots(3, 1, figsize=(12, 14)) + fig, (ax_wave, ax_energy, ax_flux, ax_ep) = plt.subplots(4, 1, figsize=(12, 18)) fig.suptitle("波形与能量分析", fontsize=16) - fig.subplots_adjust(hspace=0.40, top=0.94) + fig.subplots_adjust(hspace=0.42, top=0.95) # ── 图1:x/y/z 位移波形叠加 ── ax_wave.set_xlim(0, n_atoms - 1) @@ -397,7 +470,19 @@ def plot_wave(output_dir, save_gif=False, save_mp4=False, show=True): energy_lines.append(ln) ax_energy.legend(loc="upper right", fontsize=9) - # ── 图3:系统总能量随时间 ── + # ── 图3:能流密度 J(Hardy 公式)── + xmin_flux = bond_xpos[0] if len(bond_xpos) > 0 else 0 + xmax_flux = bond_xpos[-1] if len(bond_xpos) > 0 else n_atoms - 1 + ax_flux.set_xlim(xmin_flux, xmax_flux) + ax_flux.set_ylim(flux_ylim) + ax_flux.axhline(0, color="gray", linewidth=0.8, linestyle="--") + ax_flux.set_xlabel("位置(键中点 x 坐标)") + ax_flux.set_ylabel("能流密度 J") + ax_flux.set_title("键能流密度 J = ½ F·(vᵢ+vⱼ) (J>0 向右传播,J<0 向左传播)") + ax_flux.grid(True, alpha=0.3) + flux_line, = ax_flux.plot([], [], color="#dc2626", linewidth=1.5) + + # ── 图4:系统总能量随时间 ── ax_ep.set_xlim(t[0], t[-1]) ep_yhigh = max(e_max, p_max) ep_ylow = min(-p_max * 0.1, 0.0) @@ -430,7 +515,11 @@ def plot_wave(output_dir, save_gif=False, save_mp4=False, show=True): for i, ln in enumerate(energy_lines): ln.set_data(atom_idx, energy_arrays[i][frame]) - # 图3:系统能量(累计到当前帧) + # 图3:能流密度 + if flux.shape[1] > 0: + flux_line.set_data(bond_xpos, flux[frame]) + + # 图4:系统能量(累计到当前帧) cur_t = t[:frame + 1] ln_ek.set_data(cur_t, ek_sys[:frame + 1]) ln_us.set_data(cur_t, us_sys[:frame + 1]) @@ -441,7 +530,7 @@ def plot_wave(output_dir, save_gif=False, save_mp4=False, show=True): ax_ep.set_xlim(t[0], max(t[frame] + max(t[-1] * 0.05, 1), t[-1])) artists = wave_lines + [time_text] + energy_lines + \ - [ln_ek, ln_us, ln_et, ln_pw] + [flux_line, ln_ek, ln_us, ln_et, ln_pw] if ln_ug: artists.append(ln_ug) if ln_ugr: artists.append(ln_ugr) return artists