ローレンツアトラクタをPyOpenGLで描画

# coding: utf-8
# LorenzAttractor.py

class Calculator(object):
    # 解法の基底クラス。updateメソッドにLorenzAttractorクラスのインスタンスが渡される
    def __init__(self): pass
    def update(self, lorenz):
        raise Exception("must be overrided and not be called")

class Euler(Calculator):
    # オイラー法で解を求める
    def __init__(self):
        Calculator.__init__(self)
    def update(self, lorenz):
        x,y,z = lorenz.q
        p,b,r = lorenz.p, lorenz.b, lorenz.r
        val = [-p*x + p*y, -x*z + r*x - y, +x*y - b*z]
        return [lorenz.q[i] + lorenz.dt * val[i] for i in  [0,1,2]]

class LorenzAttractor(object):
    # nextpointメソッドでそれぞれのステップにおける点を返す
    def __init__(self, initpoint, p, r, b, dt, calculator=Euler()):
        self.p = p
        self.r = r
        self.b = b
        self.t = 0
        self.q = initpoint
        self.dt = dt
        self.calculator = calculator
    def nextpoint(self):
        self.q = self.calculator.update(self)
        self.t += self.dt
        return self.q
    def __repr__(self):
        return "<Lorenz p,r,b=%s,%s,%s, dt=%s, calculator=%s>" % (self.p, self.r, self.b, self.dt, self.calculator)

class LorenzAttractorContainer(LorenzAttractor):
    # それぞれのステップで計算した値を保持しておく
    def __init__(self, *args, **kwargs):
        LorenzAttractor.__init__(self, *args, **kwargs)
        self.points = [self.q]
    def nextpoint(self):
        p = LorenzAttractor.nextpoint(self)
        self.points.append(p)
        return p
    def __len__(self):
        return len(self.points)
    def __iter__(self):
        return iter(self.points)
    def __getitem__(self, k):
        return self.points[k]

# coding: utf-8
# lorenz attractor with PyOpenGL

import sys
from OpenGL.GL import *
from OpenGL.GLU import *
from OpenGL.GLUT import *
import LorenzAttractor

lorenz = LorenzAttractor.LorenzAttractorContainer([1,1,1], 10, 28, 8./ 3, 1e-3)
x,y,z = 52,29,-24
displistno = 100

def background():
    glMatrixMode(GL_MODELVIEW)
    glPushMatrix()
    glLoadIdentity()
    glMatrixMode(GL_PROJECTION)
    glPushMatrix()
    glLoadIdentity()
    glDisable(GL_DEPTH_TEST)
    glDisable(GL_LIGHTING)

    glBegin(GL_QUADS)
    glColor3f(.2, .7, .7)
    glVertex3f(-1, -1, 0)
    glVertex3f(+1, -1, 0)
    glColor3f(1, 1, 1)
    glVertex3f(+1, +1, 0)
    glVertex3f(-1, +1, 0)
    glEnd()
    
    glEnable(GL_LIGHTING)
    glEnable(GL_DEPTH_TEST)
    glMatrixMode(GL_PROJECTION)
    glPopMatrix()
    glMatrixMode(GL_MODELVIEW)
    glPopMatrix()


def display():
    glMatrixMode(GL_PROJECTION)
    glLoadIdentity()
    gluPerspective(60, float(glutGet(GLUT_WINDOW_WIDTH))/ glutGet(GLUT_WINDOW_HEIGHT), 1, 10000.)
    gluLookAt(x, y, z,  0, 24, 7,  0, 1, 0)
    glMatrixMode(GL_MODELVIEW)
    
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)   
    background()
    glCallList(displistno)

    glutSwapBuffers()

def reshape(w, h):
    glViewport(0, 0, w, h)

def keyboard(k, xx, yy):
    global x,y,z
    if k in ['Q', 'q', '\033']:
        sys.exit()
    if k == 'u': x += 1
    if k == 'j': x -= 1
    if k == 'i': y += 1
    if k == 'k': y -= 1
    if k == 'o': z += 1
    if k == 'l': z -= 1
    if k == 'p': print x,y ,z
    glutPostRedisplay()

def drawAttractor():
    glMaterialfv(GL_FRONT_AND_BACK, GL_AMBIENT, [0.3, 0.2, 0.7, 1])
    glMaterialfv(GL_FRONT_AND_BACK, GL_DIFFUSE, [0.3, 0.2, 0.1, 1])
    glMaterialfv(GL_FRONT_AND_BACK, GL_SPECULAR,[0.3, 0.2, 0.1, 1])
    glMaterialfv(GL_FRONT_AND_BACK, GL_SHININESS, [20])
    for i in xrange(0, len(lorenz), 15):
        p = lorenz[i]
        q = p[1], p[2], p[0] # OpenGLの座標系に直す
        glPushMatrix()
        glLoadIdentity()
        glTranslatef(q[0], q[1], q[2])
        glutSolidSphere(.7, 32, 32)
        glPopMatrix()

def init():
    global displistno
    glClearColor(.2, .7, .7, 1)
    
    while lorenz.t < 50:
        lorenz.nextpoint()
    # ディスプレイリストの作成
    displistno = glGenLists(1)
    glNewList(displistno, GL_COMPILE)
    drawAttractor()
    glEndList()
    # 照明の設定
    glEnable(GL_LIGHTING)
    glEnable(GL_LIGHT0)
    glLightfv(GL_LIGHT0, GL_POSITION, [10, 14, -16, 0])
    glLightfv(GL_LIGHT0, GL_AMBIENT , [0.8, 1, 1, 0])

    glEnable(GL_DEPTH_TEST)

def main():
    glutInit(sys.argv)
    glutInitDisplayMode(GLUT_RGB | GLUT_DEPTH | GLUT_DOUBLE)
    glutInitWindowSize(400, 400)
    glutInitWindowPosition(100, 100)
    glutCreateWindow("lorenz attractor")
    
    glutDisplayFunc(display)
    glutReshapeFunc(reshape)
    glutKeyboardFunc(keyboard)
    init()
    glutMainLoop()

if __name__ == '__main__':
    main()