Paul Martini, Autumn 2019
import numpy as np
%matplotlib inline
import matplotlib.pyplot as plt
# Routines from DHW, modified by PM for notebook
def accelerate(pot,x,y,qh=1):
"""
Compute acceleration of a particle
pot = 1 for Kepler, 2 for H.O.
x, y = particle position
returns accelerations xdd, ydd
"""
if (pot==1):
r=np.sqrt(x**2+y**2)
xdd= -x/r**3
ydd= -y/r**3
elif (pot==2):
# fill in the next two lines
xdd= 0.
ydd= 0.
else:
raise ValueError ("Invalid potential type")
return([xdd,ydd])
def potential(pot,x,y,qh=1):
"""
Compute potential energy (per unit mass)
pot = 1 for Kepler, 2 for H.O.
x, y = particle position
returns potential phi
"""
if (pot==1):
phi=-1/np.sqrt(x**2+y**2)
elif (pot==2):
# fill in the next line
phi= 0.
else:
raise ValueError ("Invalid potential type")
return(phi)
def orbit_euler(params):
'''
Orbit calculation via the Euler Method
Routine by DHW, modified by PM for notebook I/O
Parameters
----------
params : dict
Dictionary of parameter values
params['xinit'] : initial x value
params['yinit'] : initial y value
params['xdotinit'] : initial x velocity
params['ydotinit'] : initial y velocity
params['tmax'] : maximum time
params['dtout'] : output time step
params['tstep'] : calculation time step
params['pot_type'] : potential type (1 = Kepler, 2 = H.O.)
Returns
-------
out : dict
Dictionary of 1-d float arrays:
out['t'] : time
out['x'] : x positions
out['y'] : y positions
out['xdot'] : x velocities
out['ydot'] : x velocities
out['de'] : change in energy
'''
x = params['xinit']
y = params['yinit']
xdot = params['xdotinit']
ydot = params['ydotinit']
tmax = params['tmax']
dtout = params['dtout']
tstep = params['tstep']
pot_type = params['pot_type']
i = 0
t = 0.0
de = 0
# Define arrays to hold the output
N_out = int(tmax / dtout)
t_out = np.arange(0, tmax, dtout)
x_out = np.zeros(len(t_out), dtype=float)
y_out = np.zeros(len(t_out), dtype=float)
xdot_out = np.zeros(len(t_out), dtype=float)
ydot_out = np.zeros(len(t_out), dtype=float)
de_out = np.zeros(len(t_out), dtype=float)
# Store initial values
t_out[i] = 0.0
x_out[i] = x
y_out[i] = y
xdot_out[i] = xdot
ydot_out[i] = ydot
de_out[i] = 0.
e0 = potential(pot_type, x, y, qh=params['qh']) + 0.5*(xdot**2 + ydot**2)
tnext = dtout
i=1
while (t<tmax):
xdd, ydd = accelerate(pot_type, x, y, qh=params['qh'])
x+=xdot*tstep
y+=ydot*tstep
xdot+=xdd*tstep
ydot+=ydd*tstep
t+=tstep
if (t>=tnext and t < tmax):
# print(t, i)
energy=potential(pot_type, x, y, qh=params['qh']) + 0.5*(xdot**2 + ydot**2)
de_out[i]=(energy-e0)/np.abs(e0)
# print('%8.5f %7.3f %7.3f %8.2f %8.2f %7.4e' % (t,x,y,xdot,ydot,de))
t_out[i] = t
x_out[i] = x
y_out[i] = y
xdot_out[i] = xdot
ydot_out[i] = ydot
tnext+=dtout
i+=1
# Put output in a dictionary
out = {}
out['t'] = t_out
out['x'] = x_out
out['y'] = y_out
out['xdot'] = xdot_out
out['ydot'] = ydot_out
out['de'] = de_out
return out
def add_axtitle(axes, label):
'''
Helper routine to add titles above subplot panels
Parameters
----------
axes : matplotlib axes instance
label : string
'''
ylim = axes.get_ylim()
ylab = ylim[0] + 1.02*(ylim[1]-ylim[0])
xlim = axes.get_xlim()
xlab = xlim[0] + 0.1*(xlim[1]-xlim[0])
axes.text(xlab, ylab, label, fontsize=14)
def plot4xy(coords4, limit, labels4, toplabel):
'''
Create a 4-panel plot of the orbits for 4 calculations
Parameters
----------
coords4 : list
list with 4 dictionaries that are output from orbit_euler or orbit_leapfrog
limit : float
plot range (-limit, limit) for x and y axes
labels4 : list
list of 4 strings to label each panel
toplabel : string
title of plot
'''
fig, axarr = plt.subplots(nrows=2, ncols=2, figsize=(12,12), sharex=True, sharey=True)
plt.subplots_adjust(wspace=0.1, hspace=0.1)
xlab = -0.83*limit
ylab = .83*limit
axarr[0,0].plot(coords4[0]['x'], coords4[0]['y'])
axarr[0,0].set_xlim(-1*limit, limits)
axarr[0,0].set_ylim(-1*limit, limits)
add_axtitle(axarr[0,0], labels4[0])
axarr[0,1].plot(coords4[1]['x'], coords4[1]['y'])
add_axtitle(axarr[0,1], labels4[1])
axarr[1,0].plot(coords4[2]['x'], coords4[2]['y'])
add_axtitle(axarr[1,0], labels4[2])
axarr[1,1].plot(coords4[3]['x'], coords4[3]['y'])
add_axtitle(axarr[1,1], labels4[3])
ax_main = fig.add_axes([.1, .1, 0.8, 0.8], frameon=False, xticks=[], yticks=[])
ax_main.set_xlabel("X", fontsize=18)
ax_main.set_ylabel("Y", fontsize=18)
ax_main.set_title(toplabel, fontsize=20)
p1param = {}
p1param['dtout'] = 0.02
p1param['tmax'] = 40.
p1param['xinit'] = 1.
p1param['yinit'] = 0.
p1param['xdotinit'] = 0.
p1param['ydotinit'] = 1.
p1param['pot_type'] = 1
p1param['qh'] = 0.
# Case 1
p1param['tstep'] = 1.e-2
e1out = orbit_euler(p1param)
e1label = "Euler, dt={0:.2f}".format(p1param['tstep'])
# Case 3
# lf1out = orbit_leapfrog(p1param)
# next 3 lines are dummy values to illustrate the plotting script
lf1out = {}
lf1out['x'] = 0.
lf1out['y'] = 0.
lf1label = "Leapfrog, dt={0:.2f}".format(p1param['tstep'])
# Case 2
p1param['tstep'] = 1.e-4
e2out = orbit_euler(p1param)
e2label = "Euler, dt={0:.4f}".format(p1param['tstep'])
# Case 4
# lf2out = orbit_leapfrog(p1param)
# next 3 lines are dummy values to illustrate the plotting script
lf2out = {}
lf2out['x'] = 0.
lf2out['y'] = 0.
lf2label = "Leapfrog, dt={0:.4f}".format(p1param['tstep'])
coords = [e1out, e2out, lf1out, lf2out]
labels = [e1label, e2label, lf1label, lf2label]
limits = 1.8
toplabel = "Kepler Potential with $\.y_i = 1$"
plot4xy(coords, limits, labels, toplabel)
def orbit_3body(params, moon=False):
'''
Parameters
----------
params : dict
Dictionary of parameter values
params['m1'] : mass of body 1 (planet)
params['m2'] : mass of body 2 (moon)
params['x2'] : position of body 2
params['ydot1'] : initial y velocity of body 1
params['ydot2'] : initial y velocity of body 2
params['dtout'] : output time step
params['tmax'] : maximum time
params['tstep'] : calculation time step
moon : bool
True means ydot1, ydot2 are in units of v_{c,1}, v_{c,2}
Returns
-------
out : dict
Dictionary of 1-d float arrays:
out['t'] : time
out['x1'] : x positions for body 1
out['y1'] : y positions for body 1
out['x2'] : x positions for body 2
out['y2'] : y positions for body 2
out['x3'] : x positions for body 3
out['y3'] : y positions for body 3
out['de'] : change in energy
'''
# Your code here
out = {}
out['t'] = t_out
out['x1'] = x1_out
out['y1'] = y1_out
out['x2'] = x2_out
out['y2'] = y2_out
out['x3'] = x3_out
out['y3'] = y3_out
out['de'] = de_out
return out
def plot3body(params, out):
'''
Parameters
----------
params : dict
Dictionary of parameter values
params['m1'] : mass of body 1 (planet)
params['m2'] : mass of body 2 (moon)
params['x2'] : position of body 2
params['ydot1'] : initial y velocity of body 1
params['ydot2'] : initial y velocity of body 2
params['dtout'] : output time step
params['tmax'] : maximum time
params['tstep'] : calculation time step
out : dict
Dictionary of 1-d float arrays:
out['t'] : time
out['x1'] : x positions for body 1
out['y1'] : y positions for body 1
out['x2'] : x positions for body 2
out['y2'] : y positions for body 2
out['x3'] : x positions for body 3
out['y3'] : y positions for body 3
out['de'] : change in energy
'''
# Calculate the period of the planet orbit
Pplanet = 2.*out['t'][np.where(out['y2'] < 0.)][0]
# Create masks for each planet orbit
tp1mask = out['t'] < Pplanet
tp2mask = (out['t'] >= Pplanet)
tp2mask = tp2mask*(out['t'] < 2*Pplanet)
tp3mask = (out['t'] >= 2*Pplanet)
tp3mask = tp3mask*(out['t'] < 3*Pplanet)
tp4mask = (out['t'] >= 3*Pplanet) # >= 3 orbits of the planet
# Main panel
fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(10,10))
tlabel = "m$_p$ = {0:.2f}, m$_m$ = {1:.4f}, x$_m$ = {2:.2f}, $\.y_p/v_c$ = {3:.2f}, $\.y_m/y_c$ = {4:.2f}".format(params['m1'], params['m2'], params['x2'], params['ydot1'], params['ydot2'])
ax.set_title(tlabel, fontsize=16)
Plabel = "T$_p$ = {0:.4f}".format(Pplanet)
ax.text(-1.8, 1.8, Plabel, fontsize=16)
maxlabel = "t$_m$ = {0:.1f}".format(params['tmax'])
ax.text(-1.8, 1.65, maxlabel, fontsize=16)
# Draw the orbits
ax.plot(out['x1'], out['y1'], 'b', label='Star', lw=1)
ax.plot(out['x2'], out['y2'], 'k', label='Planet', lw=1)
ax.plot(out['x3'][tp1mask], out['y3'][tp1mask], 'r', label='Moon Orbit 1', lw=1)
ax.plot(out['x3'][tp2mask], out['y3'][tp2mask], 'lime', label='Moon Orbit 2', lw=1)
ax.plot(out['x3'][tp3mask], out['y3'][tp3mask], 'chocolate', label='Moon Orbit 3', lw=1)
ax.plot(out['x3'][tp4mask], out['y3'][tp4mask], 'cyan', label='Moon Orbit 4+', lw=1)
ax.set_xlim(-2, 2)
ax.set_ylim(-2, 2)
ax.legend(loc='lower left')
ax.set_xlabel('X', fontsize=16)
ax.set_ylabel('Y', fontsize=16)
# Add a panel that zooms in on the moon --
ax_moon = fig.add_axes([.7, .7, .15, .15])
ax_moon.set_title("Moon Orbit", fontsize=14)
x_moon = out['x3'] - out['x2']
y_moon = out['y3'] - out['y2']
ax_moon.plot(x_moon[tp1mask], y_moon[tp1mask], 'r')
ax_moon.plot(x_moon[tp2mask], y_moon[tp2mask], 'lime')
ax_moon.plot(x_moon[tp3mask], y_moon[tp3mask], 'chocolate')
ax_moon.plot(x_moon[tp4mask], y_moon[tp4mask], 'cyan')
val = 1.25*np.abs(1.-params['x2'])
ax_moon.set_xlim(-1.*val, val)
ax_moon.set_ylim(-1.*val, val)
# Add a panel that zooms in on the planet --
ax_planet = fig.add_axes([.7, .15, .15, .15])
ax_planet.set_title("Planet Orbit", fontsize=14)
ax_planet.plot(out['x1'], out['y1'], 'b')
ax_planet.plot(out['x2'], out['y2'], 'k')
ax_planet.plot(out['x3'][tp1mask], out['y3'][tp1mask], 'r')
ax_planet.plot(out['x3'][tp2mask], out['y3'][tp2mask], 'lime')
ax_planet.plot(out['x3'][tp3mask], out['y3'][tp3mask], 'chocolate')
ax_planet.plot(out['x3'][tp4mask], out['y3'][tp4mask], 'cyan')
ax_planet.set_xlim(1.-2.*val, 1.+2*val)
ax_planet.set_ylim(-2.*val, 2.*val)
p3params1 = {}
p3params1['dtout'] = 0.01
p3params1['tmax'] = 30.
p3params1['m1'] = 0.1 # Planet
p3params1['m2'] = 0.01 # Moon
p3params1['x2'] = 1.05
p3params1['ydot1'] = 1.
p3params1['ydot2'] = 1.
p3params1['tstep'] = 1.e-4
out1 = orbit_3body(p3params1, moon=True)
plot3body(p3params1, out1)