In [2]:
import numpy as np
import matplotlib.pyplot as plt
from scipy import integrate

%matplotlib inline

In [23]:
def Runge_Kutta(f, t, x0):
    n = len(x0)
    N = len(t)
    x = np.zeros((n, N))
    x[:, 0] = x0
    dt = t[1] - t[0]
    for i in range(N - 1):
        k1 = f(x[:, i])
        k2 = f(x[:, i] + dt * k1 / 2)
        k3 = f(x[:, i] + dt * k2 / 2)
        k4 = f(x[:, i] + dt * k3)
        x[:, i + 1] =  x[:, i] + dt * (k1 + 2 *  (k2 + k3) + k4) / 6
    return x

In [24]:
def Lorenz(X, t=0):
    return np.array([10 * (X[1] - X[0]),
                     28 * X[0] - X[1] - X[0] * X[2],
                     X[0] * X[1] - 8 * X[2] / 3])

In [27]:
N = 10 ** 5
t = np.linspace(0, 50, N)
x0 = np.array([1, 0, 0])
#x = Runge_Kutta(Lorenz, t, x0)
%timeit Runge_Kutta(Lorenz, t, x0)

In [ ]:
#plt.plot(x.T)

In [28]:
%timeit Runge_Kutta(Lorenz, t, x0)


1 loop, best of 3: 23.7 s per loop

In [35]:
X, infodict = integrate.odeint(Lorenz, x0, t, full_output=True)
infodict['message']


Out[35]:
'Integration successful.'

In [37]:
plt.plot(X)


Out[37]:
[<matplotlib.lines.Line2D at 0x7fd01ec75518>,
 <matplotlib.lines.Line2D at 0x7fd01ec756d8>,
 <matplotlib.lines.Line2D at 0x7fd01ec758d0>]

In [30]:
def trya():
    def Lorenz(x, y, z):
        return np.array([10 * (y - x),
                         28 * x - y - x * z,
                         x * y - 8 * z / 3])


    N = 10 ** 5
    t = np.linspace(0, 50, N)
    x0 = np.array([1, 0, 0])
    n = len(x0)
    N = len(t)
    x = np.zeros(N)
    y = np.zeros(N)
    z = np.zeros(N)
    #x = x0
    dt = t[1] - t[0]
    x[0] = 1
    for i in range(N - 1):
        k1 = Lorenz(x[i], y[i], z[i])
        k2 = Lorenz(x[i] + dt * k1[0] / 2, y[i] + dt * k1[1] / 2, z[i] + dt * k1[2] / 2)
        k3 = Lorenz(x[i] + dt * k2[0] / 2, y[i] + dt * k2[1] / 2, z[i] + dt * k2[2] / 2)
        k4 = Lorenz(x[i] + dt * k3[0], y[i] + dt * k3[1], z[i] + dt * k3[2])

        x[i + 1] = x[i] + dt * (k1[0] + 2 *  (k2[0] + k3[0]) + k4[0]) / 6
        y[i + 1] = y[i] + dt * (k1[1] + 2 *  (k2[1] + k3[1]) + k4[1]) / 6
        z[i + 1] = z[i] + dt * (k1[2] + 2 *  (k2[2] + k3[2]) + k4[2]) / 6
%timeit trya()


1 loop, best of 3: 11.2 s per loop

In [ ]:
mean_x += new_x / N_iter