Take the 2-minute tour ×
Code Review Stack Exchange is a question and answer site for peer programmer code reviews. It's 100% free, no registration required.

I have been trying to figure out a way to optimize the solving of ODEs in Python but haven't been able to achieve this goal. I tried getting help via a bounty on SO using Cython but nothing came of that.

The code below isn't too slow but it is an example of code I tend to run a lot. Some of the code I use can take forever since I am simulating space flight trajectories. If I can get some help with this code, I can work on the rest myself with ideas and knowledge I gain.

The code isn't short but not too long. The crux of the matter is the performance of the ODE solving. The rest is solving constants that go into the ODE or IC.

import numpy as np
from scipy.integrate import ode
import pylab
from mpl_toolkits.mplot3d import Axes3D
from scipy.optimize import brentq
from scipy.optimize import fsolve

me = 5.974 * 10 ** 24  #  mass of the earth                                         
mm = 7.348 * 10 ** 22  #  mass of the moon                                          
G = 6.67259 * 10 ** -20  #  gravitational parameter                                 
re = 6378.0  #  radius of the earth in km                                           
rm = 1737.0  #  radius of the moon in km                                            
r12 = 384400.0  #  distance between the CoM of the earth and moon                   
rs = 66100.0  #  distance to the moon SOI                                           
Lambda = np.pi / 6  #  angle at arrival to SOI                                      
M = me + mm
d = 300  #  distance the spacecraft is above the Earth                              
pi1 = me / M
pi2 = mm / M
mue = 398600.0  #  gravitational parameter of earth km^3/sec^2                      
mum = G * mm  #  grav param of the moon                                             
mu = mue + mum
omega = np.sqrt(mu / r12 ** 3)
#  distance from the earth to Lambda on the SOI                                     
r1 = np.sqrt(r12 ** 2 + rs ** 2 - 2 * r12 * rs * np.cos(Lambda))
vbo = 10.85  #  velocity at burnout                                                 
h = (re + d) * vbo  #  angular momentum                                             
energy = vbo ** 2 / 2 - mue / (re + d)  #  energy                                   
v1 = np.sqrt(2.0 * (energy + mue / r1))  #  refer to the close up of moon diagram   
#  refer to diagram for angles                                                      
theta1 = np.arccos(h / (r1 * v1))
phi1 = np.arcsin(rs * np.sin(Lambda) / r1)
#                                                                                   
p = h ** 2 / mue  #  semi-latus rectum                                              
a = -mue / (2 * energy)  #  semi-major axis                                         
eccen = np.sqrt(1 - p / a)  #  eccentricity 

nu0 = 0
nu1 = np.arccos((p - r1) / (eccen * r1))


#  Solving for the eccentric anomaly                                                
def f(E0):
    return np.tan(E0 / 2) - np.sqrt((1 - eccen) / (1 + eccen)) * np.tan(0)


E0 = brentq(f, 0, 5)


def g(E1):
    return np.tan(E1 / 2) - np.sqrt((1 - eccen) / (1 + eccen)) * np.tan(nu1 / 2)


E1 = fsolve(g, 0)


#  Time of flight from r0 to SOI                                                    
deltat = (np.sqrt(a ** 3 / mue) * (E1 - eccen * np.sin(E1)
                                   - (E0 - eccen * np.sin(E0))))



#  Solve for the initial phase angle 
def s(phi0):
    return phi0 + deltat * 2 * np.pi / (27.32 * 86400) + phi1 - nu1


phi0 = fsolve(s, 0)


nu = -phi0                                                                   

gamma = 0 * np.pi / 180  #  angle in radians of the flight path                     

vx = vbo * (np.sin(gamma) * np.cos(nu) - np.cos(gamma) * np.sin(nu))
#  velocity of the bo in the x direction                                            
vy = vbo * (np.sin(gamma) * np.sin(nu) + np.cos(gamma) * np.cos(nu))
#  velocity of the bo in the y direction                                            

xrel = (re + 300.0) * np.cos(nu) - pi2 * r12


yrel = (re + 300.0) * np.sin(nu)


u0 = [xrel, yrel, 0, vx, vy, 0]


def deriv(u, dt):
    return [u[3],  #  dotu[0] = u[3]                                                
            u[4],  #  dotu[1] = u[4]                                                
            u[5],  #  dotu[2] = u[5]                                                
            (2 * omega * u[4] + omega ** 2 * u[0] - mue * (u[0] + pi2 * r12) /
             np.sqrt(((u[0] + pi2 * r12) ** 2 + u[1] ** 2) ** 3) - mum *
             (u[0] - pi1 * r12) /
             np.sqrt(((u[0] - pi1 * r12) ** 2 + u[1] ** 2) ** 3)),
            #  dotu[3] = that                                                       
            (-2 * omega * u[3] + omega ** 2 * u[1] - mue * u[1] /
             np.sqrt(((u[0] + pi2 * r12) ** 2 + u[1] ** 2) ** 3) - mum * u[1] /
             np.sqrt(((u[0] - pi1 * r12) ** 2 + u[1] ** 2) ** 3)),
            #  dotu[4] = that                                                       
            0]  #  dotu[5] = 0                                                      


dt = np.linspace(0.0, 259200.0, 259200.0)  #  secs to run the simulation            
u = odeint(deriv, u0, dt)
x, y, z, x2, y2, z2 = u.T


fig = pylab.figure()
ax = fig.add_subplot(111, projection='3d')
ax.plot(x, y, z, color = 'r')
#  adding the moon                                                                  
phi = np.linspace(0, 2 * np.pi, 100)
theta = np.linspace(0, np.pi, 100)
xm = rm * np.outer(np.cos(phi), np.sin(theta)) + r12 - pi2 * r12
ym = rm * np.outer(np.sin(phi), np.sin(theta))
zm = rm * np.outer(np.ones(np.size(phi)), np.cos(theta))
ax.plot_surface(xm, ym, zm, color = '#696969', linewidth = 0)
ax.auto_scale_xyz([-8000, 385000], [-8000, 385000], [-8000, 385000])
#  adding the earth                                                                 
xe = re * np.outer(np.cos(phi), np.sin(theta)) - pi2 * r12
ye = re * np.outer(np.sin(phi), np.sin(theta))
ze = re * np.outer(np.ones(np.size(phi)), np.cos(theta))
ax.plot_surface(xe, ye, ze, color = '#4169E1', linewidth = 0)
ax.auto_scale_xyz([-8000, 385000], [-8000, 385000], [-8000, 385000])

pylab.savefig("test.eps", format = "eps", dpi = 1000)
pylab.show()
share|improve this question
    
Check out PyDSTool written by my good friend Rob Clewley. His code is better optimized than the MATLAB and vetted for almost 4-5 years now. It has been started while he was postdoc at Cornell. [1]ni.gsu.edu/~rclewley/PyDSTool/FrontPage.html –  Predrag Punosevac May 20 '13 at 19:24
add comment

1 Answer 1

I would manually "eliminate common sub-expressions" in the deriv:

def deriv(u, dt):
    norm1 = np.sqrt(((u[0] + pi2 * r12) ** 2 + u[1] ** 2) ** 3)
    norm2 = np.sqrt(((u[0] - pi1 * r12) ** 2 + u[1] ** 2) ** 3)
    return [u[3],  #  dotu[0] = u[3]                                                
        u[4],  #  dotu[1] = u[4]                                                
        u[5],  #  dotu[2] = u[5]                                                
        (2 * omega * u[4] + omega ** 2 * u[0] - mue * (u[0] + pi2 * r12) /
         norm1 - mum *
         (u[0] - pi1 * r12) /
         norm2),
        #  dotu[3] = that                                                       
        (-2 * omega * u[3] + omega ** 2 * u[1] - mue * u[1] /
         norm1 - mum * u[1] /
         norm2),
        #  dotu[4] = that                                                       
        0]  #  dotu[5] = 0                                                      

And BTW is the tan(0) in f(E0) actually tan(nu0 / 2)?

I am not a Python programmer.

share|improve this answer
add comment

Your Answer

 
discard

By posting your answer, you agree to the privacy policy and terms of service.

Not the answer you're looking for? Browse other questions tagged or ask your own question.