particles.py

"""
Animation of Elastic collisions with Gravity

author: Jake Vanderplas
email: [email protected]
website: http://jakevdp.github.com
license: BSD
Please feel free to use and modify this, but keep the above information. Thanks!

Some amendments made, code doesn't function 100% properly just yet

"""

import numpy as np
from scipy.spatial.distance import pdist, squareform
import math
import sys
import os

import matplotlib.pyplot as plt
import scipy.integrate as integrate
import matplotlib.animation as animation

num_particles = 2000
counter = 0
initial = True

class ParticleBox:

    def __init__(self, init_state=[[1,0,0,-1],
                                   [-0.5,0.5,0.5,0.5],
                                   [-0.5,-0.5,-0.5,0.5]],

                 bounds = [-2, 2, -2, 2],
                 size = 0.04,
                 M = 0.05):

        self.init_state = np.asarray(init_state, dtype=float)
        self.M = M * np.ones(self.init_state.shape[0])
        self.size = size
        self.state = self.init_state.copy()
        self.time_elapsed = 0
        self.bounds = bounds

    def grab_state(self):
        return self.state

    def grab_positions(self):
        return self.state[:, :2]

    def grab_x_velocities(self):
        return self.state[:, 2]

    def grab_y_velocities(self):
        return self.state[:, 3]

    def grab_elapsed_time(self):
        return self.time_elapsed


    def step(self, dt):
        global velocities
        self.time_elapsed += dt

        #update positions
        self.state[:, :2] += dt * self.state[:, 2:]

        # find pairs of particles undergoing a collision
        D = squareform(pdist(self.state[:, :2]))
        ind1, ind2 = np.where(D < 2 * self.size)
        unique = (ind1 < ind2)
        ind1 = ind1[unique]
        ind2 = ind2[unique]

        # update velocities of colliding pairs
        for i1, i2 in zip(ind1, ind2):
            # mass
            m1 = self.M[i1]
            m2 = self.M[i2]

            # location vector
            r1 = self.state[i1, :2]
            r2 = self.state[i2, :2]

            # velocity vector
            v1 = self.state[i1, 2:]
            v2 = self.state[i2, 2:]

            # relative location & velocity vectors
            r_rel = r1 - r2
            v_rel = v1 - v2

            # momentum vector of the center of mass
            v_cm = (m1 * v1 + m2 * v2) / (m1 + m2)

            # collisions of spheres reflect v_rel over r_rel
            rr_rel = np.dot(r_rel, r_rel)
            vr_rel = np.dot(v_rel, r_rel)
            v_rel = 2 * r_rel * vr_rel / rr_rel - v_rel

            # assign new velocities
            self.state[i1, 2:] = v_cm + v_rel * m2 / (m1 + m2)
            self.state[i2, 2:] = v_cm - v_rel * m1 / (m1 + m2) 

        # check for crossing boundary
        crossed_x1 = (self.state[:, 0] < self.bounds[0] + self.size)
        crossed_x2 = (self.state[:, 0] > self.bounds[1] - self.size)
        crossed_y1 = (self.state[:, 1] < self.bounds[2] + self.size)
        crossed_y2 = (self.state[:, 1] > self.bounds[3] - self.size)

        self.state[crossed_x1, 0] = self.bounds[0] + self.size
        self.state[crossed_x2, 0] = self.bounds[1] - self.size

        self.state[crossed_y1, 1] = self.bounds[2] + self.size
        self.state[crossed_y2, 1] = self.bounds[3] - self.size

        self.state[crossed_x1 | crossed_x2, 2] *= -1
        self.state[crossed_y1 | crossed_y2, 3] *= -1



#------------------------------------------------------------
# set up initial state
np.random.seed(0)
init_state = -0.5 + np.random.random((num_particles, 4))
init_state[:, :2] *= 3.9

box = ParticleBox(init_state, size=0.01)
dt = 1. / 30 # 30fps


#------------------------------------------------------------
# set up figure and animation
fig = plt.figure()
fig.subplots_adjust(left=0, right=1, bottom=0, top=1)
ax = fig.add_subplot(111, aspect='equal', autoscale_on=False,
                     xlim=(-3.2, 3.2), ylim=(-2.4, 2.4))


# particles holds the locations of the particles
particles, = ax.plot([], [], 'bo', ms=6)


# rect is the box edge
rect = plt.Rectangle(box.bounds[::2],
                     box.bounds[1] - box.bounds[0],
                     box.bounds[3] - box.bounds[2],
                     ec='none', lw=2, fc='none')
ax.add_patch(rect)

                     
def init():
    """initialize animation"""
    global box, rect
    particles.set_data([], [])
    rect.set_edgecolor('none')
    return particles, rect

def animate(i):
    """perform animation step"""
    global box, rect, dt, ax, fig, counter
    box.step(dt)

    ms = int(fig.dpi * 2 * box.size * fig.get_figwidth()
             / np.diff(ax.get_xbound())[0])

    
    v_x = np.asarray(box.grab_x_velocities())
    v_y = np.asarray(box.grab_y_velocities())
    v_res = []

    for counter in range(0, num_particles):
        v_res.append(abs(math.sqrt(math.pow(box.grab_x_velocities()[counter], 2) + math.pow(box.grab_y_velocities()[counter], 2))))
        counter += 1

    

    resolution = 100

    group_width = (max(v_res) - min(v_res)) / resolution

    v_freq = np.zeros(resolution+5)

    counter = 0

    for speed in v_res:
        current_speed = min(v_res)
        while speed > current_speed:
            current_speed += group_width
            counter += 1
            #print counter
        v_freq[counter] += 1
        counter = 0
        
            
    #line.set_data(np.arange(min(v_res), max(v_res), resolution), v_freq)
 
    
    # update pieces of the animation
    rect.set_edgecolor('k')
    particles.set_data(box.state[:, 0], box.state[:, 1])
    particles.set_markersize(ms)

    if box.grab_elapsed_time() >= 5:
        os.system("rm Boltzmann.txt")
        fob = open("Boltzmann.txt", "w")
        for element in v_freq:
            fob.write(str(element) + ",")
        fob.close()
        sys.exit()
    
    return particles, rect



ani = animation.FuncAnimation(fig, animate, frames=600,
                              interval=10, blit=True, init_func=init)


# save the animation as an mp4.  This requires ffmpeg or mencoder to be
# installed.  The extra_args ensure that the x264 codec is used, so that
# the video can be embedded in html5.  You may need to adjust this for
# your system: for more information, see
# http://matplotlib.sourceforge.net/api/animation_api.html
#ani.save('particle_box.mp4', fps=30, extra_args=['-vcodec', 'libx264'])

plt.show()