Let's try a ODE solver, the Runge-Kutta algorithm:
import matplotlib.pyplot as plt
import numba
import numpy as np
Let's setup an ODE function to solve. We can write our ODE as a system of linear first order ODE equations:
The harmonic motion equation can be written in terms of $\mathbf{f}(t, \mathbf{y}) = \dot{\mathbf{y}}$, where this is in the standard form:
$$ \mathbf{y} = \left( \begin{matrix} \dot{x} \\ x \end{matrix} \right) $$$$ \mathbf{f}(t, \mathbf{y}) = \dot{\mathbf{y}} = \left( \begin{matrix} \ddot{x} \\ \dot{x} \end{matrix} \right) = \left( \begin{matrix} -\frac{k}{m} x \\ \dot{x} \end{matrix} \right) = \left( \begin{matrix} -\frac{k}{m} y_1 \\ y_0 \end{matrix} \right) $$x_max = 1 # Size of x max
v_0 = 0
koverm = 1 # k / m
def f(_t, y):
"Y has two elements, x and v"
return np.array([-koverm * y[1], y[0]])
Now let's do the most basic, stupid thing possible: The Euler algorithm:
def euler_ivp(f, init_y, t):
steps = len(t)
order = len(init_y) # Number of equations
y = np.empty((steps, order))
y[0] = init_y # Note that this sets the elements of the first row
for n in range(steps - 1):
h = t[n + 1] - t[n]
# Compute dydt based on *current* position
dydt = f(t[n], y[n])
# Compute next velocity and position
y[n + 1] = y[n] - dydt * h
return y
We can see, when compared to the actual solution, that the Euler very quickly is destroyed by errors at each step (in this case, the frequency is not too bad, though):
ts = np.linspace(0, 40, 1000 + 1)
y = euler_ivp(f, [x_max, v_0], ts)
plt.plot(ts, np.cos(ts))
plt.plot(ts, y[:, 0], "--");
Note that $h = t_{n+1} - t_n $.
$$ \dot{y} = f(t,y) \\ \implies y = \int f(t,y) \, dt \\ \implies y_{n+1} = y_{n} + \int_{t_n}^{t_{n+1}} f(t,y) \, dt $$Now, expand $f$ in a Taylor series around the midpoint of the interval:
$$ f(t,y) \approx f(t_{n+\frac{1}{2}},y_{n+\frac{1}{2}}) + \left( t - t_{n+\frac{1}{2}}\right) \dot{f}(t_{n+\frac{1}{2}}) + \mathcal{O}(h^2) $$The second term here is symmetric in the interval, so all we have left is the first term in the integral:
$$ \int_{t_n}^{t_{n+1}} f(t,y) \, dt \approx h\, f(t_{n+\frac{1}{2}},y_{n+\frac{1}{2}}) + \mathcal{O}(h^3) $$Back into the original statement, we get:
$$ y_{n+1} \approx \color{blue}{ y_{n} + h\, f(t_{n+\frac{1}{2}},y_{n+\frac{1}{2}}) } + \mathcal{O}(h^3) \tag{rk2} $$We've got one more problem! How do we calculate $f(t_{n+\frac{1}{2}},y_{n+\frac{1}{2}})$? We can use the Euler's algorithm that we saw last time:
$$ y_{n+\frac{1}{2}} \approx y_n + \frac{1}{2} h \dot{y} = \color{red}{ y_n + \frac{1}{2} h f(t_{n},y_{n}) } $$Putting it together, this is our RK2 algorithm:
$$ \mathbf{y}_{n+1} \approx \color{blue}{ \mathbf{y}_{n} + \mathbf{k}_2 } \tag{1.0} $$$$ \mathbf{k}_1 = h \mathbf{f}(t_n,\, \mathbf{y}_n) \tag{1.1} $$$$ \mathbf{k}_2 = h \mathbf{f}(t_n + \frac{h}{2},\, \color{red}{\mathbf{y}_n + \frac{\mathbf{k}_1}{2}}) \tag{1.2} $$We've picked up bold face to indicate that we can have a vector of ODEs.
def rk2_ivp(f, init_y, t):
steps = len(t)
order = len(init_y)
y = np.empty((steps, order))
y[0] = init_y
for n in range(steps - 1):
h = t[n + 1] - t[n]
k1 = h * f(t[n], y[n]) # 1.1
k2 = h * f(t[n] + h / 2, y[n] + k1 / 2) # 1.2
y[n + 1] = y[n] + k2 # 1.0
return y
Let's try this with the same grid as before:
ts = np.linspace(0, 40, 1000 + 1)
y = rk2_ivp(f, [x_max, v_0], ts)
plt.plot(ts, np.cos(ts))
plt.plot(ts, y[:, 0], "--");
And, on a coarser grid:
ts = np.linspace(0, 40, 100 + 1)
y = rk2_ivp(f, [x_max, v_0], ts)
plt.plot(ts, np.cos(ts))
plt.plot(ts, y[:, 0], "--");
We can get the RK4 algorithm by keeping another non-zero term in the Taylor series:
$$ \mathbf{y}_{n+1} \approx \mathbf{y}_{n} + \frac{1}{6} (\mathbf{k}_1 + 2 \mathbf{k}_2 + 2 \mathbf{k}_3 + \mathbf{k}_4 ) \tag{2.0} $$$$ \mathbf{k}_1 = h \mathbf{f}(t_n,\, \mathbf{y}_n) \tag{2.1} $$$$ \mathbf{k}_2 = h \mathbf{f}(t_n + \frac{h}{2},\, \mathbf{y}_n + \frac{\mathrm{k}_1}{2}) \tag{2.2} $$$$ \mathbf{k}_3 = h \mathbf{f}(t_n + \frac{h}{2},\, \mathbf{y}_n + \frac{\mathrm{k}_2}{2}) \tag{2.3} $$$$ \mathbf{k}_4 = h \mathbf{f}(t_n + h,\, \mathbf{y}_n + \mathrm{k}_3) \tag{2.4} $$def rk4_ivp(f, init_y, t):
steps = len(t)
order = len(init_y)
y = np.empty((steps, order))
y[0] = init_y
for n in range(steps - 1):
h = t[n + 1] - t[n]
k1 = h * f(t[n], y[n]) # 2.1
k2 = h * f(t[n] + h / 2, y[n] + k1 / 2) # 2.2
k3 = h * f(t[n] + h / 2, y[n] + k2 / 2) # 2.3
k4 = h * f(t[n] + h, y[n] + k3) # 2.4
y[n + 1] = y[n] + 1 / 6 * (k1 + 2 * k2 + 2 * k3 + k4) # 2.0
return y
Let's try this with the same course grid as before:
ts = np.linspace(0, 40, 100 + 1)
y = rk4_ivp(f, [x_max, v_0], ts)
plt.plot(ts, np.cos(ts))
plt.plot(ts, y[:, 0], "--");
%%timeit
ts = np.linspace(0, 40, 1000 + 1)
y = rk4_ivp(f, [x_max, v_0], ts)
f_jit = numba.njit(f)
rk4_ivp_jit = numba.njit(rk4_ivp)
%%timeit
ts = np.linspace(0, 40, 1000 + 1)
y = rk4_ivp_jit(f_jit, np.array([x_max, v_0]), ts)
You can inspect the types if you'd like to add them after running once:
f_jit.inspect_types()