How do I optimize the solving of ODEs?

Question

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()

• Predrag Punosevac: 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

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.}