格子ボルツマン法を用いた流体力学のシミュレーション(openGL)

以前ここのページで紹介した、格子ボルツマン法(LBM:Lattice Boltzmann Method)シミュレーションのコードをopenGLによる可視化に改良してみた。

import numpy
import time
import math
import matplotlib.pyplot
import matplotlib.animation
from OpenGL.GL import *
from OpenGL.GLU import *
from OpenGL.GLUT import *

numpy.set_printoptions(threshold=numpy.inf)

#--Define constants--
# lattice dimensions
height = 80
width = 200

# fluid viscosity
viscosity = 0.02

# "relaxation" parameter
omega = 1 / (3*viscosity + 0.5)


u0 = 0.10

four9ths = 4.0/9.0
one9th   = 1.0/9.0
one36th  = 1.0/36.0
performanceData = False
count = 1


n0 = four9ths * (numpy.ones((height,width)) - 1.5*u0**2)

nN = one9th * (numpy.ones((height,width)) - 1.5*u0**2)
nS = one9th * (numpy.ones((height,width)) - 1.5*u0**2)
nE = one9th * (numpy.ones((height,width)) + 3*u0 + 4.5*u0**2 - 1.5*u0**2)
nW = one9th * (numpy.ones((height,width)) - 3*u0 + 4.5*u0**2 - 1.5*u0**2)
nNE = one36th * (numpy.ones((height,width)) + 3*u0 + 4.5*u0**2 - 1.5*u0**2)
nSE = one36th * (numpy.ones((height,width)) + 3*u0 + 4.5*u0**2 - 1.5*u0**2)
nNW = one36th * (numpy.ones((height,width)) - 3*u0 + 4.5*u0**2 - 1.5*u0**2)
nSW = one36th * (numpy.ones((height,width)) - 3*u0 + 4.5*u0**2 - 1.5*u0**2)

rho = n0 + nN + nS + nE + nW + nNE + nSE + nNW + nSW

ux = (nE + nNE + nSE - nW - nNW - nSW) / rho
uy = (nN + nNE + nNW - nS - nSE - nSW) / rho


# Initialize barriers:
barrier = numpy.zeros((height,width), bool) 
barrier[int(height/2)-8:int(height/2)+8, int(height/2)] = True 
barrier[int(height/2)-8, int(height/2)+1] = True    
barrierN = numpy.roll(barrier,  1, axis=0) 
barrierS = numpy.roll(barrier, -1, axis=0) 
barrierE = numpy.roll(barrier,  1, axis=1) 
barrierW = numpy.roll(barrier, -1, axis=1)
barrierNE = numpy.roll(barrierN,  1, axis=1)
barrierNW = numpy.roll(barrierN, -1, axis=1)
barrierSE = numpy.roll(barrierS,  1, axis=1)
barrierSW = numpy.roll(barrierS, -1, axis=1)



def stream():
  global nN, nS, nE, nW, nNE, nNW, nSE, nSW
  nN  = numpy.roll(nN,   1, axis=0)
  nNE = numpy.roll(nNE,  1, axis=0)
  nNW = numpy.roll(nNW,  1, axis=0)
  nS  = numpy.roll(nS,  -1, axis=0)
  nSE = numpy.roll(nSE, -1, axis=0)
  nSW = numpy.roll(nSW, -1, axis=0)
  nE  = numpy.roll(nE,   1, axis=1)
  nNE = numpy.roll(nNE,  1, axis=1)
  nSE = numpy.roll(nSE,  1, axis=1)
  nW  = numpy.roll(nW,  -1, axis=1)
  nNW = numpy.roll(nNW, -1, axis=1)
  nSW = numpy.roll(nSW, -1, axis=1)
  nN[barrierN] = nS[barrier]
  nS[barrierS] = nN[barrier]
  nE[barrierE] = nW[barrier]
  nW[barrierW] = nE[barrier]
  nNE[barrierNE] = nSW[barrier]
  nNW[barrierNW] = nSE[barrier]
  nSE[barrierSE] = nNW[barrier]
  nSW[barrierSW] = nNE[barrier]

def collide():
  global rho, ux, uy, n0, nN, nS, nE, nW, nNE, nNW, nSE, nSW, u2
  rho = n0 + nN + nS + nE + nW + nNE + nSE + nNW + nSW
  ux = (nE + nNE + nSE - nW - nNW - nSW) / rho
  uy = (nN + nNE + nNW - nS - nSE - nSW) / rho
  ux2 = ux * ux 
  uy2 = uy * uy
  u2 = ux2 + uy2
  omu215 = 1 - 1.5*u2
  uxuy = ux * uy
  n0 = (1-omega)*n0 + omega * four9ths * rho * omu215
  nN = (1-omega)*nN + omega * one9th * rho * (omu215 + 3*uy + 4.5*uy2)
  nS = (1-omega)*nS + omega * one9th * rho * (omu215 - 3*uy + 4.5*uy2)
  nE = (1-omega)*nE + omega * one9th * rho * (omu215 + 3*ux + 4.5*ux2)
  nW = (1-omega)*nW + omega * one9th * rho * (omu215 - 3*ux + 4.5*ux2)
  nNE = (1-omega)*nNE + omega * one36th * rho * (omu215 + 3*(ux+uy) + 4.5*(u2+2*uxuy))
  nNW = (1-omega)*nNW + omega * one36th * rho * (omu215 + 3*(-ux+uy) + 4.5*(u2-2*uxuy))
  nSE = (1-omega)*nSE + omega * one36th * rho * (omu215 + 3*(ux-uy) + 4.5*(u2-2*uxuy))
  nSW = (1-omega)*nSW + omega * one36th * rho * (omu215 + 3*(-ux-uy) + 4.5*(u2+2*uxuy))

  nE[:,0] = one9th * (1 + 3*u0 + 4.5*u0**2 - 1.5*u0**2)
  nW[:,0] = one9th * (1 - 3*u0 + 4.5*u0**2 - 1.5*u0**2)
  nNE[:,0] = one36th * (1 + 3*u0 + 4.5*u0**2 - 1.5*u0**2)
  nSE[:,0] = one36th * (1 + 3*u0 + 4.5*u0**2 - 1.5*u0**2)
  nNW[:,0] = one36th * (1 - 3*u0 + 4.5*u0**2 - 1.5*u0**2)
  nSW[:,0] = one36th * (1 - 3*u0 + 4.5*u0**2 - 1.5*u0**2)

def curl(ux, uy):
  return numpy.roll(uy,-1,axis=1) - numpy.roll(uy,1,axis=1) - numpy.roll(ux,-1,axis=0) + numpy.roll(ux,1,axis=0)

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

    count += 1
    stream()
    collide()

    if count % 30 == 0:
      glPointSize(10);
      L2 = curl(ux,uy)
      vmin = L2.min()  
      vmax = L2.max()
      vv = abs(vmax) if abs(vmax) > abs(vmin) else abs(vmin)
      L3 =  2*L2/vv

      x = y = xi = yi = 0  
      kk = 0
      glBegin(GL_POINTS)
      for i in L3:
        for j in i:
          if x == width:
            y += 1
            x = 0
          else:
            pass  
          if j >0:
            glColor3f(j,0,0)
          else:
            glColor3f(0,-j,0)
          glVertex2d(x,y)
          x += 1
        #end for
      #end for
      glEnd()
      glutSwapBuffers()
    #endif


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

def idle():
    glutPostRedisplay()

def reshape(w, h):
    glViewport(0,0, w, h)
    glLoadIdentity()
    glOrtho(-2, width+2, -2, height+2, -10.0, 10.0)


if __name__ == "__main__":


    glutInit(sys.argv)
    glutInitDisplayMode(GLUT_RGBA | GLUT_DOUBLE | GLUT_DEPTH)
    glutInitWindowSize(width*3, height*3)
    glutCreateWindow("pyOpenGL TEST")
    glutDisplayFunc(draw)
    glutReshapeFunc(reshape)
    init()

    glutIdleFunc(idle)
    glutMainLoop()