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()