Runge-Kutta implementation in Python + mpmath

rungekutta.py

#!/usr/bin/env python2

import mpmath

@mpmath.extradps(30)
def runge_kutta(f, t0, y0, h, steps):
    t = t0
    y = y0
    for n in range(steps):
        k1 = f(t, y)
        k2 = f(t + h / 2, y + k1 * (h / 2))
        k3 = f(t + h / 2, y + k2 * (h / 2))
        k4 = f(t + h, y + k3 * h)

        t += h
        y += (h / 6) * (k1 + 2 * k2 + 2 * k3 + k4)

    return t, y

def runge_kutta_system(f, g, t0, y0, z0, h, steps):
    t = t0
    y = y0
    z = z0
    for n in range(steps):
        K1 = f(t, y, z)
        L1 = g(t, y, z)

        K2 = f(t + h / 2, y + h * K1 / 2, z + h * L1 / 2)
        L2 = g(t + h / 2, y + h * K1 / 2, z + h * L1 / 2)

        K3 = f(t + h / 2, y + h * K2 / 2, z + h * L2 / 2)
        L3 = g(t + h / 2, y + h * K2 / 2, z + h * L2 / 2)

        K4 = f(t + h, y + h * K3, z + h * L3)
        L4 = g(t + h, y + h * K3, z + h * L3)

        t += h
        y += h * (K1 + 2 * K2 + 2 * K3 + K4) / 6
        z += h * (L1 + 2 * L2 + 2 * L3 + L4) / 6

    return t, y, z

if __name__ == '__main__':
    I_i = mpmath.mpf('0.05')
    A = mpmath.mpf('0.2')
    B = mpmath.mpf('0.05')
    GAMMA = mpmath.mpf('0.4')
    EPSILON = mpmath.mpf('0.0005')
    I = lambda et: I_i + et
    f = lambda t, v, w: -v * (v - A) * (v - 1) - w + I(t)
    g = lambda t, v, w: B * (v - GAMMA * w)

    t0 = mpmath.mpf('0')
    v0 = mpmath.mpf('0.0375')
    w0 = mpmath.mpf('0')
    t = []
    v = []

    for i in range(5000):
        t.append(t0)
        v.append(v0)
        t0, v0, w0 = runge_kutta_system(f, g, t0, v0, w0,
                                        mpmath.mpf('0.01'), 1)

    from matplotlib import pyplot
    pyplot.plot(t, v)
    # pyplot.xlim(0, 100)
    # pyplot.ylim(0, 10)
    pyplot.show()