Gravity train time calculation

gravitytrain.py

#!/usr/bin/env python3
"""
Calculate travel time for an object falling through a tunnel through the
earth's core, only driven by gravitation
"""
from scipy.integrate import odeint, quad
from scipy.interpolate import interp1d
from numpy import arange, linspace, where, sign, diff
import math
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import host_subplot
import mpl_toolkits.axisartist as AA
# Constant defined in PREM-Article
# For simplification use it as maximum radius
a = 6368000
G = 6.67*10**-11
def rho(r):
    """The density of the earth at radius r in meters
    Taken from Preliminary Reference Earth Model
    http://scholar.google.com/scholar?cluster=15294383875138495935
    """
    r = math.fabs(r)
    x = r/a
    r = r/1000
    rh = 0
    if r >= 6368: rh = 1.02
    elif r >= 6356: rh = 2.6
    elif r >= 6346.6: rh = 2.9
    elif r >= 6151: rh = 2.691 + 0.6924*x
    elif r >= 5971: rh = 7.1089 - 3.8045*x
    elif r >= 5771: rh = 11.2494 - 8.0298*x
    elif r >= 5701: rh = 5.3197 - 1.4836*x
    elif r >= 3480: rh = 7.9565 - 6.4761*x + 5.5283 * x**2 - 3.0807 * x**3
    elif r >= 1221.5: rh = 12.5815 - 1.2638*x - 3.6426 * x**2 - 5.5281 * x**3
    else: rh = 13.0855 - 8.8381 * x**2
    return rh*1000
def dmass(r):
    'mass differntial of the infinitesimal shell at radius r'
    return 4 * math.pi * r**2 * rho(r)
def dF(r, position):
    """force differential produced by the infinitesimal shell
    at radius r when at position
    """
    if math.fabs(r)<100: return 0
    dF = G*dmass(r)/(position**2)
    return dF
def g(position):
    'gravitational acceleration when at position'
    return quad(lambda r: dF(r, position), 0, position, limit=200)[0]
# Interpolate g to reduce computation time
pos_span = linspace(-a-1, a+1, num=1000)
g_inter = interp1d(pos_span, list(map(g, pos_span)))
# g_span = g_inter(pos_span)
def f(x, t):
    """
    vector for the ODE
    x := [ pos, d pos / dt ]
    """
    return [x[1], -g_inter(x[0])]
x0 = [a, 0]
t = arange(0, 80*60, .5)
x = odeint(f, x0, t)
pos = x.transpose()[0]
vel = x.transpose()[1]
traveltime = t[where(diff(sign(diff(pos))))][0]/60
formatstring = """
A gravitational powered cart through earth\'s core takes {0:.2f} minutes
to reach the other side of the earth.
"""
print(formatstring.format(traveltime))

host = host_subplot(111, axes_class=AA.Axes)
plt.subplots_adjust(right=.75)

par = host.twinx()
par2 = host.twinx()

offset = 80
new_fixed_axis = par2.get_grid_helper().new_fixed_axis
par2.axis['right'] = new_fixed_axis(loc='right',
                                    axes=par2,
                                    offset=(offset, 0))

par2.axis['right'].toggle(all=True)

host.set_xlim(0,80)
host.set_ylim(-8000,8000)

host.set_xlabel('time in minutes')
host.set_ylabel('position in kilometers')
par.set_ylabel('velocity in kilometers per second')
par2.set_ylabel('acceleration in meter per second sqared')

par.set_ylim(-10,10)
par2.set_ylim(-15,15)

p1, = host.plot(t/60, pos/1000, label='position')
p2, = par.plot(t/60, vel/1000, label='velocity')
p3, = par2.plot(t[:-1]/60, 2*diff(vel), label='acceleration')

host.legend(loc=4)

host.axis["left"].label.set_color(p1.get_color())
par.axis["right"].label.set_color(p2.get_color())
par2.axis["right"].label.set_color(p3.get_color())

plt.grid()
plt.draw()

plt.show()