rungekutta.py

from math import sin, cos, tau
import collections

class frange(collections.namedtuple('_', 'start stop step')):
    def __iter__(self):
        t, end, step = super().__iter__()
        while t <= end:
            yield t
            t += step

def RK(z, F, I):
    ɛ, RKs = I.step, []
    for t in I:
        RKs.append(z.real)
        z1 = F(t, z)
        z2 = F(t + ɛ/2, z + z1*ɛ/2)
        z3 = F(t + ɛ/2, z + z2*ɛ/2)
        z4 = F(t + ɛ, z + z3*ɛ)
        z += (z1 + 2*z2 + 2*z3 + z4) / 6 * ɛ
    def closure(t):
        if I.start <= t <= I.stop:
            i, tt = divmod((t - I.start) / ɛ, 1)
            vl, vu = RKs[int(i) : int(i)+2]
            return vl*tt + vu*(1-tt)
    return closure

diffeq = lambda t, z: z / 1j, frange(0, tau, 1e-5)
RKsin, RKcos = RK(1j, *diffeq), RK(1, *diffeq)
for t in frange(0, tau, 1e-2):
    s, c = RKsin(t), RKcos(t)
    print(*(format(x, '20.17f') for x in [t, s, c, s-sin(t), c-cos(t)]))