332 lines
12 KiB
Python
332 lines
12 KiB
Python
import numpy as np
|
||
import matplotlib.pyplot as plt
|
||
from matplotlib.animation import FuncAnimation
|
||
|
||
# 设置中文字体
|
||
plt.rcParams['font.sans-serif'] = ['SimHei']
|
||
plt.rcParams['axes.unicode_minus'] = False
|
||
|
||
# 简谐振子系统参数
|
||
k = 1.0 # 劲度系数
|
||
m = 1.0 # 质量
|
||
omega = np.sqrt(k/m) # 角频率
|
||
|
||
def potential_energy(x):
|
||
"""势能: (1/2)*k*x^2"""
|
||
return 0.5 * k * x**2
|
||
|
||
def kinetic_energy(v):
|
||
"""动能: (1/2)*m*v^2"""
|
||
return 0.5 * m * v**2
|
||
|
||
def total_energy(x, v):
|
||
"""总能量"""
|
||
return potential_energy(x) + kinetic_energy(v)
|
||
|
||
def acceleration(x):
|
||
"""加速度: a = -(k/m)*x"""
|
||
return -omega**2 * x
|
||
|
||
# ============ 四种数值方法 ============
|
||
|
||
def euler_explicit(x0, v0, dt, n_steps):
|
||
"""欧拉显式 (前向欧拉)"""
|
||
x = np.zeros(n_steps)
|
||
v = np.zeros(n_steps)
|
||
t = np.zeros(n_steps)
|
||
x[0] = x0
|
||
v[0] = v0
|
||
t[0] = 0
|
||
for i in range(n_steps - 1):
|
||
x[i+1] = x[i] + v[i] * dt
|
||
v[i+1] = v[i] + acceleration(x[i]) * dt
|
||
t[i+1] = t[i] + dt
|
||
return t, x, v
|
||
|
||
|
||
def euler_implicit(x0, v0, dt, n_steps):
|
||
"""欧拉隐式 (后向欧拉)"""
|
||
x = np.zeros(n_steps)
|
||
v = np.zeros(n_steps)
|
||
t = np.zeros(n_steps)
|
||
x[0] = x0
|
||
v[0] = v0
|
||
t[0] = 0
|
||
for i in range(n_steps - 1):
|
||
# 解析解：对简谐振子联立方程可直接得出
|
||
x[i+1] = (x[i] + v[i] * dt) / (1 + omega**2 * dt**2)
|
||
v[i+1] = v[i] + acceleration(x[i+1]) * dt
|
||
t[i+1] = t[i] + dt
|
||
return t, x, v
|
||
|
||
|
||
def midpoint_method(x0, v0, dt, n_steps):
|
||
"""中点差分法 (二阶 Runge-Kutta / 改进欧拉)"""
|
||
x = np.zeros(n_steps)
|
||
v = np.zeros(n_steps)
|
||
t = np.zeros(n_steps)
|
||
x[0] = x0
|
||
v[0] = v0
|
||
t[0] = 0
|
||
for i in range(n_steps - 1):
|
||
x_mid = x[i] + v[i] * dt / 2
|
||
v_mid = v[i] + acceleration(x[i]) * dt / 2
|
||
x[i+1] = x[i] + v_mid * dt
|
||
v[i+1] = v[i] + acceleration(x_mid) * dt
|
||
t[i+1] = t[i] + dt
|
||
return t, x, v
|
||
|
||
|
||
def leapfrog_method(x0, v0, dt, n_steps):
|
||
"""蛙跳法 (Leapfrog / Störmer-Verlet)
|
||
|
||
核心思想：位置和速度交错半步更新
|
||
v_{n+1/2} = v_{n-1/2} + a(x_n) * dt <- 速度在半整数时刻
|
||
x_{n+1} = x_n + v_{n+1/2} * dt <- 位置在整数时刻
|
||
|
||
启动：用欧拉半步得到 v_{1/2}，之后进入交错循环。
|
||
输出时将速度还原到整数时刻：v_n = (v_{n-1/2} + v_{n+1/2}) / 2
|
||
"""
|
||
x = np.zeros(n_steps)
|
||
v = np.zeros(n_steps) # 整数时刻速度（仅用于输出/能量计算）
|
||
t = np.zeros(n_steps)
|
||
|
||
x[0] = x0
|
||
v[0] = v0
|
||
t[0] = 0
|
||
|
||
# 用欧拉半步初始化 v_{1/2}
|
||
v_half = v0 + acceleration(x0) * dt / 2
|
||
|
||
for i in range(n_steps - 1):
|
||
# 更新位置（整数时刻）
|
||
x[i+1] = x[i] + v_half * dt
|
||
t[i+1] = t[i] + dt
|
||
|
||
# 更新半步速度（用新位置的加速度）
|
||
v_half_new = v_half + acceleration(x[i+1]) * dt
|
||
|
||
# 整数时刻速度（用于能量计算和输出）：取两个半步的平均
|
||
v[i+1] = (v_half + v_half_new) / 2
|
||
|
||
v_half = v_half_new
|
||
|
||
return t, x, v
|
||
|
||
|
||
def exact_solution(x0, v0, t):
|
||
"""精确解"""
|
||
A = np.sqrt(x0**2 + (v0 / omega)**2)
|
||
phi = np.arctan2(-v0 / omega, x0)
|
||
x_exact = A * np.cos(omega * t + phi)
|
||
v_exact = -A * omega * np.sin(omega * t + phi)
|
||
return x_exact, v_exact
|
||
|
||
|
||
# ============ 参数设置 ============
|
||
dt = 0.05 # 时间步长
|
||
T_total = 10.0 # 总时间
|
||
n_steps = int(T_total / dt) + 1
|
||
|
||
# 三个初始条件 (x0, v0)
|
||
initial_conditions = [
|
||
(1.0, 0.0), # 从最大位移释放
|
||
(0.0, 1.0), # 从平衡位置给初速度
|
||
(0.8, 0.6) # 混合初始条件
|
||
]
|
||
|
||
# 四种方法：(名称, 函数, 线型颜色, alpha)
|
||
methods = [
|
||
('欧拉显式', euler_explicit, 'r-', 0.8),
|
||
('欧拉隐式', euler_implicit, 'b-', 0.8),
|
||
('中点差分法', midpoint_method, 'g-', 0.8),
|
||
('蛙跳法', leapfrog_method, 'm-', 0.8),
|
||
]
|
||
|
||
colors_ic = ['#e63946', '#457b9d', '#2d6a4f'] # 三种初始条件的颜色
|
||
labels_ic = ['从最大位移释放', '平衡位置初速度', '混合条件']
|
||
|
||
# 四种方法对应的绘图颜色（用于位移/速度时序图）
|
||
method_colors = ['r', 'b', 'g', 'm']
|
||
|
||
print(f"时间步长 dt = {dt}s, 步数 = {n_steps}")
|
||
|
||
# ============ 计算结果 ============
|
||
results = {}
|
||
for method_name, method_func, color_style, alpha in methods:
|
||
results[method_name] = {}
|
||
for idx, (x0, v0) in enumerate(initial_conditions):
|
||
t_arr, x_arr, v_arr = method_func(x0, v0, dt, n_steps)
|
||
E = total_energy(x_arr, v_arr)
|
||
results[method_name][idx] = {
|
||
't': t_arr, 'x': x_arr, 'v': v_arr, 'E': E,
|
||
'E_initial': E[0], 'E_final': E[-1],
|
||
'E_drift': (E[-1] - E[0]) / E[0] * 100
|
||
}
|
||
|
||
# 精确解（高密度时间轴）
|
||
exact_t = np.linspace(0, T_total, 1000)
|
||
|
||
# ============ 图形布局 ============
|
||
fig = plt.figure(figsize=(18, 12))
|
||
fig.suptitle('简谐振子数值方法四方比较', fontsize=16, fontweight='bold')
|
||
|
||
gs = fig.add_gridspec(3, 4, hspace=0.40, wspace=0.30)
|
||
ax_phase = fig.add_subplot(gs[0, :]) # 第1行全宽：相空间
|
||
ax_energy = fig.add_subplot(gs[1, :]) # 第2行全宽：能量演化
|
||
ax_traj = fig.add_subplot(gs[2, 0]) # 第3行：位置-时间
|
||
ax_velocity = fig.add_subplot(gs[2, 1]) # 第3行：速度-时间
|
||
ax_err_bar = fig.add_subplot(gs[2, 2]) # 第3行：能量误差柱状图
|
||
ax_animation= fig.add_subplot(gs[2, 3]) # 第3行：相空间动画
|
||
|
||
# ============ 1. 相空间轨迹 ============
|
||
ax_phase.set_title('相空间轨迹比较 (x vs v)', fontsize=13)
|
||
ax_phase.set_xlabel('位置 x', fontsize=11)
|
||
ax_phase.set_ylabel('速度 v', fontsize=11)
|
||
ax_phase.grid(True, alpha=0.3)
|
||
ax_phase.axhline(y=0, color='k', linestyle='-', alpha=0.3)
|
||
ax_phase.axvline(x=0, color='k', linestyle='-', alpha=0.3)
|
||
|
||
line_styles_ic = ['-', '--', ':']
|
||
method_line_colors = ['r', 'b', 'g', 'm']
|
||
|
||
for midx, (x0, v0) in enumerate(initial_conditions):
|
||
x_exact, v_exact = exact_solution(x0, v0, exact_t)
|
||
ax_phase.plot(x_exact, v_exact, 'k-', alpha=0.25, linewidth=1,
|
||
label='精确解' if midx == 0 else "")
|
||
for mi, (method_name, _, _, _) in enumerate(methods):
|
||
res = results[method_name][midx]
|
||
ax_phase.plot(res['x'], res['v'],
|
||
linestyle=['-', '--', ':', '-.'][midx],
|
||
color=method_line_colors[mi], alpha=0.75, linewidth=1.4,
|
||
label=method_name if midx == 0 else "")
|
||
|
||
handles, labels_leg = ax_phase.get_legend_handles_labels()
|
||
by_label = dict(zip(labels_leg, handles))
|
||
ax_phase.legend(by_label.values(), by_label.keys(),
|
||
loc='upper right', fontsize=8, ncol=2)
|
||
|
||
# ============ 2. 能量演化（归一化） ============
|
||
ax_energy.set_title('归一化能量随时间演化 E(t)/E(0)', fontsize=13)
|
||
ax_energy.set_xlabel('时间 t (s)', fontsize=11)
|
||
ax_energy.set_ylabel('E / E0', fontsize=11)
|
||
ax_energy.grid(True, alpha=0.3)
|
||
|
||
energy_linestyles = ['-', '--', ':', '-.']
|
||
for mi, (method_name, _, _, _) in enumerate(methods):
|
||
res = results[method_name][0] # 只用第一个初始条件展示能量
|
||
E_norm = res['E'] / res['E_initial']
|
||
ax_energy.plot(res['t'], E_norm,
|
||
linestyle=energy_linestyles[mi],
|
||
color=method_line_colors[mi],
|
||
linewidth=1.8, alpha=0.9, label=method_name)
|
||
|
||
ax_energy.axhline(y=1.0, color='k', linestyle='-', alpha=0.5, linewidth=0.8)
|
||
ax_energy.set_ylim(0.85, 1.70)
|
||
ax_energy.legend(loc='upper left', fontsize=9)
|
||
|
||
# ============ 3. 位置-时间 (第一初始条件) ============
|
||
idx_show = 0
|
||
ax_traj.set_title(f'位置-时间 ({labels_ic[idx_show]})', fontsize=11)
|
||
ax_traj.set_xlabel('时间 t (s)', fontsize=10)
|
||
ax_traj.set_ylabel('位置 x', fontsize=10)
|
||
ax_traj.grid(True, alpha=0.3)
|
||
|
||
x_ex, v_ex = exact_solution(*initial_conditions[idx_show], exact_t)
|
||
ax_traj.plot(exact_t, x_ex, 'k-', alpha=0.5, linewidth=1.2, label='精确解')
|
||
for mi, (method_name, _, _, _) in enumerate(methods):
|
||
res = results[method_name][idx_show]
|
||
ax_traj.plot(res['t'], res['x'],
|
||
color=method_line_colors[mi], linewidth=1.4,
|
||
alpha=0.85, label=method_name)
|
||
ax_traj.legend(loc='upper right', fontsize=7)
|
||
|
||
# ============ 4. 速度-时间 (第一初始条件) ============
|
||
ax_velocity.set_title(f'速度-时间 ({labels_ic[idx_show]})', fontsize=11)
|
||
ax_velocity.set_xlabel('时间 t (s)', fontsize=10)
|
||
ax_velocity.set_ylabel('速度 v', fontsize=10)
|
||
ax_velocity.grid(True, alpha=0.3)
|
||
|
||
ax_velocity.plot(exact_t, v_ex, 'k-', alpha=0.5, linewidth=1.2, label='精确解')
|
||
for mi, (method_name, _, _, _) in enumerate(methods):
|
||
res = results[method_name][idx_show]
|
||
ax_velocity.plot(res['t'], res['v'],
|
||
color=method_line_colors[mi], linewidth=1.4,
|
||
alpha=0.85, label=method_name)
|
||
ax_velocity.legend(loc='upper right', fontsize=7)
|
||
|
||
# ============ 5. 能量误差柱状图 ============
|
||
ax_err_bar.set_title('能量漂移百分比 |dE/E0|*100%', fontsize=11)
|
||
ax_err_bar.set_ylabel('漂移 (%)', fontsize=10)
|
||
ax_err_bar.grid(True, alpha=0.3, axis='y')
|
||
|
||
method_names_short = ['显式', '隐式', '中点', '蛙跳']
|
||
bar_colors = method_line_colors
|
||
drifts = [abs(results[mn][0]['E_drift']) for mn, *_ in methods]
|
||
bars = ax_err_bar.bar(method_names_short, drifts, color=bar_colors, alpha=0.8, width=0.5)
|
||
for bar, val in zip(bars, drifts):
|
||
ax_err_bar.text(bar.get_x() + bar.get_width() / 2,
|
||
bar.get_height() + 0.3,
|
||
f'{val:.2f}%', ha='center', va='bottom', fontsize=9)
|
||
ax_err_bar.set_ylim(0, max(drifts) * 1.2 + 2)
|
||
|
||
# ============ 6. 动画：四种方法相空间对比 (同一初始条件) ============
|
||
ax_animation.set_xlim(-1.5, 1.5)
|
||
ax_animation.set_ylim(-1.5, 1.5)
|
||
ax_animation.set_title('四种方法同步释放 (相空间)', fontsize=11)
|
||
ax_animation.set_xlabel('位置 x', fontsize=10)
|
||
ax_animation.set_ylabel('速度 v', fontsize=10)
|
||
ax_animation.grid(True, alpha=0.3)
|
||
ax_animation.axhline(y=0, color='k', linestyle='-', alpha=0.3)
|
||
ax_animation.axvline(x=0, color='k', linestyle='-', alpha=0.3)
|
||
|
||
anim_points = []
|
||
anim_trails = []
|
||
anim_idx = 0 # 选第一组初始条件，比较四种方法的同步演化
|
||
for mi, (method_name, _, _, _) in enumerate(methods):
|
||
pt, = ax_animation.plot([], [], 'o', color=method_line_colors[mi],
|
||
markersize=7, label=method_name)
|
||
tr, = ax_animation.plot([], [], '-', color=method_line_colors[mi], alpha=0.35, linewidth=1)
|
||
anim_points.append(pt)
|
||
anim_trails.append(tr)
|
||
|
||
x_anim_exact, v_anim_exact = exact_solution(*initial_conditions[anim_idx], exact_t)
|
||
ax_animation.plot(x_anim_exact, v_anim_exact, 'k-', alpha=0.2, linewidth=0.8, label='精确解')
|
||
|
||
ax_animation.legend(loc='upper right', fontsize=8)
|
||
|
||
def animate(frame):
|
||
for mi, (method_name, _, _, _) in enumerate(methods):
|
||
res = results[method_name][anim_idx]
|
||
if frame < len(res['x']):
|
||
anim_points[mi].set_data([res['x'][frame]], [res['v'][frame]])
|
||
start = max(0, frame - 50)
|
||
anim_trails[mi].set_data(res['x'][start:frame+1], res['v'][start:frame+1])
|
||
return anim_points + anim_trails
|
||
|
||
anim = FuncAnimation(fig, animate, frames=n_steps, interval=20, blit=True)
|
||
|
||
# ============ 打印能量误差统计 ============
|
||
print("\n" + "=" * 65)
|
||
print("能量守恒误差分析 (相对漂移百分比，初始条件：从最大位移释放)")
|
||
print("=" * 65)
|
||
for method_name, _, _, _ in methods:
|
||
res0 = results[method_name][0]
|
||
print(f" {method_name:10s}：{res0['E_drift']:+.4f}%")
|
||
|
||
print("\n" + "=" * 65)
|
||
print("所有初始条件下的能量漂移")
|
||
print("=" * 65)
|
||
for method_name, _, _, _ in methods:
|
||
print(f"\n{method_name}:")
|
||
for idx, (x0, v0) in enumerate(initial_conditions):
|
||
res = results[method_name][idx]
|
||
print(f" {labels_ic[idx]}: {res['E_drift']:+.4f}%")
|
||
|
||
# 显示图形
|
||
plt.tight_layout()
|
||
plt.show()
|
||
|
||
# 可选：保存动画
|
||
# anim.save('harmonic_oscillator_4methods.gif', writer='pillow', fps=30)
|