# Simulates the positions of point masses acting under their mutual gravity
# Given their masses, initial positions, and initial velocities
# With N particles, there are N(N-1) pairs.
# This is currently set for three point masses in two dimensions. The
# code can be easily altered to accommodate three dimensions (the variable
# Nd) and more point masses (the variable N).
# My next ambition for this code is to animate the trajectories.
from math import sqrt
import numpy as np
import copy as cp
import matplotlib.pyplot as plt
Nd = 2 # Number of spatial dimensions x, y, z
N = 3 # Number of particles
maxiters = 1000 # Number of time steps
FracErr = 10.0 # Allowed fractional error (meters) in unit time
Np = (N-1)*N//2 # Number of pairwise interactions, triangle number
Ncoord = N*Nd # total count of spatial coordinates
Nv = Ncoord*2 # Number of dependent variables: Nd coordinates for
# position for each particle, and Nd coordinates for
# velocity for each particle
G = 6.672e-11 # Gravitational constant in Nm^2/kg^2
hstart = 3600. # starting time step, I suppose in seconds
def RKStep(y, h, drv): # Take one Runge-Kutta step
# Fourth order Runge-Kutta code
# y are the dependent variables ([Nv] array)
# h is the step size in the independent variable (time step)
# drv is the function computing derivatives of y with respect
# to time (components of velocity and acceleration)
# and returning them in another [Nv] array
# The k's are [Nv] arrays used in 4th order R-K
k1 = h*drv(y)
k2 = h*drv(y+0.5*k1)
k3 = h*drv(y+0.5*k2)
k4 = h*drv(y+k3)
ynew = y + (k1+2.0*(k2+k3)+k4)/6.0
return ynew
def pairindex(a, b): # Finds the location in a 1-D array of the
# pairwise interaction between particle a and item b
# First check for valid indices
i = -1 # return -1 if index failure
if (max(a,b) > N-1) or (min(a,b) < 0):
print("An index is out of range. (",a,",",b,")")
if a == b:
print("No pair of particle",a,"with itself.")
elif a > b:
print("Pair indices should be low, high, not (",a,",",b,")")
i = b*(N-1) - b*(b+1)//2 + a - 1 # return the complementary index
else: # b > a as intended
i = a*(N-1) - a*(a+1)//2 + b - 1
return i
def pairfind(a, b, propty): # Finds the pairwise property between a and b
# a and b are the particles in the pair
# propty is an [Np] array of some pair property, such as distance
if (max(a,b) > N-1) or (min(a,b) < 0):
print("An index in (",a,",",b,"is out of range (",N-1,")")
p = propty[0]
elif a == b:
print("Pairfind call with equal indices:",a)
p = propty[0]
else: # a and b are different indices
if a < b:
p = propty[pairindex(a, b)]
else: # b < a
p = -propty[pairindex(b, a)]
print("Pairfind call indices out of order (",a,",",b,")")
return p
def coordlabel(c): # Coordinate index 0, 1, or 2 for x, y, z
# Used for debugging or reporting of results
if c == 0:
dx = "x"
elif c == 1:
dx = "y"
elif c == 2:
dx = "z"
else:
dx = "unknown coordinate"
return dx
def Displacements(r): # Return displacement vectors given positions
# All of these arrays contain Np elements, one for each pair of particles
displ = np.empty((Np,Nd),float) # pairwise displacement vectors
dsqr = np.zeros(Np,float) # squared distances
dist = np.empty(Np,float) # distances
dhat = np.empty((Np,Nd),float) # Unit vectors of displacements
pair = 0 # Index of a pair
for a in range(N): # First particle in the pair
for b in range(a+1, N): # Second particle in the pair
for c in range(Nd): # Coordinates
## dispmnt = r[b*Nd+c] - r[a*Nd+c] # for debugging
displ[pair,c] = r[b*Nd+c] - r[a*Nd+c] # Displacement from a to b
dsqr[pair] += displ[pair,c]**2 # squared distance
dist[pair] = sqrt(dsqr[pair]) # Distance between a and b
dhat[pair,:] = displ[pair,:]/dist[pair] # unit vector
pair += 1 # Increment to the next pair
return displ, dsqr, dist, dhat
# displ[p] is the displacement vector of pair p
# dsqr[p] is the squared distance of pair p
# dist[p] is the distance between the particles of pair p
# dhat[p] is the unit vector in the direction of displ[p]
def Attractions(dist2, mass, dhat):
# Vectors of attractions (forces) between particles
attr = np.empty((Np,Nd), float) # array of pairwise attractions
pair = 0
for a in range(N): # First particle in pair
for b in range(a+1, N): # Second particle in pair
F = G * mass[a] * mass[b] / dist2[pair] # Newtonian gravity
attr[pair,:] = dhat[pair,:] * F # Components of the force
pair += 1 # Increment to the next pair
return attr # Array of pairwise attractin vectors
def AccelVec(AttVec, mass): # Acceleration vectors of particles
# AttVec is the [Np] array of attraction pairs, already calculated
# mass is the [N] array of particle masses
F = np.empty(Nd, float) # Net force on one particle at a time
accel = np.empty((N,Nd), float) # Acceleration Nd-vectors of the N particles
for a in range(N): # Total force and acceleration of each particle
F.fill(0.0)
for b in range(0,a):
F -= pairfind(b, a, AttVec) # Add the attraction force toward b
for b in range(a+1, N):
F += pairfind(a, b, AttVec)
## acdebug = F/mass[a] # for debugging
accel[a,:] = F/mass[a] # Components of acceleration of particle a
return accel # [N] array of acelerations of the N particles
def f(r): # The derivatives of the phase variables with respect to time
fs = np.empty(Nv, float) # array of derivatives with respect to time
fs[:Ncoord] = r[Ncoord:] # dx/dt = v
displ, dsqr, dist, dhat = Displacements(r)
att = Attractions(dsqr, m, dhat) # array of pairwise attractions
acc = AccelVec(att, m) # array of particle acceleration vectors
for a in range(N): # particle accelerations converted to system of ODEs
ai = Ncoord + a * Nd
fs[ai:ai+Nd] = acc[a,:] # copy accelerations to the derivative list
## dbg = fs[ai:ai+Nd] # for debugging
return fs # [Nv] array of time-derivatives of the phase variables
# Create and initialize the particle positions and velocities
r = np.empty(Nv,float) # Positions and velocities of particles
# in a linear array of length Nv
### Initialize positions and velocities of two bodies
##r[0] = -1.5e11 # Coordinates of particle 1, m
##r[1] = 0.
##r[2] = 7.5e10 # Coordinates of particle 2, m
##r[3] = 0.
##r[4] = 0. # Velocity of particle 1, m/s
##r[5] = -30000.
##r[6] = 0. # Velocity of particle 2, m/s
##r[7] = 15000.
# Initialize positions and velocities of three bodies
r[0] = -1.5e11 # Coordinates of particle 1, m
r[1] = 0.
r[2] = 7.5e10 # Coordinates of particle 2, m
r[3] = 0.
r[4] = 0. # Coordinates of particle 3, m
r[5] = 1.0e11
r[6] = 0. # Velocity of particle 1, m/s
r[7] = -30000.
r[8] = 0. # Velocity of particle 2, m/s
r[9] = 15000.
r[10] = 100. # Velocity of particle 3, m/s
r[11] = -50.
# Create and initialize the particle masses
m = np.empty(N, float) # Masses of particles
##m[:] = 2.0e30, 4.0e30 # Masses in kg; 1 and 2 solar masses
m[:] = 2.0e30, 4.0e30, 6.0e29 # masses in kg
tpts = [] # List of Simulation times, if you need them for anything
# Such as an animation
coord = [] # List of particle coordinates at each time step
xyz = np.empty((N,Nd),float) # One set of coordinates
rhistory = [] # List of dependent variables at each time step
h = hstart # Starting time step
t = 0.
# Save the first values of the trajectory variables
tpts.append(t) # Time
for a in range(N):
xyz[a,:] = r[a*Nd:(a+1)*Nd]
coord.append(cp.copy(xyz)) # positions in a 2D array
rhistory.append(r) # positions and velocities in a 1D array
for iters in range(maxiters): # Do the time steps
# Use adaptive time step size
r1 = cp.copy(r)
r2 = cp.copy(r)
rm = RKStep(r1, h, f)
r1 = RKStep(rm, h, f) # Two steps of R-K of size h
r2 = RKStep(r2, 2.0*h, f)# One step of R-K of size 2h
rerr = r1[:Ncoord] - r2[:Ncoord] # Differences in spatial
# coordinates returned by taking two small time steps
# vs. one large step
denom = sqrt(np.dot(rerr,rerr)) # Square root of
# sum of squares of differences (sort of standard deviation)
if denom < 1.0e-12: # Avoid division by zero
ratio = 32.
else:
ratio = 30.0 * FracErr * h / denom # Using error in position
if ratio > 1.0: # Take bigger steps
r = rm # Keep the accurate estimate
t += h # time of saved estimate
adjust = ratio**0.20 # Now compute the next step size
if adjust > 2.0:
adjust = 2.0
h *= adjust # Increase next step size, up to a factor of 2
else: # ratio < 1: Take smaller steps
h *= ratio**0.25
r = RKStep(r, h, f) # Keep this estimate
t += h # time of saved estimate
# Add to the lists of variables
tpts.append(t) # Times
for a in range(N):
xyz[a,:] = r[a*Nd:(a+1)*Nd]
coord.append(cp.copy(xyz)) # positions. Making a copy to
# circumvent mutability issues
rhistory.append(r) # Dependent variables (positions and velocities)
# Plot the trajectories for two bodies
##x1, x2, y1, y2 = [], [], [], [] # x and y coordinates
##Npoints = len(tpts)
##for i in range(Npoints):
## x1.append(coord[i][0,0])
## y1.append(coord[i][0,1])
## x2.append(coord[i][1,0])
## y2.append(coord[i][1,1])
##
##plt.plot(x1, y1, label=str(m[0])+" kg")
##plt.plot(x2, y2, label=str(m[1])+" kg")
##plt.xlabel("meters")
##plt.ylabel("meters")
##plt.title("two-body gravity")
##plt.legend()
##plt.show()
# Plot the trajectories for three bodies
x1, x2, x3, y1, y2, y3 = [], [], [], [], [], [] # x and y coordinates
Npoints = len(tpts) # Shhould be the same as maxiters
for i in range(Npoints):
x1.append(coord[i][0,0])
y1.append(coord[i][0,1])
x2.append(coord[i][1,0])
y2.append(coord[i][1,1])
x3.append(coord[i][2,0])
y3.append(coord[i][2,1])
plt.plot(x1, y1, label=str(m[0])+" kg")
plt.plot(x2, y2, label=str(m[1])+" kg")
plt.plot(x3, y3, label=str(m[2])+" kg")
plt.xlabel("meters")
plt.ylabel("meters")
plt.title("three-body gravity")
plt.legend()
plt.show()