Astronomy 8824 - Problem Set 2 Start

Paul Martini, Autumn 2019

In [1]:
import numpy as np
%matplotlib inline
import matplotlib.pyplot as plt
In [2]:
# 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)

Part 1: Kepler Potential

In [3]:
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                                               
In [4]:
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)
In [5]:
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)

Case with ydot = 1

In [6]:
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'])
In [7]:
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)

Part 3: Planet and Moon

Stub of orbit_3body to demonstrate input and output

In [ ]:
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
In [8]:
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)

3.1 Case with M_planet = 0.1, M_moon = 0.01, x_moon = 1.05, v_planet/v_c = 1., v_moon/v_c = 1

In [11]:
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)
In [12]:
plot3body(p3params1, out1)