gistfile1.py

# Visualization of particles with gravity
# Adapted from pyopencl example
# Source: http://enja.org/2010/08/27/adventures-in-opencl-part-2-particles-with-opengl/

# Changelog: playing around with OO structures--nothing big

import pyopencl as cl # OpenCL - GPU computing interface
mf = cl.mem_flags
from pyopencl.tools import get_gl_sharing_context_properties
from OpenGL.GL import * # OpenGL - GPU rendering interface
from OpenGL.GLU import * # OpenGL tools (mipmaps, NURBS, perspective projection, shapes)
from OpenGL.GLUT import * # OpenGL tool to make a visualization window
from OpenGL.arrays import vbo
 
import numpy # Number tools
import sys   # System tools (path, modules, maxint)
import time  # for, get this, time functions

class ComputeObject(object):
    _INITIALIZED = False

    def __init__(self):
        if not self._INITIALIZED:
            self.platform = cl.get_platforms()[0]
            print self.platform

            self.context = cl.Context(properties=[(cl.context_properties.PLATFORM, 
                                      self.platform)] + get_gl_sharing_context_properties())

            print self.context
            self.queue = cl.CommandQueue(self.context)
            self._INITIALIZED

class ParticleObject(ComputeObject):
    def __init__(self):
        ComputeObject.__init__(self)

    def create_program(self, kernel):
        self.program = cl.Program(self.context, kernel).build()

class Fountain(ParticleObject):
    NUM_PARTICLES = int(16*65536)
    CALCS_PER_KERNEL = 1
    GLOBAL_SIZE = (NUM_PARTICLES/CALCS_PER_KERNEL,)
    LOCAL_SIZE = None

    def __init__(self):  
        # We need to set up the GL world before we initialize the Compute context
        self.window = self.glut_window()

        (self.np_position, 
         self.np_velocity, 
         self.gl_position, 
         self.gl_color) = self.initial_buffers(self.NUM_PARTICLES)

        # Set up the compute context now
        ParticleObject.__init__(self)  

        # Create the CL buffers (the ones only used by OpenCL
        self.cl_velocity = cl.Buffer(self.context, mf.COPY_HOST_PTR, hostbuf=self.np_velocity)
        self.cl_start_position = cl.Buffer(self.context, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=self.np_position)
        self.cl_start_velocity = cl.Buffer(self.context, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=self.np_velocity)

        # Create the GL buffers (the ones used both by OpenCL and OpenGL)
        self.cl_gl_position = cl.GLBuffer(self.context, mf.READ_WRITE, int(self.gl_position.buffers[0]))
        self.cl_gl_color = cl.GLBuffer(self.context, mf.READ_WRITE, int(self.gl_color.buffers[0]))
 
        self.create_program(self.create_kernel())
        
        self.time_func = time.time
        self.start_time = self.time_func()
        self.frame_count = 0

    def create_kernel(self):
        return """
            __kernel void particle_fountain(__global float4* position, 
                                            __global float4* color, 
                                            __global float4* velocity, 
                                            __global float4* start_position, 
                                            __global float4* start_velocity, 
                                            float time_step,
                                            int calcs_per_kernel)
            {
                unsigned int j = get_global_size(0);
                unsigned int i = get_global_id(0);
                int k;

                for (k = 0; k < calcs_per_kernel*j; k += j) {
                    float4 p = position[i+k];
                    float4 v = velocity[i+k];
                    float life = velocity[i+k].w;
                    life -= time_step;
                    if (life <= 0.f)
                    {
                        p = start_position[i+k];
                        v = start_velocity[i+k];
                        life = 1.0f;    
                    }

                    v.z -= 9.8f*time_step;
                    p.x += v.x*time_step;
                    p.y += v.y*time_step;
                    p.z += v.z*time_step;
                    v.w = life;

                    position[i+k] = p;
                    velocity[i+k] = v;

                    color[i+k].w = life; /* Fade points as life decreases */
                }
            }"""

    def glut_window(self):
        self.width = 1280 
        self.height = 720 
        self.time_step = .01
        self.mouse_down = False
        self.mouse_old = {'x': 0., 'y': 0.}
        self.rotate = {'x': -38., 'y': 14., 'z': 0.}
        self.translate = {'x': 0., 'y': 0., 'z': 0.}
        self.initial_translate = {'x': 0., 'y': 0., 'z': -2.5}

        glutInit(sys.argv)
        glutInitDisplayMode(GLUT_RGBA | GLUT_DOUBLE | GLUT_DEPTH)
        glutInitWindowSize(self.width, self.height)
        glutInitWindowPosition(0, 0)
        window = glutCreateWindow("Particle Simulation")

        glutDisplayFunc(self.on_display)  # Called by GLUT every frame
        glutKeyboardFunc(self.on_key)
        glutMouseFunc(self.on_click)
        glutMotionFunc(self.on_mouse_move)
        glutTimerFunc(10, self.on_timer, 15)  # Call draw every 16 ms

        glViewport(0, 0, self.width, self.height)
        glMatrixMode(GL_PROJECTION)
        glLoadIdentity()
        gluPerspective(60., self.width / float(self.height), .1, 1000.)

        return(window)

    def initial_buffers(self, num_particles, colorful=True):
        np_position = numpy.ndarray((num_particles, 4), dtype=numpy.float32)
        np_color = numpy.ndarray((num_particles, 4), dtype=numpy.float32)
        np_velocity = numpy.ndarray((num_particles, 4), dtype=numpy.float32)

        np_position[:,0] = numpy.sin(numpy.arange(0., num_particles) * 2.001 * numpy.pi / num_particles) 
        np_position[:,0] *= numpy.random.random_sample((num_particles,)) / 3. + .2
        np_position[:,1] = numpy.cos(numpy.arange(0., num_particles) * 2.001 * numpy.pi / num_particles) 
        np_position[:,1] *= numpy.random.random_sample((num_particles,)) / 3. + .2
        np_position[:,2] = 0.
        np_position[:,3] = 1.

        if colorful:
            np_color[:,0] = numpy.random.random_sample((num_particles,))
            np_color[:,1] = numpy.random.random_sample((num_particles,))
            np_color[:,2] = numpy.random.random_sample((num_particles,))
            np_color[:,3] = 1.
        else:
            np_color[:,:] = [1.,1.,1.,1.]

        np_velocity[:,0] = np_position[:,0] * 2.
        np_velocity[:,1] = np_position[:,1] * 2.
        np_velocity[:,2] = 3.
        np_velocity[:,3] = numpy.random.random_sample((num_particles, ))
    
        gl_position = vbo.VBO(data=np_position, usage=GL_DYNAMIC_DRAW, target=GL_ARRAY_BUFFER)
        gl_position.bind()
        gl_color = vbo.VBO(data=np_color, usage=GL_DYNAMIC_DRAW, target=GL_ARRAY_BUFFER)
        gl_color.bind()

        return (np_position, np_velocity, gl_position, gl_color)

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

    def on_key(self, *args):
        if args[0] == '\033' or args[0] == 'q':
            sys.exit()

    def on_click(self, button, state, x, y):
        self.mouse_old['x'] = x
        self.mouse_old['y'] = y

    def on_mouse_move(self, x, y):
        self.rotate['x'] += (y - self.mouse_old['y']) * .2
        self.rotate['y'] += (x - self.mouse_old['x']) * .2

        self.mouse_old['x'] = x
        self.mouse_old['y'] = y
        #print rotate

    def on_display(self):
        """Render the particles"""        
        # Update or particle positions by calling the OpenCL kernel
        cl.enqueue_acquire_gl_objects(self.queue, (self.cl_gl_position, self.cl_gl_color))

        kernelargs = (self.cl_gl_position, 
                      self.cl_gl_color, 
                      self.cl_velocity, 
                      self.cl_start_position, 
                      self.cl_start_velocity, 
                      numpy.float32(self.time_step),
                      numpy.int32(self.CALCS_PER_KERNEL))

        self.program.particle_fountain(self.queue,
                                       self.GLOBAL_SIZE, 
                                       self.LOCAL_SIZE, 
                                       *(kernelargs))

        cl.enqueue_release_gl_objects(self.queue, (self.cl_gl_position, self.cl_gl_color))

        self.queue.finish()
        glFlush()

        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
        glMatrixMode(GL_MODELVIEW)
        glLoadIdentity()

        rz = self.rotate['z']
        rz += 37.
        rz = rz if rz < 360 else 0
        self.rotate['z'] = rz

        # Handle mouse transformations
        glTranslatef(self.initial_translate['x'], 
                     self.initial_translate['y'], 
                     self.initial_translate['z'])
                     
        glRotatef(self.rotate['x'], 1, 0, 0)
        glRotatef(self.rotate['y'], 0, 1, 0) #we switched around the axis so make this rotate_z
        #glRotatef(self.rotate['z'], 0, 0, 1) # krb

        glTranslatef(self.translate['x'], self.translate['y'], self.translate['z'])
    
        # Render the particles
        glEnable(GL_POINT_SMOOTH)
        glPointSize(1)
        glEnable(GL_BLEND)
        glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA)

        # Set up the VBOs
        self.gl_color.bind()
        glColorPointer(4, GL_FLOAT, 0, self.gl_color)
        self.gl_position.bind()
        glVertexPointer(4, GL_FLOAT, 0, self.gl_position)

        glEnableClientState(GL_VERTEX_ARRAY)
        glEnableClientState(GL_COLOR_ARRAY)

        # Draw the VBOs
        glDrawArrays(GL_POINTS, 0, self.NUM_PARTICLES)

        glDisableClientState(GL_COLOR_ARRAY)
        glDisableClientState(GL_VERTEX_ARRAY)

        glDisable(GL_BLEND)

        glutSwapBuffers()

        self.frame_count += 1
        current_time = self.time_func()
        elapsed_time = current_time - self.start_time

        if elapsed_time >= 1.0:
            self.start_time = current_time
            print "\r%.2f" % (self.frame_count / elapsed_time)
            self.frame_count = 0

def main():
    f = Fountain()
    glutMainLoop()

if __name__ == "__main__":
    main()