def RK45(x, v, t, n, h, F, tolx, varistep=False):

    # I can't guarantee these are all correct, better check them yourself
    a11 = 1./1.  # not used
    a21 = 1./4.
    a31 = 3./8.
    a41 = 12./13.
    a51 = 1./1.
    a61 = 1./2.

    b21 = 1./4.

    b31 = 3./32.
    b32 = 9./32.

    b41 = 1932./2197.
    b42 = -7200./2197.
    b43 = 7296./2197.

    b51 = 439./216.
    b52 = -8./1.
    b53 = 3680./513.
    b54 = -845./4104.

    b61 = -8./27.
    b62 = 2./1.
    b63 = -3544./2565.
    b64 = 1859./4104.
    b65 = -11./40.

    c1 = 25./216.
    # no c2
    c3 = 1408./2565.
    c4 = 2197./4104.
    c5 = -1./5.

    d1 = 16./135.
    #no d2
    d3 = 6656./12825.
    d4 = 28561./56430.
    d5 = -9./50.
    d6 = 2./55.

    sm = np.zeros_like(t)

    for i in range(n):  # written for readability, not speed

        kv1 = h * F( x[i] )
        kx1 = h *  ( v[i] )

        kv2 = h * F( x[i] + b21*kx1 )
        kx2 = h *  ( v[i] + b21*kv1 )

        kv3 = h * F( x[i] + b31*kx1 + b32*kx2 )
        kx3 = h *  ( v[i] + b31*kv1 + b32*kv2 )

        kv4 = h * F( x[i] + b41*kx1 + b42*kx2 + b43*kx3 )
        kx4 = h *  ( v[i] + b41*kv1 + b42*kv2 + b43*kv3 )

        kv5 = h * F( x[i] + b51*kx1 + b52*kx2 + b53*kx3 + b54*kx4 )
        kx5 = h *  ( v[i] + b51*kv1 + b52*kv2 + b53*kv3 + b54*kv4 )

        kv6 = h * F( x[i] + b61*kx1 + b62*kx2 + b63*kx3 + b64*kx4 + b65*kx5 )
        kx6 = h *  ( v[i] + b61*kv1 + b62*kv2 + b63*kv3 + b64*kv4 + b65*kv5 )

        xx  = x[i] + c1*kx1 + c3*kx3 + c4*kx4 + c5*kx5  # no c2
        vv  = v[i] + c1*kv1 + c3*kv3 + c4*kv4 + c5*kv5

        xx5  = x[i] + d1*kx1 + d3*kx3 + d4*kx4 + d5*kx5 + d6*kx6  # no c2
        vv5 =  v[i] + d1*kv1 + d3*kv3 + d4*kv4 + d5*kv5 + d6*kv6

        x[i+1]  = xx
        v[i+1]  = vv

        s = np.sqrt(np.sqrt(tolx * h / (2.*abs(xx5 - xx))))  # faster than **0.25

        smin = s.min()

        sm[i] = smin

        t[i+1] = t[i] + h

        if varistep:
            h *= smin  # usually you do this more carefully.

    return sm


def RK4(x, v, t, n, h, F):

    h_over_two = 0.5 * h
    h_over_six = (1./6.) * h

    for i in range(n):  # written for readability, not speed

        kv1 = F( x[i]                   )
        kx1 =    v[i]

        kv2 = F( x[i] + kx1 * h_over_two )
        kx2 =    v[i] + kv1 * h_over_two

        kv3 = F( x[i] + kx2 * h_over_two )
        kx3 =    v[i] + kv2 * h_over_two

        kv4 = F( x[i] + kx3 * h          )
        kx4 =    v[i] + kv3 * h

        v[i+1] = v[i] + h_over_six * (kv1 + 2.*(kv2 + kv3) + kv4)
        x[i+1] = x[i] + h_over_six * (kx1 + 2.*(kx2 + kx3) + kx4)

        t[i+1] = t[i] + h


def RK4_new(x0, v0, n, h, F):
    """RK4 based orbit integrator n time steps of fixed size h
       using the function F for acceleration"""

    hov2, hov6 = h / 2., h / 6.

    x, v, t = x0, v0, 0.

    results = [(t, x.copy(), v.copy())]
    for i in range(n):  # written for readability, not speed

        kv1 = F(x)
        kx1 = v

        kv2 = F(x + kx1 * hov2)
        kx2 =   v + kv1 * hov2

        kv3 = F(x + kx2 * hov2)
        kx3 =   v + kv2 * hov2

        kv4 = F(x + kx3 * h)
        kx4 =   v + kv3 * h

        v += hov6 * (kv1 + 2. * (kv2 + kv3) + kv4)
        x += hov6 * (kx1 + 2. * (kx2 + kx3) + kx4)
        t += h

        results.append((t, x.copy(), v.copy()))

    t, x, v = [np.array(thing) for thing in zip(*results)]

    return t, x, v



def acc(x):
    """ acceleration due to the sun's gravity (NumPy version) """
    return -Gm * x / (x**2).sum()**1.5


import matplotlib.pyplot as plt
import matplotlib.cm as cm
import numpy as np

Gm = 1.3271244002E+20  # m^3 s^-2               (Wikipedia)
e = 0.01671123         # earth eccentricity     (Wikipedia)
a = 1.49598261E+11     # earth semi-major axis  (Wikipedia)

# r = a(1-e^2)/(1+e*cos(theta)) an equation for an ellipse
r_aphelion   = a * (1. + e)
r_perihelion = a * (1. - e)

# using vis-viva equation v^2 = Gm * (2./r - 1./a) (Wikipedia)
# (this is really just conservation of energy: Kinetic + Potential = const)
v_aphelion = np.sqrt(Gm * (2./(1.+e) - 1.) / a)

t_year = 2. * np.pi * np.sqrt(a**3/Gm)

nstep = 500   # CHANGE THIS TO SEE WHAT HAPPENS!

Dt = t_year / float(nstep)   # time step



X4 = np.zeros((1000,3))
V4 = np.zeros((1000,3))
T4 = np.zeros((1000))

V4[0] = 0.0, 0.5*v_aphelion, 0.0
X4[0] = r_aphelion, 0.0, 0.0
T4[0] = 0.0

X45 = X4.copy()
V45 = V4.copy()
T45 = T4.copy()


RK4( X4,  V4,  T4,  nstep, Dt, acc)


vary = True  # try both
s = RK45(X45, V45, T45, nstep, Dt, acc, 1E-04, varistep=vary)


plt.figure()

plt.subplot(2, 3, 1)
plt.plot(X4[:nstep,0], X4[:nstep,1])
plt.title("RK4 x, y orbit")

plt.subplot(2, 3, 2)
plt.plot(X45[:nstep,0], X45[:nstep,1])
plt.title("RK45 x, y orbit")

if not vary:
    plt.subplot(2, 3, 4)
    plt.plot(X4[:nstep,0] - X45[:nstep,0])
    plt.plot(X4[:nstep,1] - X45[:nstep,1])
    plt.title("RK4-RK45 diff. x, y")

plt.subplot(2, 3, 5)
plt.plot(s[:nstep])
plt.title("RK45 s, vari = " + str(vary))

if vary:
    plt.subplot(2, 3, 6)
    plt.plot(T45[1:nstep]-T45[:nstep-1])
    plt.title("RK45 dT, vari = " + str(vary))

plt.show()