# How do I optimize the solving of ODEs?

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.plot(x, y, z, color = 'r')
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])
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()

-
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

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.

-