pyOpenGLを使って分子動力学を可視化

以前ここで試したmatplotlibを使った簡単な分子動力学計算の可視化処理ですが、描画が遅いのでOpenGLを使って書き直してみた。python用にはPyOpenGLのパッケージが用意されている。結果だがとても速い！

■結果をgifアニメに落としたものが↓。
gifアニメだと実際の描画より遅いので感動が伝わりにくい。。。

import random
from OpenGL.GL import *
from OpenGL.GLU import *
from OpenGL.GLUT import *

from PIL import Image
from PIL import ImageOps

a  = 10
Dt = 0.1

#Define xy index
PX = 0
PY = 1
VX = 2
VY = 3
FX = 4
FY = 5

perticle_num = 1000
xy = [[0 for i in range(6)] for j in range(perticle_num)]
step = 1

def capture():
  global step
  step = step + 1
  pad_step = '{0:04d}'.format(step)
  print pad_step

  width = glutGet(GLUT_WINDOW_WIDTH)
  height = glutGet(GLUT_WINDOW_HEIGHT)

  glReadBuffer(GL_FRONT)
  glPixelStorei(GL_UNPACK_ALIGNMENT, 1)
  data = glReadPixels(0, 0, width, height, GL_RGBA, GL_UNSIGNED_BYTE)

  image = Image.fromstring("RGBA", (width, height), data)
  image = ImageOps.flip(image)
  image.save("test"+pad_step+".png")

def integ():
  #verlet integration
  particle_num = len(xy)
  pnum = 0
  while pnum < particle_num:
    #update velocity
    xy[pnum][VX] = xy[pnum][VX] + 0.5*Dt*xy[pnum][FX]
    xy[pnum][VY] = xy[pnum][VY] + 0.5*Dt*xy[pnum][FY]

    #update position
    xy[pnum][PX] = xy[pnum][PX] + Dt*xy[pnum][VX]
    xy[pnum][PY] = xy[pnum][PY] + Dt*xy[pnum][VY]

    #boundary condition
    if xy[pnum][PX] > a:
       xy[pnum][PX] -= a
    if xy[pnum][PY] > a:
       xy[pnum][PY] -= a
    if xy[pnum][PX] < 0:
       xy[pnum][PX] += a
    if xy[pnum][PY] < 0:
       xy[pnum][PY] += a
    pnum += 1


def update():
    # If you want to get image...
    # capture()

    integ()
    nnum = 0
    while nnum < len(xy):
       glPointSize(3.0)
       glColor3f(0.3, 0.3, 1.0)
       glBegin(GL_POINTS)
       glVertex2d(xy[nnum][PX], xy[nnum][PY])
       glEnd()
       nnum = nnum + 1

def init_set():
  pnum = 0
  xv_sum = 0
  yv_sum = 0
  particle_num = len(xy)
  while pnum < particle_num:
    xy[pnum][PX] = 0  # 0 x-position 
    xy[pnum][PY] = 0  # 1 y-position
    xy[pnum][VX] = random.uniform(-1,1) # 2 x-velocity
    xy[pnum][VY] = random.uniform(-1,1) # 3 y-velocity
    xv_sum += xy[pnum][VX]
    yv_sum += xy[pnum][VY]
    xy[pnum][FX] = 0   # 0 x-force
    xy[pnum][FY] = 0   # 1 y-force
    pnum += 1

  xv_sum = xv_sum / particle_num
  yv_sum = yv_sum / particle_num
  pnum = 0
  while pnum < particle_num:
    xy[pnum][VX] = xy[pnum][VX] - xv_sum
    xy[pnum][VY] = xy[pnum][VY] - yv_sum
    pnum += 1
  #exit()

def draw():
    glClearColor(0.0, 0.0, 0.0, 0.0)
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)

    glLoadIdentity()

    glColor3f(1.0,1.0,0.0)
    gluLookAt(a/2, a/2, 20.0  ,a/2 ,a/2 ,0 , 0, 1.0, 0.0)

    glBegin(GL_LINE_LOOP)
    glVertex2d(0,0)
    glVertex2d(0,a)
    glVertex2d(a,a)
    glVertex2d(a,0)
    glEnd()

    update()
    glutSwapBuffers()

def init():
    glClearColor(0.7, 0.7, 0.7, 0.7)

def idle():
    glutPostRedisplay()

def reshape(w, h):
    glViewport(0, 0,w,h)
    glMatrixMode(GL_PROJECTION)
    glLoadIdentity()
    gluPerspective(30.0, w/h, 1.0, 100.0)
    glMatrixMode (GL_MODELVIEW)


if __name__ == "__main__":

    init_set()

    glutInitWindowPosition(50, 50);
    glutInitWindowSize(340, 340);
    glutInit(sys.argv)
    glutInitDisplayMode(GLUT_RGBA | GLUT_DOUBLE )
    glutCreateWindow("pyOpenGL TEST")
    glutDisplayFunc(draw)
    glutReshapeFunc(reshape)
    init()

    glutIdleFunc(idle)
    glutMainLoop()