# dendop_test_v00.py

a guest
Nov 23rd, 2015
130
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
1. def D5a(x, y):
2.     A = -(y[0]**2 + y[1]**2)**-1.5
3.     return [y[2],  y[3],  A*y[0],  A*y[1]] # dimensionless orbit equation
4.
5.
6. # as described in:
7. # https://github.com/jddmartin/scipy/tree/dense_output_from_dopri5_and_dop853
8. # https://github.com/jddmartin/dense_output_example_usage
9.
10. def dendop(theta, con):  # based on dense_dop
11.
12.     try:
13.         s0, s1 = con.shape
14.     except:
15.         raise RuntimeError("con shape does not make sense")
16.
17.     theta1 = 1.0 - theta
18.
19.     if s0 == 5:
20.         theta_ordering = (theta1, theta)   # dopri5
21.     elif s0 == 8:
22.         theta_ordering = (theta, theta1)   # dop853
23.     else:
24.         raise RuntimeError("con.shape[0] is not 5 or 8")
25.
26.     dense_output = con[-1]
27.
28.     for i in range(s0-1):
29.         dense_output = (con[-2-i] +
30.                         theta_ordering[i%2] * dense_output)
31.
32.     return dense_output
33.
34. def solout_save(nr, xold, x, y, con):
35.
36.     global X, Y, cons
37.
38.     X.append(x)
39.     Y.append(y)
40.     cons.append(con)
41.
42.
43.
44. from scipy.integrate import ode as ODE
45. import numpy as np
46. import matplotlib.pyplot as plt
47.
48. ep5 = 0.9
49. y0_D5 = [1.0-ep5, 0.0, 0.0, np.sqrt((1.0+ep5)/(1.0-ep5))]
50.
51. r = ODE(D5a).set_integrator('dop853', rtol=1E-12, atol=1E-12)
52. t0 = 0.0
53. r.set_initial_value(y0_D5, t0)
54.
55.
56. # Integrate to an array of test points (will interpolate to these points next)
57. x0 = np.linspace(0, 20, 201)
58. ao = [y0_D5]
59. to = [t0]
60. for x in x0[1:]:
61.     ao.append(r.integrate(x))
62.     to.append(r.t)
63. aon = np.array(ao)
64. ton = np.array(to)
65. print "last point: ", aon[-1]
66.
67.
68. # now just integrate directly to t=20.0 and save the interpolation coeficients
69. X, Y, cons = [], [], []
70. r = ODE(D5a).set_integrator('dop853', rtol=1E-12, atol=1E-12)
71. r.set_solout(solout_save, dense_components=(0, 1, 2, 3))
72. t0 = 0.0
73. r.set_initial_value(y0_D5, t0)
74.
75. a = r.integrate(x)
76. print "last point: ", a
77.
78. Xn, Yn, consn = np.array(X), np.array(Y), np.array(cons)
79. print "Yn[-1]: ", Yn[-1]
80.
81.
82. plt.figure()
83. x, y, vx, vy = Yn.T  # free integration results
84. plt.plot(Xn, x, '-b')
85. plt.plot(Xn, y, '-g')
86. x, y, vx, vy = aon.T   # test points
87. plt.plot(ton, x, 'or')
88. plt.plot(ton, y, 'oc')
89. plt.show()
90.
91.
92. tons = ton[1:-1]
93. ivals = [np.argmax(X>t) for t in tons]  # find locations to perform interpolate
94.
95. thetas = [(t-X[i-1])/(X[i]-X[i-1]) for i, t in zip(ivals, tons)]
96.
97. consnuse = consn[ivals]
98.
99.
100. interpolated = np.array([dendop(theta, con) for theta, con in
101.                          zip(thetas, consnuse)])
102. calculated = aon[1:-1]
103. times = ton[1:-1]
104.
105. differences = interpolated - calculated
106.
107.
108. plt.figure()
109. for thing in differences.T:
110.     plt.plot(times, thing)
111. plt.title("interpolated - calculated")
112. plt.show()