feat: add energy flux density panel to plot_wave animation

Add compute_energy_flux() using the Hardy formula:
  J_b = 1/2 * F_bond_on_i * (v_i + v_j)

New 4th subplot shows J vs bond position (x coordinate of bond midpoint).
J > 0: energy flows rightward; J = 0: standing wave; J < 0: leftward.
Ideal standing wave would show J ≈ 0 everywhere.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

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