Don't Mess with the Pelican (really, don't!)

I like python.

So when I read about Pelican, and I saw people use it and then blog about how they use it and how great it is, I thought I should try to use it too.

Even with plenty of help, I couldn't get much done. I just don't "get" Pelican. Then I started receiving some gentle comments suggesting that maybe Pelican is no Panacea for People Posting Python.

Now I'm here in google, and life is so much easier. So far everything just seems to work for me without much effort.

To add highlighted code, I found this how-to and right now use the Vim theme in hilite.me, though I really should just learn to use pygments.

1 2 3 4 5

def fib(n): a, b = 0, 1 for i in range(n-1): a, b = b, a + b return a

What about equations? Well, for now I'm going to try ASCII input for MathJax as described in this how-to. Goodies on ASCII input, including a renderer for testing, can be found at asciimath.org and there are some good examples of asciimathjax. 

MathType (used to be called Equation Editor in Word, PowerPoint, etc.) is supposed to be compatible in both directions, but so far I haven't figured out how. 

So for example, just type this p(x)=(e^-mu mu^x)/(x!) inside backquotes (backticks), that other little thing on the tilde '~' key, and I get this (after adding one line of HTML to the template):

`p(x)=(e^-mu mu^x)/(x!)`

If your equations stop rendering properly, you may have accumulated bits of extra html in your post floating around after some combination of copy/pasting and formatting.

Using MathJax seem to be easy, and for me, intuitive. The equations displayed in the browser can be accessed by right clicking (or control clicking).

Just for example, you can set it so the equations will zoom on command:

Or you can capture the codes in various formats:

Orbits II (SciPy odeint)

Last night I was talking to a friend who asked me why I used python for science, and not something standard like FORTRAN or C. I said something like "they took all the most commonly used FORTRAN and C functions in the world, and wrapped them up in python for us" or something dreamy-eyed like that. While NumPy has been crafted with care for python users, a lot of SciPy is really standard, widely recognized packages, all just ready for us to invoke with a few lines of python.

Personally, I found it more satisfying to write and run and debug and use the Runge-Kutta algorithms (RK4, RK45) first, before I went ahead and just used the imensely powerful scipy.inegrate package odeint. There is a lot going on under the hood in odeint, and this is great. You can read this tutorial and the reference documentaiton.

Just for example, with dynamic step sizes (can save orders of magnitude of time) how do you get the position at fixed, predetermined times? With the explicit RK scripts in the previous post, I have to include tests and provisional one-step calculations. However, odeint takes care of all of this for us, either by stopping, backing-up, or interpolating between steps, whichever is numerically prudent. However, I still use explicit Runge-Kutta here and there, especially if SciPy is not available for some reason, or if I'm doing "weird" stuff.

If you have read the first Post (Orbits part 1), then this will look even better.

Anyway, here it is!

And this is the python, pretty nice don't you think?

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18

from scipy.integrate import odeint as ODEint X0 = [r_aphelion, 0.0, 0.0, 0.0, 0.5*v_aphelion, 0.0] def func_ODE(X, t): f = np.zeros(6) # x, y, z, vx, vy, vz f[:3] = X[3:] # dx/dt = v one_over_r_cubed = ((X[:3]**2).sum())**-1.5 f[3:] = - Gm * X[:3] * one_over_r_cubed return f times = np.linspace(0.0, t_year, nstep+1) # stopping points to return data X, output_dict = ODEint(func_ODE, X0, timess, full_output=True) stepsizes = output_dict['hu']

How does it work?  We're solving 6 first-order differential equations based on a state vector:

X = [x, y, z, vx, vy, vz]

`(dx)/(dt) = v_x,  (dy)/(dt) = v_y,  (dz)/(dt) = v_z,`    is    f[:3] = X[3:],  and

`(dv_x)/(dt) = a_x,  (dv_y)/(dt) = a_y,  (dv_z)/(dt) = z_z,`    is    f[3:] = -Gm * X[:3] / r**3.

The algorithm will report back to you the calculated values of the state vector at the points in time you specify in times.  You can get a lot of information about what happened "under the hood" by setting full_output = True and examining the goodies returned in output_dict.

Orbits I (some basics)

"So... that's my man, Isaac Newton"

You can watch the video here (English subtitles available)

Newton (along with others) invented calculus, in part, to understand what’s going on behind Keppler’s laws – e.g. 'why are the orbits of the planets elliptical, with the sun at one focus'? He showed that this is the result of an inverse square behavior of gravity. 

On earth, the strength of gravity seemed about the same everywhere. But, by studying the motion of the planets (including earth) it was possible to see that the force of gravity drops off as the square of the distance from the sun. Here is how this is expressed in both scalar and vector form: 

`F=G(m_1 m_2)/r^2`    (scalar)

`bb F=G(m_1 m_2)/|bb r_12|^2 hat bbr_12`    (vector),   where

`bb r_12=bb r_2-bb r_1`,     and     `hat bbr_12=(bb r_2-bbr_1)/|bb r_2-bbr_1|`.

Sometimes it's easier to code if you just write it this way:

`bbF=-Gm_1m_2(bb r)/r^3`

NOTE: In this post, we're assuming that the sun is fixed in space, the earth's mass is so small that it doesn't affect anything, and there are no other planets (or moons!). We're just getting started on this fishing expedition.

A very simple way to try to calculate orbits (with a computer) would be to just use finite differences.

`(d bb v)/(d t) = bb a = (bb F)/(m_2) = - Gm_1(bb r)/r^3`,        `(d bb x)/(d t) = bb v`

`Delta bb v = Gm(bb r)/(r^3) bb Delta t`,        `Delta bb x = bb v bb Delta t`

`bb x_(i+1) = bb x + bb Delta bb x_i`,        `bb v_(i+1) = bb v + bb Delta bb v_i`

Note that Gm (the product of G and the mass of the sun) is known much more precisely than either G or m alone. Can you guess why? Imagine trying to measure one without the other!

Also, 

`r_(ellipse) = ( a (1 - e^2)) / (1 + e cos(theta))`

`r_(aphelion) = a ( 1 + e);     theta = pi`

`r_(perihelion) = a ( 1 - e);     theta = 0`

and the vis-viva equation (which is really just conservation of energy):

`v^2 = Gm( (2)/(r) - (1)/(a) )`

So here's a simple python script to try: (no NumPy or SciPy yet)

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56

import matplotlib.pyplot as plt import matplotlib.cm as cm import math 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 = math.sqrt(Gm * (2./(1.+e) - 1.) / a) t_year = 2. * math.pi * math.sqrt(a**3/Gm) nstep = 100 # CHANGE THIS TO SEE WHAT HAPPENS! Dt = t_year / float(nstep) # time step V = [(0.0, v_aphelion)] X = [(r_aphelion, 0.0 )] T = [0.0] x, y = X[0] vx, vy = V[0] t = T[0] for i in range(nstep): one_over_r3 = (x**2 + y**2)**-1.5 vx += -Dt * Gm * x * one_over_r3 vy += -Dt * Gm * y * one_over_r3 x += Dt * vx y += Dt * vy t += Dt V.append((vx, vy)) X.append(( x, y)) T.append(t) gap = math.sqrt((X[-1][0] - X[0][0])**2 + (X[-1][1] - X[0][1])**2) v_gap = math.sqrt((V[-1][0] - V[0][0])**2 + (V[-1][1] - V[0][1])**2) print nstep, gap, v_gap plot_it = True save_it = True save_name = 'nicefig' if nstep > 10**5: plot_it = False # don't be silly!

...and the next zillion lines just generate and formatting the plots.

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100

if plot_it: r = [math.sqrt(x**2 + y**2) for x,y in X] x, y = zip(*X) vx, vy = zip(*V) tmin, tmax = min(T), max(T) tcolor = [(tt-tmin)/(tmax-tmin) for tt in T] tc2 = [tt*0.6 + 0.2 for tt in tcolor] # avoid dark red/blue cmap = cm.jet c, cfirst, clast = cmap(tc2), cmap(min(tc2)), cmap(max(tc2)) dgrey = (0.5, 0.5, 0.5) # darker gray lgrey = (0.7, 0.7, 0.7) # lighter gray fsl = 16 # font size for axis labels fs = 14 # font size for tick labels params = ['lines.color', 'xtick.color', 'ytick.color', 'axes.labelcolor'] for param in params: plt.rcParams[param] = lgrey plt.rcParams['axes.edgecolor'] = dgrey # this one darker fig = plt.figure(figsize = (7.7,2.45)) # here is the figure ax1 = fig.add_axes([0.08, 0.15, 0.25, 0.8]) # axes for three plots ax2 = fig.add_axes([0.41, 0.15, 0.15, 0.8]) ax3 = fig.add_axes([0.64, 0.40, 0.30, 0.55]) fig.patch.set_alpha(0.0) # FIRST PLOT line = ax1.scatter(x, y, s=16) # scatter plot of all points line.set_facecolors(c) # go back and change color of each one line.set_edgecolors(c) l = ax1.scatter(x[0], y[0], s=50) # make a big blue one for starting point l.set_facecolors('b') l = ax1.scatter(0.0, 0.0, s=120) # the sun! l.set_facecolors((1.0, 1.0, 0.0)) # is yellow? l.set_edgecolors((1.0, 0.7, 0.0)) # with an orange photosphere? not really! ax1.arrow(x[0], y[0], 0.0, 0.25*a, head_width=0.08*a, head_length=0.16*a, fc='k', ec='k') # hwbox = 1.6E+11 ax1.set_xlim([0.0 - hwbox, 0.0 + hwbox]) ax1.set_ylim([0.0 - hwbox, 0.0 + hwbox]) ax1.set_xlabel('x(m)', fontsize = fsl, color=lgrey) ax1.set_ylabel('y(m)', fontsize = fsl, color=lgrey) ax1.tick_params(axis='both', which='major', labelsize = fs, colors=lgrey) # SECOND PLOT hwbox = 2E+10 ax2.plot(x, y, '-k') l = ax2.scatter(x[0], y[0], s=240) l.set_facecolors(c[0]) line = ax2.scatter(x, y, s=60) line.set_facecolors(c) line.set_edgecolors(c) ax2.arrow(x[0], y[0], 0.0, 0.25*hwbox, head_width=0.05*hwbox, head_length=0.1*hwbox, fc='k', ec='k') # ax2.set_xlim([x[0] - 0.5*hwbox, x[0] + 0.5*hwbox]) ax2.set_ylim([y[0] - hwbox, y[0] + hwbox]) ax2.set_xlabel('x (m)', fontsize = fsl, color=lgrey) ax2.set_ylabel('y (m)', fontsize = fsl, color=lgrey) ax2.tick_params(axis='both', which='major', labelsize = fs, colors=lgrey) # THIRD PLOT ax3.plot(T, r, '-b', linewidth=6) ax3.plot( [min(T), max(T)], [r_aphelion, r_aphelion], '-k') ax3.plot( [min(T), max(T)], [r_perihelion, r_perihelion], '-k') ax3.set_ylabel('r (m)', fontsize = fsl, color=lgrey) ax3.set_xlabel('time (sec)', fontsize = fsl, color=lgrey) ax3.tick_params(axis='both', which='major', labelsize = fs, colors=lgrey) for ax in [ax1, ax2, ax3]: ax.patch.set_facecolor(dgrey) ax.patch.set_alpha(1.0) line_1 = 'nstep = ' + str(nstep) line_2 = 'gap = ' + str(round(gap,1)) + ' m' line_3 = 'v_gap = ' + str(round(v_gap, 6)) + ' m/s' fig.text(0.62, 0.19, line_1, horizontalalignment = 'left', verticalalignment = 'center', fontsize = 16, color = lgrey) fig.text(0.62, 0.12, line_2, horizontalalignment = 'left', verticalalignment = 'center', fontsize = 16, color = lgrey) fig.text(0.62, 0.05, line_3, horizontalalignment = 'left', verticalalignment = 'center', fontsize = 16, color = lgrey) if save_it: plt.savefig(save_name) plt.show()

Just to note: the following values (Wikipedia):

Gm = 1.3271244002E+20 m^3 s^-2

e = 0.01671123

a = 1.49598261E+11

give the period of an orbit of 31558319.5 seconds, or about 365.25833 days. In reality, a year on earth is about 365.25636 days. This difference may be due to our simplifying assumption that the sun just sits in one place and the earth rotates around it. In reality the masses of the sun and earth are roughly 1.989E+30 and 5.972E+24 (Wikipedia) and if they were alone, the'd rotate around a common barycenter. You can see the difference in the year, and the ratio of the masses are both of the order 10E-06. But instead of looking at reduced mass yet, let's just see what happens with this approximation.

Here are some results. Distances are in meters, time in seconds, etc. The total time is the theoretical period (one year) based on the values used for Gm, e, and a, assuming the earth's mass is very small, and no other planets. If the calculation were correct, then the earth should return to exactly the same position after a period of time given by t_year.

Color is used to represent time (from blue to red - indicated by arrow also). The plot on the left is the whole orbit (sun at one focus). The center plot is a zoom in on the starting position. On the right, the radial distance to the sun is plotted against time. The horizontal lines indicate perihelion and aphelion.

In this first simulation (above) where only 50 time steps were used for an entire orbit, the results are obviously extremely bad. The earth doesn’t even return to the same position after one orbit (center plot), and the variation of the distance to the sun (right plot) is right out! The "gap" values shown below the right plot are the scalar distances between initial and final values. Should be zero if we are orbiting correctly.

So let's try increasing the number of steps an see how quickly it improves.

With 200 and 1000 steps, it's not quite dead - it's getting better, and it doesn't want to go on the cart. We know how that story ends - it's all down hill. But if we plot the improvement, it's not really going down hill fast enough, especially if we want to start calculating more complicated things.

Orbits with large eccentricity are much more demanding because the forces are changing quickly. A quick way to make the orbit elliptical is to just cut the initial velocity in half, and let the planet fall towards the sun.

So clearly the orbits aren't going to come even close to returning to the same spot with a reasonable number of steps. Luckily there are better methods that can be extended to generalized multi-body problems.

Scene II:  Enter Runge and Kutta

Newton did find some better ways to solve for the position of an object in an elliptical orbit as a function of time, but we’re going to do some more complicated calculations (multi-body, 3D) so let’s move on a few centuries.

Around 1900, Carl Runge and Martin Wilhelm Kutta developed a system for finding better algorithms for solving ODEs involving time-varying effects. Algorithms in this family are called Runge-Kutta algorithms. You can read about these everywhere. 

The simplest or most common is called RK4, or “the” Runge-Kutta algorithm. Here it is implemented for simultaneous velocity and position expressions. The function a is acceleration (F/m). Be careful to note the x and v subscripts as they are written.

I've left off the terms for time dependence here.

Don't forget, these are all still vectors.

`a()` is a function returning the acceleration vector `(dv)/(dt)`

`k_(v1) = a(x_i)`

`k_(x1) = v_i`

`k_(v2) = a(x_i + (h)/(2) k_(x1))`

`k_(x2) = v_i + (h)/(2) k_(v1)`

`k_(v3) = a (x_i + (h)/(2) k_(x2) )`
`k_(x3) = v_i + (h)/(2) k_(v2)`

`k_(v4) = a ( x_i + h k_(x3) )`
`k_(x4) = v_i + h k_(v3)`

and finally to iterate:

`v_(i+1) = v_i + (h)/(6) ( k_(v1) + 2 (k_(v2) + k_(v3)) + k_(v4) )`
`x_(i+1) = x_i + (h)/(6) ( k_(x1) + 2 (k_(x2) + k_(x3)) + k_(x4) )`

A second, popular algorithm is the Runge-Kutta-Fehlberg or RK45. This method is also fourth order, but more accurate. It uses many more coefficients and floating point operations per step, but with a few more calculations it provides a fifth order solution as well. This can be compared to estimate the approximate size of the error. This comparison can be used to adjust the step size h dynamically, so time is not wasted “over-calculating”. But you have to think about how often and when you change the step size so that it makes sense for your calculation.

`k_(v1) = h a( x_i )`
`k_(x1) = h  v_i `

`k_(v2) = h a( x_i + b_(21) k_(x1) )`
`k_(x2) = h   ( v_i + b_(21) k_(v1) )`

`k_(v3) = h a( x_i + b_(31) k_(x1) + b_(32) k_(x2) )`
`k_(x3) = h   ( v_i + b_(31) k_(v1) + b_(32) k_(v2) )`

`k_(v4) = h a( x_i + b_(41) k_(x1) + b_(42) k_(x2) + b_(43) k_(x3) )`
`k_(x4) = h   ( v_i + b_(41) k_(v1) + b_(42) k_(v2) + b_(43) k_(v3) )`

`k_(v5) = h a( x_i + b_(51) k_(x1) + b_(52) k_(x2) + b_(53) k_(x3) + b_(54) k_(x4) )`
`k_(x5) = h   ( v_i + b_(51) k_(v1) + b_(52) k_(v2) + b_(53) k_(v3) + b_(54) k_(v4) )`

`k_(v6) = h a( x_i + b_(61) k_(x1) + b_(62) k_(x2) + b_(63) k_(x3) + b_(64) k_(x4) + b_(65) k_(x5) )`
`k_(x6) = h   ( v_i + b_(61) k_(v1) + b_(62) k_(v2) + b_(63) k_(v3) + b_(64) k_(v4) + b_(65) k_(v5) )`

and finally to iterate:

`x_(i+1) = x_i + c_1 k_(x1) + c_3 k_(x3) + c_4 k_(x4) + c_5 k_(x5)`      (no  `c_2`)

`v_(i+1) = v_i + c_1 k_(v1) + c_3 k_(v3) + c_4 k_(v4) + c_5 k_(v5)`

These methods are fourth order, so the error (per step) is O(5). Since the number of steps must increase as step size decreases, we can expect the total error to drop as step size to the fourth power.

Yep, for 100 steps, the position error after one orbit is 10^6 meters, or 1000 km, and for 1000 steps it drops to about 30 meters! Remember thought, this is just the numerical precision for this case. The input values have uncertainties, and the assumptions are unrealistic. We're just fishing with python here. Still it's nice to see what a few extra lines can do!

So we can see both methods are the same order of accuracy. Each step is O(5), and the total number of steps increases as the step size decreases, so the final error drops as step-size to the minus four power.  The RK45 method is about a factor of ~30 more accurate in this particular example. You may wonder why you might want to use it. You can get the same accuracy (at least in this example) with RK4 by using `10^(1/4)` or about twice the number of steps, which may be faster than using RK45 with roughly 5 times the number of floats and 50% more calls to the subroutine for acceleration.

With just a few more calculations, the RK45 method can also provide a more accurate fifth order approximation.

`x_(5th) = x_i + d_1 k_(x1) + d_3 k_(x3) + d_4 k_(x4) + d_5 k_(x5) + d_6 k_(x6)`      (no  `d_2`)
`v_(5th) = v_i + d_1 k_(v1) + d_3 k_(v3) + d_4 k_(v4) + d_5 k_(v5) + d_6 k_(v6)`

The best use of this is often to "check" the fourth order solution. If they agree to within a pre-determined tolerance, everything is OK. However if they disagree beyond the tolerance, the size of the disagreement can be used to choose a smaller step size. In the same way, if they agree much more closely than is necessary, a larger step size can be chosen. We'll see later that for some crazy orbits, most of the time you can use a large step size, but when something exciting happens (like a "swing-by") it's good to be able to dynamically (and automatically) reduce the step size during that short event.

`s = ( (delta_x h)/(2 |x_(5th)-x| ) )^(1/4)`

If you want to try it, the RK45 method (Runge-Kutta-Fehlberg) is very popular and you can read about it wherever ODEs are sold.

Here's the half-velocity elliptical orbit test again with both RK4 and RK45 methods. You can see the step size for the RK45 vary by roughly a factor of 50, and so it ends up at a different place. In a serious calculation, you would pay more attention to the algorithm that is actually changing step size. There are different scenarios out there to choose from.

Here is what they look like.

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203

def numburz(): """ possibly the coeficients for the RK45 (Runge-Kutta-Fehlberg) method, but you really better look these up for yourself!! """ global a11, a21, a31, a41, a51, a61 global b21 global b31, b32 global b41, b42, b43 global b51, b52, b53, b54 global b61, b62, b63, b64, b65 global c1, c3, c4, c5 # no c2 global d1, d3, d4, d5, d6 # no d2 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. def RK45(x, v, t, n, h, F, tolx, varistep=False): 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 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 numburz() 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()