ideal_gas_simulation_version_4.py

import numpy as np
from matplotlib import pyplot as plt
from matplotlib.animation import FuncAnimation
from scipy.spatial import distance

# switch to different window than Pycharm's.
# since pycharm doesn't support animation
plt.switch_backend("Qt5Agg")

# constants
box_size = 100  # the size of the box
radius = 5  # radius of the molecules
mass = 1  # the mass of every molecule
molecule_amount = 106  # the amount of molecules
dt = 0.001  # time step for Euler's method
step_amount = 50  # the amount times to do Euler's method per frame


# functions
def init_molecules():
    """
    Initialize three molecules with set positions and velocities

    Returns
    -------
    molecules : (np.ndarray, np.ndarray)
        a tuple of the positions of the molecules and their velocities in array form.
        with the x-values in the first column and the y-values in the second column
        each row is a separate molecule
    """

    # set the type to floats since it will error otherwise because
    # the velocity and time_step are also floats.
    positions = np.array([[0, 0], [10, 10], [-5, 5]], dtype=float)
    velocities = np.array([[1, 2], [-1, 2], [0.5, -1.5]])

    return positions, velocities


def init_molecules_random(N, box_size, radius):
    """
    Randomly generate N molecules within the confines of the box.

    Parameters
    ----------
    N : int
        the number of molecules
    box_size : int
        the size of the box
    radius : float
        the radius of each molecule

    Returns
    -------
    molecules : (np.ndarray, np.ndarray)
        a tuple of the positions of the molecules and their velocities in array form.
        with the x-values in the first column and the y-values in the second column
        each row is a separate molecule
    """
    # a velocity from -5 to 5 sounds nice
    velocity_range = 5

    # a spawn range from -box_size/2 + radius to box_size/2 - radius
    # so that we don't spawn in/outside the borders
    spawn_range = box_size - 2 * radius

    # generate arrays with N rows and 2 columns with random numbers from 0 to 1
    # multiply such that we get [-box_size/2+radius, box_size/2-radius)
    positions = np.random.rand(N, 2) * spawn_range - (spawn_range/2 - radius)

    # multiply such that we get [-velocity_range, velocity_range)
    velocities = np.random.rand(N, 2) * 2 * velocity_range - velocity_range

    return positions, velocities


def wall_collisions(r, v, radius, box_size):
    """
    Process collisions of the molecules with the walls.\n
    If a molecule comes in contact with the wall and is still moving
    in the direction of the wall it inverts the respective velocity component.

    Parameters
    ----------
    r : np.ndarray
        the positions of the molecules, with each row a molecule and in the columns the x and y component
    v : np.ndarray
        the velocities of the molecules, with each row a molecule and in the columns the x and y component
    radius : float
        the radius of each molecule
    box_size : int
        the size of the box

    Returns
    -------
    v : np.ndarray
        the updated velocities of the molecules after they collided with the walls

    """
    # invert the velocity of the corresponding axis
    # if it is on the edge or passed it
    # and it is going further over the edge

    # if we're crossing the right boundary and are still going right invert the x-velocity
    v[((r[:, 0] + radius) >= box_size / 2) & (v[:, 0] > 0), 0] *= -1

    # if we're crossing the left boundary and are still going left invert the x-velocity
    v[((r[:, 0] - radius) <= -box_size / 2) & (v[:, 0] < 0), 0] *= -1

    # if going we're crossing the upper boundary and are still going up invert the y-velocity
    v[((r[:, 1] + radius) >= box_size / 2) & (v[:, 1] > 0), 1] *= -1

    # if going we're crossing the lower boundary and are still going down invert the y-velocity
    v[((r[:, 1] - radius) <= -box_size / 2) & (v[:, 1] < 0), 1] *= -1

    return v


def resolve_collisions(N, r, v, radius):
    """
    Calculates the resultant velocity of the particles if they are colliding
    and assigns their new velocities. It calculates this using the formula given
    in the exercises document.

    Parameters
    ----------
    N : int
        the number of molecules
    r : np.ndarray
        the positions of the molecules, with each row a molecule and in the columns the x and y component
    v : np.ndarray
        the velocities of the molecules, with each row a molecule and in the columns the x and y component
    radius : float
        the radius of each molecule

    Returns
    -------
    v : np.ndarray
        the updated velocities of the molecules after they collided with each other
    """

    # calculate the distance for each particle to every particle (including self)
    # yields an NxN array
    distances = distance.cdist(r, r)

    # np.where yields a tuple with in one the row indexes and in the other the column indexes
    # where the given statement is True
    # checks if the molecules are colliding
    # unpack the two tuples are two separate arguments for zip
    for i, j in zip(*np.where(distances < 2 * radius)):

        # we don't need to check for any j lower than or equal to i
        # for equal it is the same particle so no collision exists
        # any lower j we already checked when we had a lower i.
        if j > i:
            # get current positions and velocities of both objects
            r1, r2 = r[i], r[j]
            v1, v2 = v[i], v[j]

            # calculate dr the difference in position
            dr1 = r1 - r2

            # the same in the other direction
            dr2 = -dr1

            # calculate the difference in velocity
            dv = v1 - v2

            # could've squared array and called sum but this is faster
            length_squared = dr1[0] ** 2 + dr1[1] ** 2

            # check if they are moving towards each other (formula from exercise document)
            if dr2 @ dv > 0:
                # calculate new velocities with simplified formula from exercise document
                # the simplification is elaborated on in the results document.
                common_term = (dv @ dr1) / length_squared * dr1

                v1_new = v1 - common_term
                v2_new = v2 + common_term

                # use a view to change the velocities of the objects in the original array
                v[i], v[j] = v1_new, v2_new

    # return the newly calculated velocities
    return v


def animate(num):
    """
    function that is called every frame-update from FuncAnimation.\n
    Updates the molecule plot artist and returns the artist as an iterable
    for re-rendering.
    Parameters
    ----------
    num : int
        the frame number, unused
    Returns
    -------
    artists : iterable_of_artists
        Iterable of artists
    """
    # we will re-assign both r and v so we need them from the global namespace
    # and not the local one
    global r, v

    # execute Euler's method 'step_amount' times
    for _ in range(step_amount):
        # check for molecule collisions and assign new velocities
        v = resolve_collisions(molecule_amount, r, v, radius)

        # check for wall collisions and assign new velocities
        v = wall_collisions(r, v, radius, box_size)

        # update positions for some time-step dt
        r += v * dt

    # set the new data points (new molecule positions)
    line.set_data(r[:, 0], r[:, 1])

    # FuncAnimation requires an iterable so add a trailing comma to make it a tuple
    # which is iterable
    return line,


# initialize 'molecule_amount' molecules with random positions and velocities
r, v = init_molecules_random(molecule_amount, box_size, radius)

# make the plot and store the references for animation for part c
# higher dpi to make the plot bigger if you don't change this it will screw with the rendering
# causing particles to appear inside each other or the wall to be misplaced for the bouncing
# fillstyle none so they are hollow
fig = plt.figure(dpi=150)
line, = plt.plot(r[:, 0], r[:, 1], 'o', fillstyle='none')

# set the size for the plot of the box
plt.xlim(-box_size/2, box_size/2)
plt.ylim(-box_size/2, box_size/2)

# set the aspect ration to equal to, so it is square
fig.gca().set_aspect("equal", "box")

# calculate and set marker size for part d
ms = 2 * radius * fig.gca().get_window_extent().height / box_size * 72./fig.dpi
line.set_markersize(ms)

# animates the plot with a 10ms interval with blit so we don't redraw the entire frame
# using our the figure we made earlier 'fig' and using our animate function to update
# the frames
FuncAnimation(fig, animate, repeat=False, interval=10, blit=True)
plt.show()