def forward_euler_generic(rhs, u0, times, omega):
u = np.zeros((len(times), len(u0)))
u[0,:] = u0
for (i, t) in enumerate(times[1:]):
u[i+1,:] = u[i,:] + (times[i+1] - times[i]) * rhs(u[i,:], omega)
return u
def ssprk22_generic(rhs, u0, times, omega):
u = np.zeros((len(times), len(u0)))
u[0,:] = u0
for (i, t) in enumerate(times[1:]):
u_1 = u[i,:] + (times[i+1] - times[i]) * rhs(u[i,:], omega)
u[i+1,:] = 0.5 * u[i,:]+ 0.5 * u_1 + 0.5 * (times[i+1] - times[i]) * rhs(u_1, omega)
return u
def rk2_generic(rhs, u0, times, omega):
# Ralston's method
u = np.zeros((len(times), len(u0)))
u[0,:] = u0
for (i, t) in enumerate(times[1:]):
dt = (times[i+1] - times[i])
k1 = rhs(u[i,:], omega)
k2 = rhs(u[i,:] + 0.75 * k1 * dt, omega)
u[i+1,:] = u[i,:] + dt * (0.25 * k1 + 0.75 * k2)
return u
def harmonic_rhs(u, omega):
return np.array([u[1], -omega**2 * u[0]])
def energy(u, omega):
return 0.5 * (u[:,1]**2 + omega**2 * u[:,0]**2)
fig, axes = plt.subplots(figsize=(12.0, 4.2), ncols=2)
ax = axes[0]
ax_e = axes[1]
t_max = 5.0
omega = 3.0
times = np.linspace(0, t_max, 100)
u0 = [0.5, 0.5]
analytical = np.array([u0[0] * np.cos(omega * times) + (u0[0]/omega) * np.sin(omega * times),
-u0[0] * omega * np.sin(omega * times) + u0[0] * np.cos(omega * times)]).T
times_fe1 = np.linspace(0, t_max, 100)
dt1 = times_fe1[1] - times_fe1[0]
sol_fe1 = forward_euler_generic(harmonic_rhs, u0, times_fe1, omega)
times_fe2 = np.linspace(0, t_max, 400)
dt2 = times_fe2[1] - times_fe2[0]
sol_fe2 = forward_euler_generic(harmonic_rhs, u0, times_fe2, omega)
times_ssprk2 = np.linspace(0, t_max, 80)
dt3 = times_ssprk2[1] - times_ssprk2[0]
sol_ssprk2 = ssprk22_generic(harmonic_rhs, u0, times_ssprk2, omega)
times_rk2 = np.linspace(0, t_max, 80)
dt4 = times_rk2[1] - times_rk2[0]
sol_rk2 = rk2_generic(harmonic_rhs, u0, times_rk2, omega)
ax.plot(sol_fe1[:,0], sol_fe1[:,1],
color='tab:red', label=f"F.E. dt = {dt1:.3f}")
ax.plot(sol_fe2[:,0], sol_fe2[:,1],
color='tab:purple', label=f"F.E. dt = {dt2:.3f}")
ax.plot(sol_ssprk2[:,0], sol_ssprk2[:,1],
color='tab:blue', label=f"SSPRK2 dt = {dt3:.3f}")
ax.plot(sol_rk2[:,0], sol_rk2[:,1],
color='tab:green', label=f"RK2 dt = {dt4:.3f}")
ax.plot(analytical[:,0], analytical[:,1],
'k-', label=f"Analytical")
ax_e.plot(times, energy(analytical, omega), 'k')
ax_e.plot(times_fe1, energy(sol_fe1, omega), 'tab:red')
ax_e.plot(times_fe2, energy(sol_fe2, omega), 'tab:purple')
ax_e.plot(times_ssprk2, energy(sol_ssprk2, omega), 'tab:blue')
ax_e.plot(times_rk2, energy(sol_rk2, omega), 'tab:green')
ax.legend(fontsize=12)
ax.set_xlabel("x(t)", fontsize=10)
ax.set_ylabel("v(t)", fontsize=10)
ax.tick_params(axis='both', labelsize=8)
print(f"dE F.E. (dt={dt1:.3f}) at t={t_max}: {(energy(analytical, omega)[-1] - energy(sol_fe1, omega)[-1]):.3e}")
print(f"dE F.E. (dt={dt2:.3f}) at t={t_max}: {(energy(analytical, omega)[-1] - energy(sol_fe2, omega)[-1]):.3e}")
print(f"dE SSPRK2 (dt={dt3:.3f}) at t={t_max}: {(energy(analytical, omega)[-1] - energy(sol_ssprk2, omega)[-1]):.3e}")
print(f"dE RK2 (dt={dt4:.3f}) at t={t_max}: {(energy(analytical, omega)[-1] - energy(sol_rk2, omega)[-1]):.3e}")
ax_e.set_yscale("log")
plt.show()