Homework 8
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, PillowWriter
# NOTE: We're cheating a bit here to avoid having to deal with the true scaling differences between e.g. an individual star and a supermassive black hole.
# We're also adjusting for the fact that we'd like to deal with reasonable size masses.
G = 6.6743 * 10 ** -5
def distance(vec1, vec2):
"""Calculate the distance between 2 points vec1 and vec2.
Note: we only care about the first 2 fields in vec1 and vec2 corresponding to the x,y positions
Args:
vec1: 1x5 vector of (x,y,mass,vx,vy)
vec2: 1x5 vector of (x,y,mass,vx,vy)
"""
return np.linalg.norm(vec1[:2] - vec2[:2])
def force(m1, m2, s):
"""Calculate the force between m1 and m2 over distance s
Args:
m1: int/float mass
m2: int/float mass
s: int/float distance
"""
# note, we're "cheating" a bit here to produce a better animation.
# if the distances are too small, we quantize them
return (G * m1 * m2) / (np.max([s, 10.0]) ** 2)
def angle(vec1, vec2):
"""Determine the angle between vec1 and vec2, in 2D
Note: we only care about the first 2 fields in vec1 and vec2 corresponding to the x,y positions
Args:
vec1: 1x5 vector of (x,y,mass,vx,vy)
vec2: 1x5 vector of (x,y,mass,vx,vy)
"""
return np.arctan2(vec2[1] - vec1[1], vec2[0] - vec1[0])
def forceVec(vec1, vec2):
"""Calculate the force vector between vec1 and vec2
Args:
vec1: 1x5 vector of (x,y,mass,vx,vy)
vec2: 1x5 vector of (x,y,mass,vx,vy)
The force vector should be of the form [Fx, Fy]
"""
mass1 = vec1[2]
mass2 = vec2[2]
d = distance(vec1, vec2)
F = force(mass1, mass2, d)
θ = angle(vec1, vec2)
return F * np.cos(θ), F * np.sin(θ) # tuple
def sumOfForces(objects):
"""Return an Nx2 array of the cumulative force acting on each object in x and y
Args:
objects: an Nx5 matrix of (x,y,mass,vx,vy)
Returns:
a (`count`, 2) array of Fx, Fy
The return value should look like:
[[Fx1, Fy1],
[Fx2, Fy2],
...
]
"""
# Non-vectorized solution
all_object_forces = []
for obj in objects:
object_forces = []
# First calculate the forces
for other_obj in objects:
if (obj == other_obj).all():
# same object, skip
continue
object_forces.append(forceVec(obj, other_obj))
# then sum them up:
sum_x = sum(x[0] for x in object_forces)
sum_y = sum(x[1] for x in object_forces)
all_object_forces.append([sum_x, sum_y])
return all_object_forces
def centerOfMass(objects):
"""Calculate the center of mass of the objects
Args:
objects: an Nx5 matrix of (x,y,mass,vx,vy)
Returns:
a single (x,y) pair corresponding to the center of mass
"""
return (objects[:, 0] * objects[:, 2]).sum() / objects[:, 2].sum(), (
objects[:, 1] * objects[:, 2]
).sum() / objects[:, 2].sum()
def generateObjects(count, center=0, width=10, include_blackhole=False):
"""Generate a `count` x 5 matrix of (x, y) coordinates corresponding to `count` number of arbitrarily placed elements
with initial velocities roughly orthogonal to the center of mass in a e.g. rotating
(x,y) coordinates should be generated according to a normal distribution.
Mass should be randomly generated integer between 10 and 100.
(vx,vy) velocities should be 0 for this homework.
Args:
count (int): the number of (x, y) pairs to generate
center (int): the center of our distribution (average if `method`=normal)
width (int): the width of our distribution (covariance if `method`=normal)
include_blackhole (bool): put a large mass in the center and give the bodies an initial angular velocity
Returns:
a (`count`, 5) array of positions, mass, velocities
The return value should look like:
[[x1, y1, mass1, vx1, vy1],
[x2, y2, mass2, vx2, vy2],
...
]
"""
# first, generate the position and mass distribution
positions = np.random.multivariate_normal(
[center, center],
[
[width, 0],
[0, width],
],
count,
)
weights = np.random.randint(10, 100, (count, 1))
velocities = np.zeros((count, 2))
data = np.concatenate((positions, weights, velocities), axis=1)
if include_blackhole:
# now calculate the center of mass
com = centerOfMass(data)
# now get the angles between the com and the data points
angles = [angle(com, datum[:2]) for datum in data]
# if each point is θ from com, then 90+θ (π/2+θ rad) should be good
angles = angles + np.ones(count) * np.pi / 2
# v = root(GM/r)
velocity_magnitudes = np.sqrt(np.array([vec[2] / distance(com, vec[:2]) for vec in data]))
velocities = np.array(
[
[np.cos(theta) * mag, np.sin(theta) * mag]
for theta, mag in zip(angles, velocity_magnitudes)
]
)
# slice velocities to data
data[:, 3:5] = velocities
# as a last step, put a stationary black hole at the center
data[0] = np.array([center, center, data[:, 2].max() * 1e4, 0, 0])
return data
def plotObjects(objects):
"""Draw a matplotlib chart of our objects.
We should use a scatter plot to show the objects in 2D space.
The size of the object should be proportional to its mass.
Finally we should annotate the center of mass with a red dot.
To do this last part, you should be able to run:
com = centerOfMass(objects)
scatter(com[0], com[1], s=25, c="red")
Args:
objects: an Nx5 matrix of (x,y,mass,vx,vy)
"""
objects = np.array(objects)
# Create a new figure
plt.figure(figsize=(14, 8))
# Scatter the main objects
plt.scatter(objects[:, 0], objects[:, 1], s=objects[:, 2], alpha=0.5)
# calculate the center of mass
com = centerOfMass(objects)
# plot it as a red dot
plt.scatter(com[0], com[1], s=25, c="red")
# show the plot
plt.show()
def animateObjectsFrame(objects, ax, lim=None):
"""Animate a matplotlib chart of our objects. `objects` argument is an Nx2 matrix of (x,y) coordinates"""
# clear the frame
ax.clear()
# scatter the data, use max size for object size
sizes = [min(100, x) for x in objects[:, 2]]
ax.scatter(objects[:, 0], objects[:, 1], s=sizes, alpha=0.5)
# plot the center of mass in red
com = centerOfMass(objects)
ax.scatter(com[0], com[1], s=25, c="red")
if lim:
ax.set_xlim(lim)
ax.set_ylim(lim)
_animations = []
def animate(objects, time, delta_t, lim=None, save_gif=False):
_animations.clear()
frames = int(time / delta_t)
fig, ax = plt.subplots(figsize=(14, 8))
def animateStep(i):
simulateStep(objects, delta_t)
animateObjectsFrame(objects, ax, lim)
ani = FuncAnimation(fig, animateStep, interval=100, frames=frames)
_animations.append(ani)
if save_gif:
# Save as GIF
writergif = PillowWriter(fps=30)
ani.save("plot.gif", writer=writergif)
else:
# Show the plot
plt.show()
def nextPosition(vec, delta_t):
"""Calculate the next position of `vec` given forces `forcevec`
Args:
vec: 1x5 vector of (x,y,mass,vx,vy)
Returns: (new x position, new y position)
"""
return vec[0] + vec[3] * delta_t, vec[1] + vec[4] * delta_t
def nextVelocity(vec, forcevec, delta_t):
"""Calculate the next position of `vec` given forces `forcevec`
Args:
vec: 1x5 vector of (x,y,mass,vx,vy)
forcevec: 1x2 vector of (Fx, Fy)
Returns: (new x velocity, new y velocity)
"""
return (
vec[3] + (forcevec[0] / vec[2]) * delta_t,
vec[4] + (forcevec[1] / vec[2]) * delta_t,
)
def nextPositionAndVelocity(vec, forcevec, delta_t):
"""Calculate the next velocity and position of vector based on the forces acting on it"""
next_x, next_y = nextPosition(vec, delta_t)
next_vx, next_vy = nextVelocity(vec, forcevec, delta_t)
vec[0] = next_x
vec[1] = next_y
vec[3] = next_vx
vec[4] = next_vy
def simulateStep(objects, delta_t):
"""Given the objects, which have the form [x, y, mass, xvel, yvel], calculate the sum of forces
acting on each object and update the object with their next position and velocity"""
# calculate net forces
forces = sumOfForces(objects)
# update positions
for i, vector in enumerate(objects):
nextPositionAndVelocity(vector, forces[i], delta_t)
def main():
time = 1000
delta_t = 10
# testing with a binary system with 0 net initial velocity is a good bet
# data = generateObjects(2, width=10, center=0, include_blackhole=False)
# animate(data, time, delta_t, lim=(-100, 100))
# now lets test with a realistic galaxy with a black hole
# (note that i picked specific numbers here and played with the function's default arguments
# in order to generate pretty picture. feel free to play around with your own values)
data = generateObjects(100, width=1e5, include_blackhole=True)
animate(data, time, delta_t, lim=(-1e3, 1e3))
# To save a gif, commment the above and uncommment this
# animate(data, time, delta_t, lim=(-1e3, 1e3), save_gif=True)
if __name__ == "__main__":
main()