In [ ]:
%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
import math
The Plummer potential for mass $M_p$ and core radius $r_c$ is given by $$ \Phi = - { M_p \over { {(r_c^2 + r^2)}^{1/2} } } \tag{1a} $$ and is also used to described softened gravity of a point mass (think of the case $r_c = 0$) for N-body calculations.
The force is the gradient of the potential $$ f = -\nabla \Phi \tag{1b} $$
We also want to record the total energy (kinetic and potential): $$ E = { 1\over 2} v^2 + \Phi \tag{1c} $$ and angular momentum $$ J = r \times v \tag{1d} $$ although we will only be using the Z component of this vector since we are computing orbits restricted to the Z plane: $$ J_z = x \times v_y - y \times v_x \tag{1e} $$
In [ ]:
def radius(x):
"""length of a vector"""
if len(x.shape) == 1:
return math.sqrt(np.inner(x,x))
# elif len(x.shape) == 2:
def potential(pos):
"""potential, defined as a negative number"""
r = radius(pos)
y1 = 1+r*r
return -1.0/math.sqrt(y1)
def angmomz(pos,vel):
"""Angular momentum in Z"""
return pos[0]*vel[1] - pos[1]*vel[0]
def energy(pos,vel):
"""Kinetic and Potential energy"""
return 0.5*np.inner(vel,vel) + potential(pos)
def force(pos):
"""force/acceleration (in our units mass=1 scale-length=1)"""
# note we might be able to use sympy
r = radius(pos)
y2 = 1.0/math.sqrt(1+r*r)
return -pos*y2*y2*y2
Here we write a step function to solve $$ { d\boldsymbol{x} \over dt } = \boldsymbol{v} \tag{2a} $$ and $$ { d\boldsymbol{v} \over dt } = f(\boldsymbol{x}) \tag{2b} $$ in the discretization we write this as (in a simple first order Euler algorithm)
xnew = xold + dt * vold
vnew = vold + dt * fold
In [ ]:
def step0(pos,vel, dt):
"""step0: simple first order Euler"""
old = pos
pos = pos + dt*vel
vel = vel + dt*force(old)
return (pos,vel)
def step1(pos,vel, dt):
"""step1: simple first order Euler - updating position first"""
pos = pos + dt*vel
vel = vel + dt*force(pos)
return (pos,vel)
def step2(pos,vel, dt):
"""step2: simple first order Euler - updating velocity first"""
vel = vel + dt*force(pos)
pos = pos + dt*vel
return (pos,vel)
def step4(pos,vel,dt):
"""step4: Runge Kutta 4 """
# not implemented yet
return None
In [ ]:
def show_stats(data):
"""Show some stats of a numpy array"""
m = data.mean()
s = data.std()
dmin = data.min()
dmax = data.max()
rmin = (dmin-m)/s
rmax = (dmax-m)/s
print("Mean/Std:",m,s)
print("Min/Max:",dmin,dmax)
print("Rmin/Rmax:",rmin,rmax)
print("Goodness: ",s/m)
In [ ]:
x0 = 1.0 # initial X coordinate
v0 = 0.1 # initial Y launch velocity (0.5946 would be a circular orbit)
n = 200 # number of steps to take
dt = 0.1 # integration time step
step = step1 # pick an integration method
print(step.__doc__)
# Derived variables for the remainder
t = 0.0 # always start at t=0
pos = np.array([x0, 0.0, 0.0]) # keeps the current pos
vel = np.array([0.0, v0, 0.0]) # and vel
e = energy(pos,vel)
j = angmomz(pos,vel)
time = np.zeros(1) # time array (we'll append to this)
time[0] = t
phase = np.concatenate(([t,e,j],pos,vel)).reshape(1,9) # watch this peculiar
print("e0 =",e)
print("phase = ",phase)
In [ ]:
# at x0=1.0 this should be the correct speed for a circular orbit
print("v0_circular=",1/math.pow(2.0,0.75))
In [ ]:
%%time
for i in range(n):
(pos,vel) = step(pos,vel,dt)
t = t + dt
e = energy(pos,vel)
j = angmomz(pos,vel)
#print(i,pos,vel)
p = np.concatenate(([t,e,j],pos,vel)).reshape(1,9)
phase = np.concatenate((phase, p),axis=0)
time = np.append(time,t)
#print(phase)
In [ ]:
plt.scatter(phase[:,3],phase[:,4],c=time)
plt.axis('equal')
plt.title("Orbit")
In [ ]:
x = phase[:,3]
y = phase[:,4]
rad = np.sqrt(x*x+y*y)-1
plt.scatter(phase[:,0],rad)
In [ ]:
plt.scatter(phase[:,0], phase[:,1])
plt.title("Conserving Energy?")
show_stats(phase[:,1])
In [ ]:
plt.scatter(phase[:,0], phase[:,2])
plt.title("Conserving Angular Momentum?")
show_stats(phase[:,2])
There are many good and less ideal ways to save data. In astronomy standard formats such has FITS and HDF5 are common. For our work here we use a simple and fast native python method, called pickle. You can save whole objects, and reading them back in will ensure the whole object structure and hierarchy is preserved.
In [ ]:
try:
import cPickle as pickle
print("using cPickle")
except:
import pickle
print("using pickle")
# write it
pickle.dump(phase,open("orbit1.p","wb"))
# read it again
phase2 = pickle.load(open("orbit1.p","rb"))
print(phase[0])
print(phase2[0])
For many scientific applications there are canned routines made available by the community. The scipy package is one such module. We will derive the same orbit integration using scipy.odeint
See e.g. https://docs.scipy.org/doc/scipy-0.18.1/reference/generated/scipy.integrate.odeint.html#scipy.integrate.odeint or better https://docs.scipy.org/doc/scipy-0.18.1/reference/tutorial/integrate.html
However, this function uses the usual ODE notation (cf. (2a) and (2b)) $$ { d\boldsymbol{y} \over dt } = f(\boldsymbol{y},t) \tag{3} $$
In [ ]:
from scipy.integrate import odeint
In [ ]:
def ofunc(y,t):
""" function to integrate
Note we are re-using the force() function from the first part of this notebook
"""
pos = y[0:3]
vel = y[3:]
return np.concatenate((vel,force(pos)))
In [ ]:
n=200
phase0 = np.array([x0,0,0, 0,v0,0]) # initial conditions
times = np.arange(0.0,(n+1)*dt,dt) # requested times where we want a solution
# times = np.linspace(0.0,n*dt,n+1)
print(ofunc(phase0,0.0))
In [ ]:
%%time
orbit = odeint(ofunc, phase0, times)
In [ ]:
plt.scatter(orbit[:,0],orbit[:,1],c=times)
plt.axis('equal')
#plt.scatter(phase[:,3],phase[:,4])
plt.title("Orbit")
In [ ]:
# plot the old one again
plt.scatter(phase[:,3],phase[:,4],c=time)
plt.axis('equal')
plt.title("Orbit")
In [ ]:
# compare the last
p1 = phase[-1,3:]
p2 = orbit[-1,:]
#
print(phase[-1,0],p1)
print(times[-1],p2)
print(0.0,p1-p2)
In [ ]:
# e = energy(pos,vel)
# j = angmomz(pos,vel)
et = np.zeros(len(times))
jt = np.zeros(len(times))
for i in range(len(times)):
pos=orbit[i,:3]
vel=orbit[i,3:]
et[i] = energy(pos,vel)
jt[i] = angmomz(pos,vel)
plt.plot(times,et)
plt.show()
plt.plot(times,jt)
plt.show()
show_stats(et)
show_stats(jt)
In [ ]:
Clearly our initial design of the "phase" structure does not match that of the "orbit" very well. How should we improve this, so that we can write tools that use either output of our own code, with the one that comes out of scipy.integrate.odeint?
How about numpy vector operations to get Et and Jt, instead of element by element?