runge_kutta.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# vim: filetype=pyopencl.python

from __future__ import division

import sys
import math
import random
import itertools
import pyopencl as cl
from pyopencl.tools import get_gl_sharing_context_properties
from pyopencl import mem_flags
from pyopencl import clmath
from pyopencl import clrandom
from OpenGL.GL import *
from OpenGL.GLU import *
from OpenGL.GLUT import *
from OpenGL.arrays.vbo import VBO
import numpy

src = '''//CL//
float2 vtdt(float x, float y, float t, float dt, float tau)
{
    float2 v;
    const float sin_pi_x = sin(M_PI * x);
    const float sin_pi_y = sin(M_PI * y);
    const float cos_pi_x = cos(M_PI * x);
    const float cos_pi_y = cos(M_PI * y);
    const float cos_pi_t = 2.0 * cos(M_PI * t / tau);
    v.x = -cos_pi_t * sin_pi_x * sin_pi_x * cos_pi_y * sin_pi_y * dt;
    v.y = cos_pi_t * cos_pi_x * sin_pi_x * sin_pi_y * sin_pi_y * dt;
    return v;
}

__kernel void rungekutta(
    __global float4* pos,
    float t,
    float dt,
    float tau
)
{
    const int i = get_global_id(0);
    const float4 p = pos[i];
    const float2 d = vtdt(p.x, p.y, t, dt, tau);
    p.x += d.x;
    p.y += d.y;
    pos[i] = p;
}

__kernel void rungekutta4(
    __global float4* pos,
    float t,
    float dt,
    float tau
)
{
    const int i = get_global_id(0);
    const float4 p = pos[i];
    float2 d;
    float xm, ym, xs, ys;
    float xg = p.x;
    float yg = p.y;

    d = vtdt(xg, yg, t, dt, tau);
    xm = xg + 0.5 * d.x;
    ym = yg + 0.5 * d.y;
    xs = xg + d.x / 6.0;
    ys = yg + d.y / 6.0;

    d = vtdt(xm, ym, t + 0.5 * dt, dt, tau);
    xm = xg + 0.5 * d.x;
    ym = yg + 0.5 * d.y;
    xs += d.x / 3.0;
    ys += d.y / 3.0;

    d = vtdt(xm, ym, t + 0.5 * dt, dt, tau);
    xm = xg + d.x;
    ym = yg + d.y;
    xs += d.x / 3.0;
    ys += d.y / 3.0;

    d = vtdt(xm, ym, t + dt, dt, tau);
    p.x = xs + d.x / 6.0;
    p.y = ys + d.y / 6.0;
    pos[i] = p;
}
'''

def colors(n):
    d = 1 / (n - 1)
    vals = [d * i for i in range(n)]
    return [c + (1.,) for c in itertools.product(*((vals,) * 3))][1:-1]

class RungeKutta(object):
    def __init__(self, N, dt, tau, W, H, steps, groups, loop=None):
        self.N = N
        self.dt = numpy.float32(dt)
        self.tau = numpy.float32(tau)
        self.W = W
        self.H = H
        self.steps = steps
        self.colors = colors(3)
        random.shuffle(self.colors)
        self.t = 0
        self.running = False
        self.loop = loop

        # init OpenGL
        glutInit(sys.argv)
        glutInitDisplayMode(GLUT_RGBA | GLUT_DOUBLE | GLUT_DEPTH)
        glutInitWindowSize(W, H)
        glutInitWindowPosition(0, 0)
        self.window = glutCreateWindow(b'PyCon JP 2012 Sprints Demo')

        # init scene
        glViewport(0, 0, W, H)
        glMatrixMode(GL_PROJECTION)
        glLoadIdentity()
        gluPerspective(60., W / H, 0.1, 1000.)
        glMatrixMode(GL_MODELVIEW)

        # register callback
        glutDisplayFunc(self.draw)
        glutTimerFunc(30, self.timer, 30)
        glutKeyboardFunc(self.on_key)

        # generate data
        position,color = self.generate_data(1., 1., 0.15, groups)
        self.position_vbo = VBO(position)
        self.position_vbo.bind()
        self.color_vbo = VBO(color)
        self.color_vbo.bind()

        # init OpenCL
        properties = get_gl_sharing_context_properties()
        self.context = cl.Context(properties=properties)
        self.queue = cl.CommandQueue(self.context)
        program = cl.Program(self.context, src).build(cache_dir=False)
        self.rungekutta = program.rungekutta4

        # init buffer
        flag = mem_flags.READ_WRITE
        self.buffer = cl.GLBuffer(self.context, flag, int(self.position_vbo))

    def generate_data(self, W, H, R, groups=1):
        pos = numpy.empty((self.N, 4), dtype=numpy.float32)
        color = numpy.empty((self.N, 4), dtype=numpy.float32)
        n = self.N / groups
        for i in range(groups):
            w = random.random() * (W - 2 * R) + R
            h = random.random() * (H - 2 * R) + R
            r = R * numpy.random.rand(n).astype(numpy.float32)
            theta = 2 * math.pi * numpy.random.rand(n).astype(numpy.float32)
            pos[n * i:n * (i + 1),0] = r * numpy.cos(theta) + w
            pos[n * i:n * (i + 1),1] = r * numpy.sin(theta) + h
            pos[n * i:n * (i + 1),2] = 0.
            pos[n * i:n * (i + 1),3] = 1.
            color[n * i:n * (i + 1),:] = self.colors[i]
        return pos,color

    def timer(self, t):
        glutTimerFunc(t, self.timer, t)
        glutPostRedisplay()

    def on_key(self, *args):
        self.running = not self.running

    def draw(self):
        # calculate positions
        if self.running:
            cl.enqueue_acquire_gl_objects(self.queue, [self.buffer])
            for _ in range(self.steps):
                t = numpy.float32(self.t * self.dt)
                if self.loop and self.loop <= t / self.tau:
                    break
                args = (self.buffer,t,self.dt,self.tau)
                self.rungekutta(self.queue, (self.N,), (1,), *args)
                self.t += 1
            cl.enqueue_release_gl_objects(self.queue, [self.buffer])
            self.queue.finish()
            glFlush()

        # prepare GL
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
        glMatrixMode(GL_MODELVIEW)
        glLoadIdentity()
        gluLookAt(0.5, 0.5, 1, 0.5, 0.5, 0, 0, 1, 0)

        # draw points
        glPointSize(1)
        self.color_vbo.bind()
        glColorPointer(4, GL_FLOAT, 0, self.color_vbo)
        self.position_vbo.bind()
        glVertexPointer(4, GL_FLOAT, 0, self.position_vbo)
        glEnableClientState(GL_VERTEX_ARRAY)
        glEnableClientState(GL_COLOR_ARRAY)
        glDrawArrays(GL_POINTS, 0, self.N)
        glDisableClientState(GL_COLOR_ARRAY)
        glDisableClientState(GL_VERTEX_ARRAY)
        glutSwapBuffers()

    def start(self):
        glutMainLoop()

if __name__ == '__main__':
    main = RungeKutta(10000, 0.001, 8.0, 500, 500, 10, 4, None)
    main.start()