Introduction to numerical simulations: The 2 Body Problem

Many problems in statistical physics and astrophysics require solving problems consisting of many particles at once (sometimes on the order of thousands or more!) This can't be done by the traditional pen and paper techniques you would encounter in a physics class. Instead, we must implement numerical solutions to these problems.

Today, you will create your own numerical simulation for a simple problem is that solvable by pen and paper already, the 2 body problem in 2D. In this problem, we will describe the motion between two particles that share a force between them (such as Gravity). We'll design the simulation from an astronomer's mindset with astronomical units in mind. This simulation will be used to confirm the general motion of the earth around the Sun, and later will be used to predict the motion between two stars within relatively close range.


We will guide you through the physics and math required to create this simulation.

First, a brief review of the kinematic equations (remembering Order of Operations or PEMDAS, and that values can be positive or negative depending on the reference frame):

  • new time = old time + time change ($t = t_0 + \Delta t$)

  • new position = old position + velocity x time change ($x = x_0 + v \times \Delta t$)

  • new velocity = old velocity + acceleration x time change ($v = v_0 + a \times \Delta t$)

The problem here is designed to use the knowledge of scientific python you have been developing this week.

Like any code in python, The first thing we need to do is import the libraries we need. Go ahead and import Numpy and Pyplot below as np and plt respectively. Don't forget to put matplotlib inline to get everything within the notebook.


In [ ]:
%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt

Now we will define the physical constants of our system, which will also establish the unit system we have chosen. We'll use SI units here. Below, I've already created the constants. Make sure you understand what they are before moving on.


In [ ]:
#Physical Constants (SI units)
G=6.67e-11 #Universal Gravitational constant in m^3 per kg per s^2
AU=1.5e11 #Astronomical Unit in meters = Distance between sun and earth
daysec=24.0*60*60 #seconds in a day

Next, we will need parameters for the simulation. These are known as initial condititons. For a 2 body gravitation problem, we'll need to know the masses of the two objects, the starting posistions of the two objects, and the starting velocities of the two objects.

Below, I've included the initial conditions for the earth (a) and the Sun (b) at the average distance from the sun and the average velocity around the sun. We also need a starting time, and ending time for the simulation, and a "time-step" for the system. Feel free to adjust all of these as you see fit once you have built the system!



a note on dt: As already stated, numeric simulations are approximations. In our case, we are approximating how time flows. We know it flows continiously, but the computer cannot work with this. So instead, we break up our time into equal chunks called "dt". The smaller the chunks, the more accurate you will become, but at the cost of computer time.


In [ ]:
#####run specific constants. Change as needed#####
#Masses in kg
Ma=6.0e24 #always set as smaller mass
Mb=2.0e30 #always set as larger mass

#Time settings
t=0.0 #Starting time
dt=.01*daysec #Time set for simulation
tend=300*daysec #Time where simulation ends

#Initial conditions (position [m] and velocities [m/s] in x,y,z coordinates)
#For Ma
xa=1.0*AU
ya=0.0

vxa=0.0
vya=30000.0

#For Mb
xb=0.0
yb=0.0

vxb=0.0
vyb=0.0

It will be nice to create a function for the force between Ma and Mb. Below is the physics for the force of Ma on Mb. How the physics works here is not important for the moment. Right now, I want to make sure you can translate the math shown into a python function. (I'll show a picture of the physics behind this math for those interested.)

$$\vec{F_g}=\frac{-GM_aM_b}{r^3}\vec{r}$$

and $$\vec{r}=(x_b-x_a)\hat{x}+ (y_b-y_a)\hat{y}$$ $$r^3=((x_b-x_a)^2+(y_b-y_a)^2)^{3/2}$$

If we break Fg into the x and y componets we get: $$F_x=\frac{-GM_aM_b}{r^3}r_x$$ $$F_y=\frac{-GM_aM_b}{r^3}r_y$$



So, $Fg$ will only need to be a function of xa, xb, ya, and yb. The velocities of the bodies will not be needed. Create a function that calculates the force between the bodies given the positions of the bodies. My recommendation here will be to feed the inputs as separate components and also return the force in terms of components (say, fx and fy). This will make your code easier to write and easier to read.


In [ ]:
#Function to compute the force between the two objects
def Fg(Ma,Mb,G,xa,xb,ya,yb):
    #Compute rx and ry between Ma and Mb
    rx=xb-xa
    ry=yb-ya#Write it in
    
    #compute r^3, remembering r=sqrt(rx^2+ry^2)
    r3=np.sqrt(rx**2+ry**2)**3 #Write in r^3 using the equation above. Make use of np.sqrt()
    
    #Compute the force in Newtons. Use the equations above as a Guide!
    fx=-G*Ma*Mb*rx/r3 #Write it in
    fy=-G*Ma*Mb*ry/r3 #Write it in
    
    return fx,fy #What do we return?

Now that we have our force function, we will make a new function which does the whole simulation for a set of initial conditions. We call this function 'simulate' and it will take all the initial conditions as inputs. It will loop over each time step and call the force function to find the new positions for the asteroids at each time step.

The first part of our simulate function will be to initialize the loop and choose a loop type, for or while. Below is the general outline for how each type of loop can go.


For loop:

  • initialize position and velocity arrays with np.zeros or np.linspace for the amount of steps needed to go through the simulation (which is numSteps=(tend-t)/dt the way we have set up the problem). The for loop condition is based off time and should read rough like: for i in range(numSteps)


    While loop:
  • initialize posistion and velocity arrays with np.array([]) and use np.append() to tact on new values at each step like so, xaArray=np.append(xaArray,NEWVALUE). The while condition should read, while t<tend

My preference here is while since it keeps my calculations and appending separate. But, feel free to use which ever feels best for you!

Now for the actual simulation. This is the hardest part to code in. The general idea behind our loop is that as we step through time, we calculate the force, then calculate the new velocity, then the new position for each particle. At the end, we must update our arrays to reflect the new changes and update the time of the system. The time is super important! If we don't change the time (say in a while loop), the simulation would never end and we would never get our result. :(

Outline for the loop (order matters here)

  • Calculate the force with the last known positions (use your function!)

  • Calculate the new velocities using the approximation: vb = vb + dt*fg/Mb and va= va - dt*fg/Ma Note the minus sign here, and the need to do this for the x and y directions!

  • Calculate the new positions using the approximation: xb = xb + dt*Vb (same for a and for y's. No minus problem here)

  • Update the arrays to reflect our new values

  • Update the time using t=t+dt



    Now when the loop closes back in, the cycle repeats in a logical way. Go one step at a time when creating this loop and use comments to help guide yourself. Ask for help if it gets tricky!


In [ ]:
def simulate(Ma,Mb,G,xa,ya,vxa,vya,xb,yb,vxb,vyb):
    t=0
    #Run a loop for the simulation. Keep track of Ma and Mb posistions and velocites
    #Initialize vectors (otherwise there is nothing to append to!)
    xaAr=np.array([])
    yaAr=np.array([])

    vxaAr=np.array([])
    vyaAr=np.array([])

    xbAr=np.array([])#Write it in for Particle B
    ybAr=np.array([])#Write it in for Particle B

    vxbAr=np.array([])
    vybAr=np.array([])

    #using while loop method with appending. Can also be done with for loops
    while t<tend: #Write the end condition here.
        #Compute current force on Ma and Mb. Ma recieves the opposite force of Mb
        fx,fy=Fg(Ma,Mb,G,xa,xb,ya,yb)
    
        #Update the velocities and positions of the particles
        vxa=vxa-fx*dt/Ma
        vya=vya-fy*dt/Ma#Write it in for y
    
        vxb=vxb+fx*dt/Mb#Write it in for x
        vyb=vyb+fy*dt/Mb
        
        xa=xa+vxa*dt
        ya=ya+vya*dt#Write it in for y
    
        xb=xb+vxb*dt#Write it in for x
        yb=yb+vyb*dt
    
        #Save data to lists
        xaAr=np.append(xaAr,xa)
        yaAr=np.append(yaAr,ya)
    
        xbAr=np.append(xbAr,xb)#How will we append it here?
        ybAr=np.append(ybAr,yb)
    
        #update the time by one time step, dt
        t=t+dt
    return(xaAr,yaAr,xbAr,ybAr)

Now we will call our simulate function with the initial conditions we defined earlier! We will take the output of simulate and store the x and y positions of the two particles.


In [ ]:
#####Reminder of specific constants. Change as needed#####
#Masses in kg
Ma=6.0e24 #always set as smaller mass
Mb=2.0e30 #always set as larger mass

#Time settings
t=0.0 #Starting time
dt=.01*daysec #Time set for simulation
tend=300*daysec #Time where simulation ends

#Intial conditions (posistion [m] and velocities [m/s] in x,y,z coordinates)
#For Ma
xa=1.0*AU
ya=0.0

vxa=0.0
vya=30000.0

#For Mb
xb=0.0
yb=0.0

vxb=0.0
vyb=0.0


#Do simulation with these parameters
xaAr,yaAr,xbAr,ybAr = simulate(Ma,Mb,G,xa,ya,vxa,vya,xb,yb,vxb,vyb)#Insert the variable for y position of B particle)

Now for the fun part (or not so fun part if your simulation has an issue), plot your results! This is something well covered in previous lectures. Show me a plot of (xa,ya) and (xb,yb). Does it look sort of familiar? Hopefully you get something like the below image (in units of AU).


In [ ]:
from IPython.display import Image
Image("Earth-Sun-averageResult.jpg")

In [ ]:
plt.figure()
plt.plot(xaAr/AU,yaAr/AU)
plt.plot(xbAr/AU,ybAr/AU)#Add positions for B particle)
plt.show()

Challenge #1: Random Sampling of Initial Simulation Conditions

Now let's try to plot a few different asteroids with different initial conditions at once! Let's first produce the orbits of three asteroids with different masses. Suppose the masses of all asteroids in the main asteroid belt follow a Gaussian distribution. The parameters of the distribution of asteroid masses are defined below.


In [ ]:
#Mass distribution parameters 
Mave=7.0e24 #The average asteroid mass
Msigma=1.0e24 #The standard deviation of asteroid masses
Size=3 #The number of asteroids we wish to simulate

We now wish to draw a random sample of asteroid masses from this distribution (Hint: Look back at Lecture #3).


In [ ]:
#Draw 3 masses from normally distributed asteroid mass distribution
MassAr = Msigma * np.random.randn(Size) + Mave #Add your normal a.k.a. Gaussian distribution function, noting that the input to your numpy random number generator function will be: (Size)

Now let's loop over our random asteroid sample, run simulate and plot the results, for each one!


In [ ]:
plt.figure()

for mass in MassAr:#What array should we loop over?:
    xaAr,yaAr,xbAr,ybAr=simulate(mass,Mb,G,xa,ya,vxa,vya,xb,yb,vyb,vyb)
    plt.plot(xaAr/AU,yaAr/AU,label='Mass = %.2e'%mass) #Provide labels for each asteroid mass so we can generate a legend.
    #Pro tip: The percent sign replaces '%.2e' in the string with the variable formatted the way we want!
    
plt.legend()
plt.show()

Going further:

Can you make a plot with 5 asteroid masses instead of 3? If you've got some extra time, now is a great chance to experiment with plotting various initial conditions and how the orbits change! What happens if we draw some random initial velocities instead of random masses, for example?


In [ ]:
#draw 5 normally distributed mass values using the above parameters:
Size=5
MassAr = Msigma * np.random.randn(Size) + Mave

plt.figure()

for mass in MassAr:
    xaAr,yaAr,xbAr,ybAr=simulate(mass,Mb,G,xa,ya,vxa,vya,xb,yb,vyb,vyb)
    plt.plot(xaAr/AU,yaAr/AU,label='Mass = %.2e'%mass) 

plt.legend()
plt.show()

#Draw 3 velocities from normally distributed asteroid mass distribution
Size = 3
Dimensions = 2
Vave=20000 #The average asteroid velocity in m 
Vsigma=6000 #The standard deviation of asteroid velocities in m 

#You can make normal arrays with different dimensions! See: https://docs.scipy.org/doc/numpy-1.15.1/reference/generated/numpy.random.randn.html

VelAr = Vsigma * np.random.randn(Size,Dimensions) + Vave #a 2D array

for v in VelAr:
    xaAr,yaAr,xbAr,ybAr=simulate(mass,Mb,G,xa,ya,v[0],v[1],xb,yb,vxb,vyb)
    plt.plot(xaAr/AU,yaAr/AU,label='Velocity of Ma: vx = %.2e, vy = %.2e'%(v[0],v[1])) 

plt.legend()
plt.show()

Challenge #2: Fancy Plotting Fun!

When showing off your results to people unfamiliar with your research, it helps to make them more easy to understand through different visualization techniques (like legends, labels, patterns, different shapes, and sizes). You may have found that textbooks or news articles are more fun and easy when concepts are illustrated colorfully yet clearly, such as the example figure below, which shows different annotations in the form of text:


In [ ]:
from IPython.display import Image
Image(filename="fig_example.jpg")

Additionally, publications won't always be printed in color, and not all readers have the ability to distinguish colors or text size in the same way, so differences in style improve accessibility as well.

Luckily, Matplotlib can do all of this and more! Let's experiment with some variations in how we can make our plots. We can use the 'marker =' argument in plt.plot to choose a marker for every datapoint. We can use the 'linestyle = ' argument to have a dotted line instead of a solid line. Try experimenting with the extra arguments in the below plotting code to make it look good to you!

Now, add some plotting arguments to your loop like those you experimented with above. Can you make your plot more interesting and clear by changing the plotting parameters, or adding new plotting commands?

See the jupyter notebook called Plotting Demos in this same folder for some more examples of ways to make your plots pop!


In [ ]:
SMALL_SIZE = 12
MEDIUM_SIZE = 15
BIGGER_SIZE = 20

plt.rc('font', size=SMALL_SIZE)          # controls default text sizes
plt.rc('axes', titlesize=BIGGER_SIZE)     # fontsize of the figure title
plt.rc('axes', labelsize=MEDIUM_SIZE)    # fontsize of the x and y labels
plt.rc('xtick', labelsize=SMALL_SIZE)    # fontsize of the x number labels
plt.rc('ytick', labelsize=SMALL_SIZE)    # fontsize of the y numer labels
plt.rc('legend', fontsize=MEDIUM_SIZE)    # legend fontsize

colors=['black','blue','orange'] 
markers=['x','*','+']
styles=['--','-',':']
plt.figure(figsize=(8,6))

dt=10*daysec #Increase time set for simulation to better show markers individually 

for mass,color,mrk,sty in zip(MassAr,colors,markers,styles):
    xaAr,yaAr,xbAr,ybAr=simulate(mass,Mb,G,xa,ya,vxa,vya,xb,yb,vyb,vyb)
    plt.plot(xaAr/AU,yaAr/AU,label='Mass = %.2e'%mass, color=color, marker=mrk,linestyle=sty,linewidth=mass/Mave) #weighting width of lines by mass
    
plt.legend()
plt.title('Asteroid Trajectories')
plt.xlabel('x position (m)')
plt.ylabel('y position (m)')
plt.show()

In [ ]: