Skip to content

Instantly share code, notes, and snippets.

@likr
Created December 15, 2012 11:57
Show Gist options
  • Save likr/4294168 to your computer and use it in GitHub Desktop.
Save likr/4294168 to your computer and use it in GitHub Desktop.
#!/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()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment