Ideas Gas in box simulation

#!/usr/bin/env python3.3
# -*- coding: utf8 -*-

from OpenGL.GL import *
from OpenGL.GLU import *
from OpenGL.GLUT import *
import sys
import numpy as np
from scipy.spatial.distance import pdist, squareform

# Window size
width = height = 720
global axrng
axrng = 1.0


dt=0.001
# Animation control and state
global anim
anim = 0

# FPS calculation
global frameCount
global previousTime
PrintFPS = 1
previousTime = 0
frameCount = 0

def init():
    glClearColor(0.0,0.0,0.0,1.0)
    glOrtho(-axrng,axrng,-axrng,axrng,axrng,-axrng)

def idle():
    global anim
    if anim == 1:
        if PrintFPS:
            calculateFPS()
        glutPostRedisplay()

def idle2():
    global anim
    if anim == 1:
        glutPostRedisplay()

def calculateFPS():
    global frameCount,previousTime
    frameCount+=1
    currentTime = glutGet(GLUT_ELAPSED_TIME);
    timeInterval = currentTime - previousTime;
    if timeInterval > 1000:
        fps = frameCount / (timeInterval / 1000.)
        previousTime = currentTime;
        frameCount = 0;
        print (fps)

def keyboard(key,x_,y_):
    global anim
    if key == (b'\x1b') or key == (b'q') :
        sys.exit()
    if key == (b'a'):
        anim = 1
    if key == (b's'):
        anim = 0

def anim():
    glClear(GL_COLOR_BUFFER_BIT|GL_DEPTH_BUFFER_BIT)
    Gas.draw()
    Gas.step()
    glutSwapBuffers()

def anim2():
    glClear(GL_COLOR_BUFFER_BIT|GL_DEPTH_BUFFER_BIT)
    Gas.drawHIST()
    glutSwapBuffers()

def reshape(w,h):
    if h == 0:
        h = 1
    glViewport(0,0,w,h)
    glMatrixMode(GL_PROJECTION)
    glLoadIdentity()
    if w <= h:
        glOrtho(-axrng,axrng,-axrng*h/w,axrng*h/w,axrng,-axrng)
    else:
        glOrtho(-axrng*w/h,axrng*w/h,-axrng,axrng,axrng,-axrng)
    glMatrixMode(GL_MODELVIEW)
    glLoadIdentity()

def main():
    glutInit(sys.argv)
    glutInitDisplayMode(GLUT_RGB|GLUT_DOUBLE|GLUT_DEPTH)
    glutInitWindowPosition(0,0)
    glutInitWindowSize(width,height)
    glutCreateWindow(b'Particles in the box')
    init()
    glutReshapeFunc(reshape)
    glutDisplayFunc(anim)
    glutKeyboardFunc(keyboard)
    glutIdleFunc(idle)
    glutInitWindowPosition(720,0)
    glutInitWindowSize(600,600)
    glutCreateWindow(b'Speed distribution')
    init()
    glutDisplayFunc(anim2)
    glutReshapeFunc(reshape)
    glutIdleFunc(idle2)
    glutMainLoop()

class gas:
    def __init__(self,size):
        self.time=0
        self.size=size
        self.rad=0.005
        self.G=-9.8
        self.state=np.random.random((self.size,4))
        self.state[:,0:2]=-1+(2)*(self.state[:,0:2]-self.state[:,0.:2].max())/(-self.state[:,0.:2].max()-self.state[:,0.:2].min())
        #self.state[:,2:4]=self.state[:,2:4]*20-10
        self.state[:,2:4]=np.zeros(self.state[:,2:4].shape)
    def drawHIST(self):
        glPushMatrix()
        glColor4f(1.0,1.0,1.0,0.5)
        glBegin(GL_LINES)
        glVertex2f(-axrng*0.95,-axrng*0.95)
        glVertex2f(axrng*0.95,-axrng*0.95)
        glVertex2f(-axrng*0.95,-axrng*0.95)
        glVertex2f(-axrng*0.95,axrng*0.95)
        glEnd()
        glColor4f(0.0,0.0,1.0,0.5)
        HISTarray,bins=np.histogram(np.sqrt(self.state[:,2:3]*self.state[:,2:3]+self.state[:,3:4]*self.state[:,3:4]),bins=self.size/25)
        HISTarray=-axrng*0.95+(axrng*0.95+axrng*0.95)*(HISTarray-HISTarray.min())/(HISTarray.max()-HISTarray.min())
        bins=-axrng*0.95+(axrng*0.95+axrng*0.95)*(bins-bins.min())/(bins.max()-bins.min())
        width=bins[1]-bins[0]
        glBegin(GL_QUADS)
        for a in np.arange(len(HISTarray)):
            glVertex2f(bins[a]-width/2,-axrng*0.95)
            glVertex2f(bins[a]-width/2,HISTarray[a])
            glVertex2f(bins[a]+width/2,HISTarray[a])
            glVertex2f(bins[a]+width/2,-axrng*0.95)
        glEnd()
        glPopMatrix()
        glutPostRedisplay()
    def draw(self):
        glPushMatrix()
        glColor4f(1.0,1.0,1.0,0.5)
        glEnable(GL_VERTEX_ARRAY)
        glVertexPointer(2,GL_FLOAT,0,self.state[:,:2])
        glDrawArrays(GL_POINTS,0,self.size)
        glPopMatrix()
    def step(self):
        self.state[:,3:4]+=self.G/2*dt
        self.state[:,2:3][ self.state[:,0:1]+self.state[:,2:3]*dt >= axrng ]*=-1.
        self.state[:,2:3][ self.state[:,0:1]+self.state[:,2:3]*dt <=-axrng ]*=-1.
        self.state[:,3:4][ self.state[:,1:2]+self.state[:,3:4]*dt >= axrng ]*=-1.
        self.state[:,3:4][ self.state[:,1:2]+self.state[:,3:4]*dt <=-axrng ]*=-1.
        self.state[:,0:1]+=self.state[:,2:3]*dt
        self.state[:,1:2]+=self.state[:,3:4]*dt
        self.state[:,3:4]+=self.G/2*dt
        self.time+=dt
        # Piece of collision detection code taken from internet
        D = squareform(pdist(self.state[:, :2]))
        ind1, ind2 = np.where(D < 2 * self.rad)
        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)
            v_cm = (v1 + v2) / (2.)

            # 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  / (2.)
            self.state[i2, 2:] = v_cm - v_rel  / (2.)
        #print (str(self.time)+" "+str((self.state[:,2:3]*self.state[:,2:3]+self.state[:,3:4]*self.state[:,3:4]).sum()/2)+" "+str( (-self.G*self.state[:,1:2]).sum() ))

if __name__ == "__main__":
    Gas = gas(1000)
    main()