Week 5 starting notebook - PHYS 633 Astrophysical Dynamics

This notebook will start you off in Week 5. It sets up a disk of planetesimals with uniform surface density in circular orbits around a central star.

Things to note:

The code starts up the built-in REBOUND visualization web-server with sim.start_server(port=1234). If you go to http://localhost:1234 in your browser, you should see an animation of the particles as the simulation runs, together with some useful information such as the number of steps per second. Press ‘h’ on this screen to see a list of keyboard commands.
After the simulation ends, the code uses the stored particle positions to generate plots of particle positions at each timestep. You will need to mkdir png to create a directory for the png files. You’ll find a shell script images_to_movie.sh in the main directory of the github repository that you can use to turn these png files in to a movie. The syntax is images_to_movie.sh 'png/plot%3d.png' movie.mp4. The script uses ffmpeg which you will need to install if you don’t have it (I used brew install ffmpeg).
Because we have many particles, the forces are calculated using a tree code (sim.gravity = "tree") in which far away particles are grouped together for the purposes of calculating the gravitational force. This improves how the simulation time scales with number of particles from $\propto N^2$ to $\propto N\log N$ which gives a big speed up.
The tree code requires us to define a boxsize for the simulation. We use 'open' boundary conditions, so that particles that leave the box are removed. If you look in the top left of the visualization window you will see that the total number of particles drops slightly over time as particles leave the simulation box.
We also set a softening length (sim.softening = 0.01) which softens the gravitational force at small separations (the interparticle potential is $\propto 1/\sqrt{r^2+\epsilon^2}$ rather than $\propto 1/r$ ). This avoids numerical errors that arise when two particles pass very close to one another and experience very large accelerations.

%config InlineBackend.figure_format = 'retina'
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
import rebound
import time
import numpy as np
import random

plt.rcParams['xtick.direction'] = 'in'
plt.rcParams['ytick.direction'] = 'in'
plt.rcParams['xtick.top'] = True
plt.rcParams['ytick.right'] = True

def show_progress(i):
    # shows progress during long simulations
    print(".", end="", flush=True)
    if not (i+1)%80:
        print() # start a new line
    return

def setup_disk(N=10000, mass=1e-2):
    # initialize a circular disk of N planetesimals
    for i in range(N):
        # choose a random radial location and angle around the orbit
        r = np.random.uniform(0.3, 3.0)
        theta = np.random.uniform(0, 2*np.pi)
        sim.add(
            # set the planetesimal mass so that we get the correct total mass for the disk
            m=mass/N,
            # circular orbit
            x=r*np.cos(theta),
            y=r*np.sin(theta),
            vx=-np.sin(theta)/np.sqrt(r),
            vy=np.cos(theta)/np.sqrt(r)
        )
    return

start = time.time()

sim = rebound.Simulation()

# start the visualization server: in your browser go to http://localhost:1234
sim.start_server(port=1234)

# use a tree code since we will have many particles
sim.integrator = "leapfrog"
sim.gravity = "tree"
sim.dt = 0.01
# soften the interaction for close approaches
sim.softening = 0.01

# define a box size for the simulation and add "open" boundary conditions
# particles that leave the box will be removed
boxsize = 10
sim.configure_box(boxsize)
sim.boundary = "open"

# Star
# we'll use a hash to label it so we can find it later
sim.add(m=1.0, hash="star")

# Planetesimals
Nparticles = 10000
setup_disk(N=Nparticles, mass=1e-2)

sim.move_to_com()

Noutputs = 20
Norbits = 10

print("Running simulation:")
times = np.linspace(0, Norbits * 2*np.pi, Noutputs)
dt = times[-1]/(Noutputs-1)
xy = np.zeros((Noutputs, Nparticles+1, 2))
xs = np.zeros(Noutputs)
ys = np.zeros(Noutputs)
for i,t in enumerate(times):
    sim.integrate(t, exact_finish_time=1)
    show_progress(i)
    # store the particle positions
    for j, p in enumerate(sim.particles):
        # store the (x,y) locations of each particle
        xy[i][j] = [p.x, p.y]
    # the star position is included in xy[][] but we'll also store it separately:
    xs[i] = sim.particles["star"].x
    ys[i] = sim.particles["star"].y

# shut down the visualization server; this will stop it complaining next time we run the simulation
sim.stop_server(port=1234)
print('\nSim time = %.3f s' % (time.time()-start,))

# The (x,y) limits for the snapshots
L = boxsize/2

print("Making plots:")
for i,t in enumerate(times):
    show_progress(i)
    
    # plot the current configuration of the particles in the planet's frame
    plt.figure(figsize=(6,6), dpi=150)    
    plt.xlim((-L,L))
    plt.ylim((-L,L))
    
    plt.plot(xy[i,:,0],xy[i,:,1],'ko',ms=1,markeredgewidth=0)
    plt.plot(xs,ys,'bo',ms=3)

    plt.gca().set_aspect('equal', adjustable='box')
    plt.savefig('png/plot%03d.png' % (i,))
    if i==Noutputs-1: # show the last plot
        plt.show()
    plt.close()

print('\nTotal time = %.3f s' % (time.time()-start,))

# to make a movie you can use:
# images_to_movie.sh 'png/plot%3d.png' movie.mp4