ylniu/dynamics
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

docs: 更新 examples/Readme.md 并新增 Readme.html

Browse Source 
- 覆盖全部 10 个案例（原 Readme 只到 case06）
- 新增案例选择指南表格
- Readme.html 为深色主题独立 HTML 页面
  （含卡片布局、标签分类、代码高亮、响应式设计）
- 各案例详情对齐最新配置参数
This commit is contained in:
Ying-Li Niu
2026-06-17 15:33:49 +08:00
parent ea99f09f9b
commit 0e636e275d
59 changed files with 7302 additions and 418 deletions
Download Patch File Download Diff File Expand all files Collapse all files
plot_wave.py
+420 -63
View File
@@ -108,9 +108,21 @@ def _load_wave_dataset(output_dir):
"gravity_strength": float(header.get("gravity_strength", 1.0)),
"G": gravity_vec,
"driving_force": int(header.get("driving_force", 0)),
"display_amp": _parse_display_amp(header.get("display_amp", "")),
}
def _parse_display_amp(raw):
if not raw or not str(raw).strip():
return np.ones(3)
try:
import ast as _ast
v = np.array(_ast.literal_eval(str(raw).strip()), dtype=np.float64)
return v if v.shape == (3,) else np.ones(3)
except Exception:
return np.ones(3)
def compute_energy(x, y, z, vx, vy, vz, masses, mass_arr,
bond_pairs, bond_stiffness, bond_rest_lengths,
gravity_field, G, gravity_interaction, gravity_strength):
@@ -320,6 +332,70 @@ def compute_energy_flux(x, y, z, vx, vy, vz,
return flux, bond_xpos
def compute_driver_work_power(x, y, z, vx, vy, vz,
bond_pairs, bond_stiffness, bond_rest_lengths,
atom_ids, driver_info):
"""计算每个驱动原子通过键对系统（非驱动原子）做功的功率。
对于驱动原子 d 与系统原子 j 之间的键：
P_{d→j} = F_{d→j} · v_j
其中 F_{d→j} 是键对系统原子 j 的弹簧力。
Returns:
drv_powers: dict {atom_id: (n_frames,)} 每个驱动原子的瞬时功率
total_power: (n_frames,) 所有驱动原子功率之和
"""
if bond_pairs is None or len(bond_pairs) == 0:
n_frames = x.shape[0]
return {}, np.zeros(n_frames)
id_to_idx = {int(aid): i for i, aid in enumerate(atom_ids)}
driven_idx = {id_to_idx[aid] for aid in driver_info if aid in id_to_idx}
n_frames = x.shape[0]
drv_powers = {}
for b in range(len(bond_pairs)):
ii, jj = int(bond_pairs[b, 0]), int(bond_pairs[b, 1])
i_drv = ii in driven_idx
j_drv = jj in driven_idx
if i_drv == j_drv: # 两端同为驱动或同为自由，跳过
continue
drv_loc = ii if i_drv else jj # 驱动端 index
sys_loc = jj if i_drv else ii # 系统端 index
drv_aid = int(atom_ids[drv_loc])
dx_ = x[:, jj] - x[:, ii]
dy_ = y[:, jj] - y[:, ii]
dz_ = z[:, jj] - z[:, ii]
dist = np.sqrt(dx_**2 + dy_**2 + dz_**2)
dist = np.maximum(dist, 1e-12)
k = bond_stiffness[b]
r0 = bond_rest_lengths[b]
fac = k * (dist - r0) / dist # 标量弹力因子
# 作用于系统原子的弹簧力：指向驱动原子方向
if i_drv: # drv=i, sys=j: 力方向 j→i，即 -(dx_, dy_, dz_)
fx = -fac * dx_
fy = -fac * dy_
fz = -fac * dz_
else: # drv=j, sys=i: 力方向 i→j，即 +(dx_, dy_, dz_)
fx = fac * dx_
fy = fac * dy_
fz = fac * dz_
power = fx * vx[:, sys_loc] + fy * vy[:, sys_loc] + fz * vz[:, sys_loc]
if drv_aid not in drv_powers:
drv_powers[drv_aid] = np.zeros(n_frames)
drv_powers[drv_aid] += power
total_power = sum(drv_powers.values()) if drv_powers else np.zeros(n_frames)
return drv_powers, total_power
def plot_wave(output_dir, save_gif=False, save_mp4=False, show=True):
"""主绘图函数：读取 display.txt 并生成波形+能量动画。
@@ -370,26 +446,106 @@ def plot_wave(output_dir, save_gif=False, save_mp4=False, show=True):
dy = y - pos_0[np.newaxis, :, 1]
dz = z - pos_0[np.newaxis, :, 2]
# ── 系统总能量（用于右下时间图）──
ek_sys, us_sys, ug_sys, ugr_sys = compute_energy(
x, y, z, vx, vy, vz, masses, masses,
bond_pairs, bond_stiffness, bond_rest_lengths,
gravity_field, G, gravity_interaction, gravity_strength)
e_total = ek_sys + us_sys + ug_sys + ugr_sys
power = np.gradient(e_total, t)
# ── 每粒子能量 ──
# ── 每粒子能量（图2 与图3 共用同一套计算）──
ek_atom, pe_atom, et_atom = compute_per_atom_energy(
x, y, z, vx, vy, vz, masses,
bond_pairs, bond_stiffness, bond_rest_lengths,
atom_ids, driver_info)
# ── 系统总能量 = 各粒子求和（与图2 完全一致）──
ek_sys = np.sum(ek_atom, axis=1)
us_sys = np.sum(pe_atom, axis=1)
e_total = np.sum(et_atom, axis=1)
power = np.gradient(e_total, t)
# 重力势能：仍用原有函数提供（若启用重力场）
_, _, ug_sys, ugr_sys = compute_energy(
x, y, z, vx, vy, vz, masses, masses,
bond_pairs, bond_stiffness, bond_rest_lengths,
gravity_field, G, gravity_interaction, gravity_strength)
if gravity_field or gravity_interaction:
e_total = e_total + ug_sys + ugr_sys
power = np.gradient(e_total, t)
# ── 能流密度 ──
flux, bond_xpos = compute_energy_flux(
x, y, z, vx, vy, vz,
bond_pairs, bond_stiffness, bond_rest_lengths)
# ── y 轴范围 ──
# ── 驱动做功功率 ──
drv_powers, total_drv_power = compute_driver_work_power(
x, y, z, vx, vy, vz,
bond_pairs, bond_stiffness, bond_rest_lengths,
atom_ids, driver_info)
# ── 原子可视化预计算 ──
display_amp = np.array(data.get("display_amp", [1.0, 1.0, 1.0]), dtype=np.float64)
eq_x_vis = pos_0[:, 0]
eq_z_vis = pos_0[:, 2]
# 视觉坐标 = 平衡位置 + 放大的位移
x_vis = eq_x_vis + (x - eq_x_vis) * display_amp[0] # (n_frames, n_atoms)
z_vis = eq_z_vis + (z - eq_z_vis) * display_amp[2]
# 找边界原子（与驱动原子成键的系统原子）及对应键
id_to_idx_vis = {int(aid): i for i, aid in enumerate(atom_ids)}
driven_set_vis = {id_to_idx_vis[aid] for aid in driver_info if aid in id_to_idx_vis}
bond_boundary_list = [] # (drv_idx, sys_idx, bond_b)
for _b in range(len(bond_pairs)):
_ii, _jj = int(bond_pairs[_b, 0]), int(bond_pairs[_b, 1])
if (_ii in driven_set_vis) ^ (_jj in driven_set_vis):
_drv = _ii if _ii in driven_set_vis else _jj
_sys = _jj if _ii in driven_set_vis else _ii
bond_boundary_list.append((_drv, _sys, _b))
# 唯一边界原子索引列表
_bnd_set = {}
for _drv, _sys, _b in bond_boundary_list:
if _sys not in _bnd_set:
_bnd_set[_sys] = len(_bnd_set)
boundary_atom_idx = np.array(list(_bnd_set.keys()), dtype=int)
n_boundary = len(boundary_atom_idx)
_lat = (eq_x_vis[-1] - eq_x_vis[0]) / max(n_atoms - 1, 1)
_z_all = z_vis.reshape(-1)
_z_min, _z_max = np.min(_z_all), np.max(_z_all)
_z_mg = max((_z_max - _z_min) * 0.2, _lat * 2)
_z_range = max((_z_max - _z_min) + 2 * _z_mg, _lat * 4)
_arrow_len = _z_range * 0.50 # 箭头最大显示长度 = 纵坐标范围的 50%
# 预计算边界原子受到的驱动力
# 方向：沿显示坐标下的键方向（消除坐标轴比例失真）；大小：胡克力模 k|d-r0|
bnd_fx_scaled = np.zeros((n_frames, max(n_boundary, 1)))
bnd_fz_scaled = np.zeros((n_frames, max(n_boundary, 1)))
_f_mag_all = []
for _drv, _sys, _b in bond_boundary_list:
_bi = _bnd_set[_sys]
# 物理键长
_dx3 = x[:, _drv] - x[:, _sys]
_dy3 = y[:, _drv] - y[:, _sys]
_dz3 = z[:, _drv] - z[:, _sys]
_dist = np.maximum(np.sqrt(_dx3**2 + _dy3**2 + _dz3**2), 1e-12)
# 有符号力大小（正 = 拉向驱动原子，负 = 推离）
_f_signed = bond_stiffness[_b] * (_dist - bond_rest_lengths[_b])
# 显示坐标下的键方向（x-z 平面）
_dx_d = x_vis[:, _drv] - x_vis[:, _sys]
_dz_d = z_vis[:, _drv] - z_vis[:, _sys]
_disp_len = np.maximum(np.sqrt(_dx_d**2 + _dz_d**2), 1e-12)
bnd_fx_scaled[:, _bi] += _f_signed * _dx_d / _disp_len
bnd_fz_scaled[:, _bi] += _f_signed * _dz_d / _disp_len
_f_mag_all.append(np.abs(_f_signed))
_f_max = np.max(_f_mag_all) if _f_mag_all else 1.0
_f_max = _f_max if _f_max > 1e-20 else 1.0
bnd_fx_scaled = bnd_fx_scaled / _f_max * _arrow_len
bnd_fz_scaled = bnd_fz_scaled / _f_max * _arrow_len
# 边界原子速度（方向沿实际速度，大小归一化）
bnd_vx_raw = vx[:, boundary_atom_idx] if n_boundary > 0 else np.zeros((n_frames, 1))
bnd_vz_raw = vz[:, boundary_atom_idx] if n_boundary > 0 else np.zeros((n_frames, 1))
_v_max = np.max(np.sqrt(bnd_vx_raw**2 + bnd_vz_raw**2)) if n_boundary > 0 else 1.0
_v_max = _v_max if _v_max > 1e-20 else 1.0
bnd_vx_scaled = bnd_vx_raw / _v_max * _arrow_len
bnd_vz_scaled = bnd_vz_raw / _v_max * _arrow_len
# y 轴范围 ──
def get_ylim(arr):
vmax = np.max(np.abs(arr))
if vmax < 1e-10:
@@ -413,8 +569,9 @@ def plot_wave(output_dir, save_gif=False, save_mp4=False, show=True):
energy_vmax = energy_vmax if energy_vmax > 1e-12 else 1.0
energy_ylim = (0.0, energy_vmax * 1.2)
e_max = max(np.max(e_total), 0.01) * 1.3
p_max = max(np.max(np.abs(power)) * 1.3, 0.01)
e_max = max(np.max(e_total), 1e-12)
e_min = min(np.min(e_total), 0.0)
p_max = max(np.percentile(np.abs(power), 95) if len(power) > 0 else 0.0, 0.0)
# 能流 y 轴范围（对称，正负各半）
if flux.size > 0:
@@ -426,100 +583,282 @@ def plot_wave(output_dir, save_gif=False, save_mp4=False, show=True):
atom_idx = np.arange(n_atoms)
# ── 图形布局：4 行 × 1 列，纵向排列 ──
# ── 驱动/非驱动粒子能量（右下图）──
id_to_idx = {int(aid): i for i, aid in enumerate(atom_ids)}
driven_idx = np.array([id_to_idx[aid] for aid in driver_info if aid in id_to_idx], dtype=int)
free_idx = np.setdiff1d(np.arange(n_atoms), driven_idx)
has_driver = len(driven_idx) > 0
if has_driver:
ek_drv = np.sum(ek_atom[:, driven_idx], axis=1)
ep_drv = np.sum(pe_atom[:, driven_idx], axis=1)
ek_free = np.sum(ek_atom[:, free_idx], axis=1)
ep_free = np.sum(pe_atom[:, free_idx], axis=1)
# ── 图形布局：左3行、右3行（subplot_mosaic）──
plt.rcParams['font.sans-serif'] = ['Microsoft YaHei', 'SimHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False
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.42, top=0.95)
from matplotlib.collections import LineCollection as _LC
# ── 图1：x/y/z 位移波形叠加 ──
fig, axes = plt.subplot_mosaic(
[['atoms', 'ep'],
['wave', 'drv'],
['energy', 'pwr']],
figsize=(20, 15))
ax_atoms = axes['atoms']
ax_wave = axes['wave']
ax_ep = axes['ep']
ax_energy = axes['energy']
ax_drv = axes['drv']
ax_pwr = axes['pwr']
fig.subplots_adjust(hspace=0.45, wspace=0.32, top=0.97)
# ── 左上：原子位置 + 键 + 力/速度箭头 ──
_x_min, _x_max = eq_x_vis[0], eq_x_vis[-1]
ax_atoms.set_xlim(_x_min - _lat, _x_max + _lat)
ax_atoms.set_ylim(_z_min - _z_mg, _z_max + _z_mg)
ax_atoms.set_xlabel("位置 $x$")
ax_atoms.set_ylabel("位移 $z$（放大 {:.0f}×）".format(display_amp[2]))
ax_atoms.set_title("原子运动（红=驱动，箭头：红=驱动力，蓝=边界速度）")
ax_atoms.set_aspect('auto')
ax_atoms.grid(True, alpha=0.2)
# 键线段（LineCollection，初始帧）
def _make_bond_segs(frame_idx):
segs = []
for _b in range(len(bond_pairs)):
_ii, _jj = int(bond_pairs[_b, 0]), int(bond_pairs[_b, 1])
segs.append([(x_vis[frame_idx, _ii], z_vis[frame_idx, _ii]),
(x_vis[frame_idx, _jj], z_vis[frame_idx, _jj])])
return segs
_bond_lc = _LC(_make_bond_segs(0), colors='#888888', linewidths=0.8, zorder=1)
ax_atoms.add_collection(_bond_lc)
# 散点：自由原子（黑）
_free_mask = np.array([i not in driven_set_vis for i in range(n_atoms)])
_scat_free, = ax_atoms.plot(
x_vis[0, _free_mask], z_vis[0, _free_mask],
'o', color='black', ms=4, zorder=3)
# 散点：驱动原子（红）
_drv_mask = ~_free_mask
_scat_drv, = ax_atoms.plot(
x_vis[0, _drv_mask], z_vis[0, _drv_mask],
'o', color='red', ms=6, zorder=4)
# 力箭头（红，边界原子）
_q_force = ax_atoms.quiver(
x_vis[0, boundary_atom_idx] if n_boundary > 0 else [],
z_vis[0, boundary_atom_idx] if n_boundary > 0 else [],
bnd_fx_scaled[0] if n_boundary > 0 else [],
bnd_fz_scaled[0] if n_boundary > 0 else [],
color='red', angles='xy', scale_units='xy', scale=1,
width=0.007, headwidth=5, headlength=5, zorder=5)
# 速度箭头（蓝，边界原子）
_q_vel = ax_atoms.quiver(
x_vis[0, boundary_atom_idx] if n_boundary > 0 else [],
z_vis[0, boundary_atom_idx] if n_boundary > 0 else [],
bnd_vx_scaled[0] if n_boundary > 0 else [],
bnd_vz_scaled[0] if n_boundary > 0 else [],
color='blue', angles='xy', scale_units='xy', scale=1,
width=0.007, headwidth=5, headlength=5, zorder=5)
# ── 左上：x/y/z 位移波形 ──
ax_wave.set_xlim(0, n_atoms - 1)
ax_wave.set_ylim(disp_ylim)
ax_wave.set_xlabel("原子序号")
ax_wave.set_ylabel("位移")
ax_wave.set_title("粒子位移（x / y / z 方向）")
ax_wave.set_ylabel("位移 $u$")
ax_wave.set_title("粒子位移（$x$ / $y$ / $z$ 方向）")
ax_wave.grid(True, alpha=0.3)
wave_disps = [dx, dy, dz]
wave_labels = ["x 方向（纵波）", "y 方向（横波）", "z 方向（横波）"]
wave_labels = ["$u_x$（纵波）", "$u_y$（横波）", "$u_z$（横波）"]
wave_colors = ["#2563eb", "#ea580c", "#16a34a"]
wave_lines = []
for label, color in zip(wave_labels, wave_colors):
ln, = ax_wave.plot([], [], color=color, linewidth=1.5, label=label)
wave_lines.append(ln)
ax_wave.legend(loc="upper right", fontsize=9)
time_text = ax_wave.text(0.02, 0.95, "", transform=ax_wave.transAxes,
fontsize=10, verticalalignment="top")
_dt_frame = (t[1] - t[0]) if len(t) > 1 else 0.0
_t_total_str = f"{t[-1] + _dt_frame:.2f} s"
_time_axes = [ax_atoms, ax_wave, ax_energy, ax_ep, ax_drv, ax_pwr]
time_texts = [
ax.text(0.02, 0.97, "", transform=ax.transAxes,
fontsize=9, verticalalignment="top",
bbox=dict(boxstyle="round,pad=0.2", fc="white", alpha=0.7))
for ax in _time_axes
]
time_text = time_texts[1] # 保留旧名兼容下面的代码
# ── 图2：每粒子动能、势能、总能叠加 ──
# ── 左下：每粒子能量（左轴）+ 能流密度（右轴）──
ax_energy.set_xlim(0, n_atoms - 1)
ax_energy.set_ylim(energy_ylim)
ax_energy.set_xlabel("原子序号")
ax_energy.set_ylabel("能量")
ax_energy.set_title("每粒子能量（动能 / 势能 / 总能）")
ax_energy.set_xlabel("原子序号 / 键位置")
ax_energy.set_ylabel("能量 $E$")
ax_energy.set_title(
r"每粒子能量（$E_k$/$E_p$/$E_{tot}$）与能流密度 $J$"
)
ax_energy.grid(True, alpha=0.3)
energy_arrays = [ek_atom, pe_atom, et_atom]
energy_labels = ["动能", "势能", "总能"]
energy_colors = ["#1d4ed8", "#b45309", "#7c3aed"]
energy_labels = ["$E_k$（动能）", "$E_p$（势能）", "$E_{tot}$（总能）"]
energy_colors = ["#16a34a", "#b45309", "#7c3aed"]
energy_lines = []
for label, color in zip(energy_labels, energy_colors):
ln, = ax_energy.plot([], [], color=color, linewidth=1.5, label=label)
energy_lines.append(ln)
ax_energy.legend(loc="upper right", fontsize=9)
# ── 图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 = ax_energy.twinx()
ax_flux.set_ylim(flux_ylim)
ax_flux.set_ylabel("能流密度 $J$", color="#dc2626")
ax_flux.tick_params(axis='y', labelcolor="#dc2626")
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)
flux_line, = ax_flux.plot([], [], color="#dc2626", linewidth=1.5,
label="$J$（能流密度）")
handles_e, labels_e = ax_energy.get_legend_handles_labels()
handles_f, labels_f = ax_flux.get_legend_handles_labels()
ax_energy.legend(handles_e + handles_f, labels_e + labels_f,
loc="upper right", fontsize=9)
# ── 图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)
ax_ep.set_ylim(ep_ylow, ep_yhigh)
ax_ep.set_xlabel("时间 (s)")
ax_ep.set_ylabel("能量 / 功率")
ep_margin = (e_max - e_min) * 0.15 if e_max > e_min else e_max * 0.15
ax_ep.set_ylim(e_min - ep_margin, e_max + ep_margin)
ax_ep.set_clip_on(True)
ax_ep.set_xlabel("时间 $t$ (s)")
ax_ep.set_ylabel("能量 $E$ / 功率 $P$")
ax_ep.set_title("系统能量与输入功率")
ax_ep.grid(True, alpha=0.3)
ln_ek, = ax_ep.plot([], [], "b-", lw=1.5, label="动能")
ln_us, = ax_ep.plot([], [], "orange", lw=1.5, label="弹性势能")
ln_et, = ax_ep.plot([], [], "r--", lw=1.5, label="总能量")
ln_pw, = ax_ep.plot([], [], "g-", lw=1.5, alpha=0.7, label="输入功率 (dE/dt)")
ln_ek, = ax_ep.plot([], [], "b-", lw=1.5, label="$E_k$（动能）")
ln_us, = ax_ep.plot([], [], "orange", lw=1.5, label="$E_s$（弹性势能）")
ln_et, = ax_ep.plot([], [], "r--", lw=1.5, label="$E_{tot}$（总能量）")
ln_pw, = ax_ep.plot([], [], "g-", lw=1.5, alpha=0.7, label=r"$P_{in}=dE/dt$")
ln_ug = None
ln_ugr = None
if gravity_field:
ln_ug, = ax_ep.plot([], [], "purple", lw=1.0, alpha=0.5, label="重力势能")
ln_ug, = ax_ep.plot([], [], "purple", lw=1.0, alpha=0.5, label="$E_g$（重力势能）")
if gravity_interaction and n_atoms <= 200:
ln_ugr, = ax_ep.plot([], [], "brown", lw=1.0, alpha=0.5, label="万有引力势能")
ax_ep.legend(loc="upper left", fontsize=9)
ln_ugr, = ax_ep.plot([], [], "brown", lw=1.0, alpha=0.5, label="$E_{gr}$（万有引力势能）")
ax_ep.legend(loc="upper right", fontsize=9)
# ── 右下：驱动/非驱动粒子能量随时间 ──
ax_drv.set_xlim(t[0], t[-1])
ax_drv.set_xlabel("时间 $t$ (s)")
ax_drv.set_ylabel("能量 $E$")
ax_drv.grid(True, alpha=0.3)
ln_ek_drv = ln_ep_drv = ln_ek_free = ln_ep_free = None
ln_et_drv = ln_et_free = None
if has_driver:
et_drv = ek_drv + ep_drv
et_free = ek_free + ep_free
drv_ids = sorted(driver_info.keys())
ax_drv.set_title(f"驱动粒子（序号 {drv_ids}）向系统做功")
ln_ek_drv, = ax_drv.plot([], [], color="#dc2626", lw=1.2, linestyle="--",
label=r"$E_k^{drv}$（驱动动能）")
ln_ep_drv, = ax_drv.plot([], [], color="#f97316", lw=1.2, linestyle="--",
label=r"$E_p^{drv}$（驱动势能）")
ln_et_drv, = ax_drv.plot([], [], color="#7f1d1d", lw=2.0,
label=r"$E_{tot}^{drv}$（驱动总能）")
ln_ek_free, = ax_drv.plot([], [], color="#2563eb", lw=1.2, linestyle="--",
label=r"$E_k^{sys}$（系统动能）")
ln_ep_free, = ax_drv.plot([], [], color="#16a34a", lw=1.2, linestyle="--",
label=r"$E_p^{sys}$（系统势能）")
ln_et_free, = ax_drv.plot([], [], color="#1e3a5f", lw=2.0,
label=r"$E_{tot}^{sys}$（系统总能）")
ax_drv.legend(loc="upper right", fontsize=8)
# y 轴一次定好
_drv_all = np.concatenate([ek_drv, ep_drv, et_drv, ek_free, ep_free, et_free])
_dy_max = np.max(_drv_all)
_dy_min = np.min(_drv_all)
_dy_mg = (_dy_max - _dy_min) * 0.15 if _dy_max > _dy_min else abs(_dy_max) * 0.15 + 1e-12
ax_drv.set_ylim(_dy_min - _dy_mg, _dy_max + _dy_mg)
else:
ax_drv.set_title("驱动粒子能量（无驱动力）")
ax_drv.text(0.5, 0.5, "无驱动力", transform=ax_drv.transAxes,
ha="center", va="center", fontsize=12, color="gray")
# ── 右下：驱动做功功率 ──
ax_pwr.set_xlim(t[0], t[-1])
ax_pwr.set_xlabel("时间 $t$ (s)")
ax_pwr.set_ylabel("功率 $P$")
ax_pwr.set_title("驱动原子对系统做功的功率 $P = \\mathbf{F}_{bond}\\cdot\\mathbf{v}_{sys}$")
ax_pwr.axhline(0, color="gray", linewidth=0.8, linestyle="--")
ax_pwr.grid(True, alpha=0.3)
pwr_colors = ["#dc2626", "#2563eb", "#16a34a", "#f97316", "#7c3aed"]
ln_pwr_each = {} # aid -> Line2D
if drv_powers:
for idx_d, (aid, _) in enumerate(sorted(drv_powers.items())):
color = pwr_colors[idx_d % len(pwr_colors)]
ln, = ax_pwr.plot([], [], color=color, lw=1.2, linestyle="--",
label=f"$P_{{drv,{aid}}}$（原子 {aid}）")
ln_pwr_each[aid] = ln
ln_pwr_total, = ax_pwr.plot([], [], color="black", lw=2.0,
label=r"$P_{total}$（总功率）")
ax_pwr.legend(loc="upper right", fontsize=9)
# y 轴一次定好
if drv_powers:
_pw_all = np.concatenate(list(drv_powers.values()) + [total_drv_power])
_pw_max = np.max(_pw_all)
_pw_min = np.min(_pw_all)
_pw_mg = (_pw_max - _pw_min) * 0.15 if _pw_max > _pw_min else abs(_pw_max) * 0.15 + 1e-12
ax_pwr.set_ylim(_pw_min - _pw_mg, _pw_max + _pw_mg)
# ── 动画更新 ──
def update(frame):
# 图1：位移波形
# 每轮开始时清屏
# 左上：原子位置动画
_bond_lc.set_segments(_make_bond_segs(frame))
_scat_free.set_xdata(x_vis[frame, _free_mask])
_scat_free.set_ydata(z_vis[frame, _free_mask])
_scat_drv.set_xdata(x_vis[frame, _drv_mask])
_scat_drv.set_ydata(z_vis[frame, _drv_mask])
if n_boundary > 0:
_q_force.set_offsets(
np.column_stack([x_vis[frame, boundary_atom_idx],
z_vis[frame, boundary_atom_idx]]))
_q_force.set_UVC(bnd_fx_scaled[frame], bnd_fz_scaled[frame])
_q_vel.set_offsets(
np.column_stack([x_vis[frame, boundary_atom_idx],
z_vis[frame, boundary_atom_idx]]))
_q_vel.set_UVC(bnd_vx_scaled[frame], bnd_vz_scaled[frame])
if frame == 0:
all_clear = list(wave_lines) + list(energy_lines) + [flux_line]
all_clear += [ln for ln in [ln_ek, ln_us, ln_et, ln_pw, ln_ug, ln_ugr]
if ln is not None]
if has_driver:
all_clear += [ln for ln in [ln_ek_drv, ln_ep_drv, ln_et_drv,
ln_ek_free, ln_ep_free, ln_et_free]
if ln is not None]
all_clear += list(ln_pwr_each.values()) + [ln_pwr_total]
for ln in all_clear:
ln.set_data([], [])
_tstr0 = f"t = {t[0]:.2f} s / {_t_total_str} | 帧 1/{n_frames}"
for _tt in time_texts:
_tt.set_text(_tstr0)
return all_clear + time_texts + [_bond_lc, _scat_free, _scat_drv,
_q_force, _q_vel]
# 左中：位移波形
for i, ln in enumerate(wave_lines):
ln.set_data(atom_idx, wave_disps[i][frame])
time_text.set_text(f"t = {t[frame]:.2f} s | 帧 {frame+1}/{n_frames}")
_tstr = f"t = {t[frame]:.2f} s / {_t_total_str} | 帧 {frame+1}/{n_frames}"
for _tt in time_texts:
_tt.set_text(_tstr)
# 图2：每粒子能量
# 左下：每粒子能量 + 能流密度
for i, ln in enumerate(energy_lines):
ln.set_data(atom_idx, energy_arrays[i][frame])
# 图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])
@@ -527,15 +866,33 @@ def plot_wave(output_dir, save_gif=False, save_mp4=False, show=True):
ln_pw.set_data(cur_t, power[:frame + 1])
if ln_ug: ln_ug.set_data(cur_t, ug_sys[:frame + 1])
if ln_ugr: ln_ugr.set_data(cur_t, ugr_sys[:frame + 1])
ax_ep.set_xlim(t[0], max(t[frame] + max(t[-1] * 0.05, 1), t[-1]))
artists = wave_lines + [time_text] + energy_lines + \
[flux_line, ln_ek, ln_us, ln_et, ln_pw]
if ln_ug: artists.append(ln_ug)
if ln_ugr: artists.append(ln_ugr)
# 右中：驱动/系统粒子能量（累计）
if has_driver:
ln_ek_drv.set_data( cur_t, ek_drv[:frame + 1])
ln_ep_drv.set_data( cur_t, ep_drv[:frame + 1])
ln_et_drv.set_data( cur_t, et_drv[:frame + 1])
ln_ek_free.set_data(cur_t, ek_free[:frame + 1])
ln_ep_free.set_data(cur_t, ep_free[:frame + 1])
ln_et_free.set_data(cur_t, et_free[:frame + 1])
# 右下：驱动做功功率（累计）
for aid, ln in ln_pwr_each.items():
ln.set_data(cur_t, drv_powers[aid][:frame + 1])
ln_pwr_total.set_data(cur_t, total_drv_power[:frame + 1])
artists = (wave_lines + time_texts + energy_lines +
[flux_line, ln_ek, ln_us, ln_et, ln_pw])
if ln_ug: artists.append(ln_ug)
if ln_ugr: artists.append(ln_ugr)
if has_driver:
artists += [ln_ek_drv, ln_ep_drv, ln_et_drv,
ln_ek_free, ln_ep_free, ln_et_free]
artists += list(ln_pwr_each.values()) + [ln_pwr_total]
artists += [_bond_lc, _scat_free, _scat_drv, _q_force, _q_vel]
return artists
ani = FuncAnimation(fig, update, frames=n_frames, interval=50, blit=True)
ani = FuncAnimation(fig, update, frames=n_frames, interval=50, blit=True, repeat=True)
# ── 输出文件 ──
gif_path = None