def plot_energy(model,
initial_condition,
energy_function,
index=0,
integrators=[
ode.euler_1, ode.euler_2, ode.runge_kutta_4, verlet,
euler_1_modified_v, euler_1_modified_x
],
**kwargs):
""" Compares solutions for different ODE methods in one graph.
Arguments:
model -- a function returning the right-hand side of the ODE
energy_function -- function that calculates energy
index -- for a set of ODE the index of the function to be plotted
kwargs -- addtitional parameters for ode_solve function (step, times, etc)
"""
for integrator in integrators:
ys, ts = ode.ode_solve(model,
initial_condition,
integrator=integrator,
**kwargs)
es = [energy_function(*xv) for xv in ys]
plt.plot(ts, es, label=integrator.__name__)
# Axes labels
plt.xlabel("t")
plt.ylabel("E")
plt.legend()
plt.show()
def plot_compare_steps(model,
initial_condition,
integrator,
index=0,
analytical_solution=None,
dts=(0.5, 0.2, 0.1),
**kwargs):
""" Compares solutions for different ODE methods in one graph.
Arguments:
model -- a function returning the right-hand side of the ODE
analytical_solution -- exact solution of the ODE (if available)
index -- for a set of ODE the index of the function to be plotted
kwargs -- addtitional parameters for ode_solve function (step, times, etc)
"""
show_error = analytical_solution is not None
figure, axes = plt.subplots(2 if show_error else 1, 1)
main_panel = axes[0] if show_error else axes
error_panel = axes[1] if show_error else None
ts = None
for dt in dts:
ys, ts = ode.ode_solve(model,
initial_condition,
integrator=integrator,
dt=dt,
**kwargs)
if ys.ndim != 1:
ys = ys[:, index]
main_panel.plot(ts, ys, label=f"$dt={dt}$")
if show_error:
error_panel.plot(ts, local_error(ys, ts, analytical_solution))
if analytical_solution is not None:
ts = np.linspace(min(ts), max(ts), 200)
main_panel.plot(ts,
analytical_solution(ts),
label=analytical_solution.__name__)
# Axes labels
main_panel.set_xlabel("t")
main_panel.set_ylabel("y")
main_panel.legend()
if show_error:
error_panel.set_xlabel("t")
error_panel.set_ylabel(r'$\epsilon$')
figure.suptitle(integrator.__name__)
plt.show()
def plot_cummulative_error(
model,
initial_condition,
analytical_solution,
index=0,
integrators=[ode.euler_1, ode.euler_2, ode.runge_kutta_4],
dts=np.linspace(0.002, 0.1, 100),
**kwargs):
""" Compares solutions for different ODE methods in one graph.
Arguments:
model -- a function returning the right-hand side of the ODE
analytical_solution -- exact solution of the ODE (if available)
index -- for a set of ODE the index of the function to be plotted
kwargs -- addtitional parameters for ode_solve function (step, times, etc)
"""
plt.figure()
for integrator in integrators:
errors = []
for dt in dts:
ys, ts = ode.ode_solve(model,
initial_condition,
integrator=integrator,
dt=dt,
**kwargs)
if ys.ndim != 1:
ys = ys[:, index]
errors.append(
average_cummulative_error(ys, ts, analytical_solution))
slope, intercept, r_value, p_value, std_err = linregress(
np.log(dts), np.log(errors))
plt.loglog(dts,
errors,
label=f"{integrator.__name__} $\\alpha$={slope:.3}")
plt.xlabel("$\Delta$t")
plt.ylabel("$\mathcal{E}$")
plt.legend()
plt.show()
import numpy as np
import matplotlib.pyplot as plt
import ode
import graphs
def model(y, t):
return np.sin(y * t)
ys, ts = ode.ode_solve(model, 1, dt=1, integrator=ode.euler_1)
plt.plot(ts, ys)
plt.xlabel("t")
plt.ylabel("y")
plt.title(r"$\frac{dy}{dt}=sin(yt)$")
plt.show()
rho = 28
beta = 8 / 3
def lorenz(y, t):
""" Lorenz model of atmospheric convection """
x, ys, z = y
dx = sigma * (ys - x)
dy = x * (rho - z) - ys
dz = x * ys - beta * z
return np.array([dx, dy, dz])
ys, ts = ode.ode_solve(lorenz, (1, 1, 1), dt=0.01, maxt=100)
plt.figure()
plt.plot(ys[:, 0], ys[:, 2])
plt.xlabel("x")
plt.ylabel("z")
plt.title("Lorenz attractor")
plt.show()
# Thanks Samuel J. for this part of the code (plots a 3D interactive graph)
try:
import plotly.express as px
fig = px.line_3d(x=ys[:, 0], y=ys[:, 1], z=ys[:, 2])
fig.show()
except: