Homework 7

import numpy as np
import matplotlib.pyplot as plt


# 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 generateObjects(count, center=0, width=10):
    """Generate a `count` x 5 matrix of (x, y) coordinates corresponding to `count` number of arbitrarily placed elements

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

    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)
    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 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
    """
    return (G * m1 * m2) / (s ** 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
                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 main():
    # Generate a random set of objects
    objects = generateObjects(500, center=0, width=100)
    plotObjects(objects)

    # Calculate the center of mass
    com = centerOfMass(objects)
    print(f"Center of Mass: {com}")

    # Calculate the forces acting on each object
    forces = sumOfForces(objects)

    # Loop through the first 10 rows and print their positions, masses, and forces
    for i, obj in enumerate(objects[:10,:]):
        print(f"Object {i}: Position: {obj[:2]}, Mass: {obj[2]}, Forces: {forces[i]}")

if __name__ == "__main__":
    main()