dendop_test_v00.py

def D5a(x, y):
    A = -(y[0]**2 + y[1]**2)**-1.5
    return [y[2],  y[3],  A*y[0],  A*y[1]] # dimensionless orbit equation


# as described in:
# https://github.com/jddmartin/scipy/tree/dense_output_from_dopri5_and_dop853
# https://github.com/jddmartin/dense_output_example_usage

def dendop(theta, con):  # based on dense_dop

    try:
        s0, s1 = con.shape
    except:
        raise RuntimeError("con shape does not make sense")

    theta1 = 1.0 - theta

    if s0 == 5:
        theta_ordering = (theta1, theta)   # dopri5
    elif s0 == 8:
        theta_ordering = (theta, theta1)   # dop853
    else:
        raise RuntimeError("con.shape[0] is not 5 or 8")

    dense_output = con[-1]

    for i in range(s0-1):
        dense_output = (con[-2-i] +
                        theta_ordering[i%2] * dense_output)

    return dense_output

def solout_save(nr, xold, x, y, con):

    global X, Y, cons

    X.append(x)
    Y.append(y)
    cons.append(con)



from scipy.integrate import ode as ODE
import numpy as np
import matplotlib.pyplot as plt

ep5 = 0.9
y0_D5 = [1.0-ep5, 0.0, 0.0, np.sqrt((1.0+ep5)/(1.0-ep5))]

r = ODE(D5a).set_integrator('dop853', rtol=1E-12, atol=1E-12)
t0 = 0.0
r.set_initial_value(y0_D5, t0)


# Integrate to an array of test points (will interpolate to these points next)
x0 = np.linspace(0, 20, 201)
ao = [y0_D5]
to = [t0]
for x in x0[1:]:
    ao.append(r.integrate(x))
    to.append(r.t)
aon = np.array(ao)
ton = np.array(to)
print "last point: ", aon[-1]


# now just integrate directly to t=20.0 and save the interpolation coeficients
X, Y, cons = [], [], []
r = ODE(D5a).set_integrator('dop853', rtol=1E-12, atol=1E-12)
r.set_solout(solout_save, dense_components=(0, 1, 2, 3))
t0 = 0.0
r.set_initial_value(y0_D5, t0)

a = r.integrate(x)
print "last point: ", a

Xn, Yn, consn = np.array(X), np.array(Y), np.array(cons)
print "Yn[-1]: ", Yn[-1]


plt.figure()
x, y, vx, vy = Yn.T  # free integration results
plt.plot(Xn, x, '-b')
plt.plot(Xn, y, '-g')
x, y, vx, vy = aon.T   # test points
plt.plot(ton, x, 'or')
plt.plot(ton, y, 'oc')
plt.show()


tons = ton[1:-1]
ivals = [np.argmax(X>t) for t in tons]  # find locations to perform interpolate

thetas = [(t-X[i-1])/(X[i]-X[i-1]) for i, t in zip(ivals, tons)]

consnuse = consn[ivals]


interpolated = np.array([dendop(theta, con) for theta, con in
                         zip(thetas, consnuse)])
calculated = aon[1:-1]
times = ton[1:-1]

differences = interpolated - calculated


plt.figure()
for thing in differences.T:
    plt.plot(times, thing)
plt.title("interpolated - calculated")
plt.show()