diff --git a/numeric.py b/numeric.py index 8dfdd03..55b9135 100644 --- a/numeric.py +++ b/numeric.py @@ -3,7 +3,7 @@ import matplotlib.pyplot as plt threshold = 100 -def RK_n(y0, t: np.ndarray, n=4): +def RK4(y0, t): ODE = lambda t, y: t**3 + y**3 #note: t is a vector of time points, y0 is the initial value steps = len(t) @@ -16,17 +16,13 @@ def RK_n(y0, t: np.ndarray, n=4): t0_idx = np.where(t >= 0)[0][0] y[t0_idx] = y0 # set the initial condition at t[0] - #forward RK for i in range(t0_idx, steps - 1): - #n-th order Runge-Kutta method - k = np.zeros(n) - k[0] = h * ODE(t[i], y[i]) - for j in range(1, n): - t_j = t[i] + h * (j / n) - y_j = y[i] + k[j - 1] / n - k[j] = h * ODE(t_j, y_j) - y[i + 1] = y[i] + np.sum(k) / n - return y + k1 = h * ODE(t[i], y[i]) + k2 = h * ODE(t[i] + h / 2, y[i] + k1 / 2) + k3 = h * ODE(t[i] + h / 2, y[i] + k2 / 2) + k4 = h * ODE(t[i] + h, y[i] + k3) + y[i + 1] = y[i] + (k1 + 2 * k2 + 2 * k3 + k4) / 6 + return y def euler(f, a, b, y0, h): xs = [a] @@ -84,7 +80,7 @@ for idx, (y0, xm, xM, ym, yM) in enumerate(lims): Te, Ye = euler(ODE, 0, xM, y0, h) Ti, Yi = improvedEuler(ODE, 0, xM, y0, h) Tr = np.linspace(0, xM, int((xM - 0)/h)) - Yr = RK_n(y0, Tr) + Yr = RK4(y0, Tr) axes[idx][0].plot(Te, Ye, label=f'$\Delta t$={h}', linewidth=1) axes[idx][1].plot(Ti, Yi, label=f'$\Delta t$={h}', linewidth=1)