Extract GL moments, particle buffers and add texture buffer

2 files changed, 180 insertions, 6 deletions

diff --git a/utility/opengl.py b/utility/opengl.py
new file mode 100644
index 0000000..afbfa00
--- /dev/null
+++ b/utility/opengl.py
@@ -0,0 +1,136 @@
+import pyopencl as cl
+mf = cl.mem_flags
+
+import numpy
+import sympy
+
+from mako.template import Template
+from pathlib import Path
+
+from OpenGL.GL import *
+from OpenGL.arrays import vbo
+
+class MomentsVertexBuffer:
+    def __init__(self, lattice):
+        self.lattice = lattice
+
+        self.np_moments = numpy.ndarray(shape=(self.lattice.memory.volume, 4), dtype=self.lattice.memory.float_type)
+        self.gl_moments = vbo.VBO(data=self.np_moments, usage=GL_DYNAMIC_DRAW, target=GL_ARRAY_BUFFER)
+        self.gl_moments.bind()
+        self.cl_gl_moments  = cl.GLBuffer(self.lattice.context, mf.READ_WRITE, int(self.gl_moments))
+
+        self.build_kernel()
+
+    def build_kernel(self):
+        program_src = Template(filename = str(Path(__file__).parent/'../template/opengl.mako')).render(
+            descriptor = self.lattice.descriptor,
+            geometry   = self.lattice.geometry,
+            memory     = self.lattice.memory,
+
+            moments_subexpr    = self.lattice.moments[0],
+            moments_assignment = self.lattice.moments[1],
+            collide_subexpr    = self.lattice.collide[0],
+            collide_assignment = self.lattice.collide[1],
+
+            float_type = self.lattice.float_type[1],
+
+            ccode = sympy.ccode
+        )
+        self.program = cl.Program(self.lattice.context, program_src).build(self.lattice.compiler_args)
+
+    def bind(self):
+        self.gl_moments.bind()
+        glEnableClientState(GL_VERTEX_ARRAY)
+        glVertexPointer(4, GL_FLOAT, 0, self.gl_moments)
+
+    def collect(self):
+        cl.enqueue_acquire_gl_objects(self.lattice.queue, [self.cl_gl_moments])
+
+        if self.lattice.tick:
+            self.program.collect_gl_moments(
+                self.lattice.queue,
+                self.lattice.grid.size(),
+                self.lattice.layout,
+                self.lattice.memory.cl_pop_b,
+                self.lattice.memory.cl_material,
+                self.cl_gl_moments)
+        else:
+            self.program.collect_gl_moments(
+                self.lattice.queue,
+                self.lattice.grid.size(),
+                self.lattice.layout,
+                self.lattice.memory.cl_pop_a,
+                self.lattice.memory.cl_material,
+                self.cl_gl_moments)
+
+
+class MomentsTexture:
+    def __init__(self, lattice):
+        self.lattice = lattice
+        self.gl_texture_buffer = numpy.ndarray(shape=(self.lattice.memory.volume, 4), dtype=self.lattice.memory.float_type)
+        self.gl_texture_buffer[:,:] = 0.0
+
+        self.gl_moments = glGenTextures(1)
+        self.gl_texture_type = {2: GL_TEXTURE_2D, 3: GL_TEXTURE_3D}.get(self.lattice.descriptor.d)
+        glBindTexture(self.gl_texture_type, self.gl_moments)
+
+        if self.gl_texture_type == GL_TEXTURE_3D:
+            glTexImage3D(self.gl_texture_type, 0, GL_RGBA, self.lattice.memory.size_x, self.lattice.memory.size_y, self.lattice.memory.size_z, 0, GL_RGBA, GL_FLOAT, self.gl_texture_buffer)
+            glTexParameteri(self.gl_texture_type, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
+            glTexParameteri(self.gl_texture_type, GL_TEXTURE_MAG_FILTER, GL_NEAREST)
+            glTexParameteri(self.gl_texture_type, GL_TEXTURE_WRAP_T,     GL_CLAMP_TO_EDGE)
+            glTexParameteri(self.gl_texture_type, GL_TEXTURE_WRAP_S,     GL_CLAMP_TO_EDGE)
+            glTexParameteri(self.gl_texture_type, GL_TEXTURE_WRAP_R,     GL_CLAMP_TO_EDGE)
+            self.cl_gl_moments  = cl.GLTexture(self.lattice.context, mf.READ_WRITE, self.gl_texture_type, 0, self.gl_moments, 3)
+        elif self.gl_texture_type == GL_TEXTURE_2D:
+            glTexImage2D(self.gl_texture_type, 0, GL_RGBA, self.lattice.memory.size_x, self.lattice.memory.size_y, 0, GL_RGBA, GL_FLOAT, self.gl_texture_buffer)
+            glTexParameteri(self.gl_texture_type, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
+            glTexParameteri(self.gl_texture_type, GL_TEXTURE_MAG_FILTER, GL_NEAREST)
+            glTexParameteri(self.gl_texture_type, GL_TEXTURE_WRAP_T,     GL_CLAMP_TO_EDGE)
+            glTexParameteri(self.gl_texture_type, GL_TEXTURE_WRAP_S,     GL_CLAMP_TO_EDGE)
+            self.cl_gl_moments  = cl.GLTexture(self.lattice.context, mf.READ_WRITE, self.gl_texture_type, 0, self.gl_moments, 2)
+
+        self.build_kernel()
+
+    def build_kernel(self):
+        program_src = Template(filename = str(Path(__file__).parent/'../template/opengl.mako')).render(
+            descriptor = self.lattice.descriptor,
+            geometry   = self.lattice.geometry,
+            memory     = self.lattice.memory,
+
+            moments_subexpr    = self.lattice.moments[0],
+            moments_assignment = self.lattice.moments[1],
+            collide_subexpr    = self.lattice.collide[0],
+            collide_assignment = self.lattice.collide[1],
+
+            float_type = self.lattice.float_type[1],
+
+            ccode = sympy.ccode
+        )
+        self.program = cl.Program(self.lattice.context, program_src).build(self.lattice.compiler_args)
+
+    def bind(self, location = GL_TEXTURE0):
+        glEnable(self.gl_texture_type)
+        glActiveTexture(location);
+        glBindTexture(self.gl_texture_type, self.gl_moments)
+
+    def collect(self):
+        cl.enqueue_acquire_gl_objects(self.lattice.queue, [self.cl_gl_moments])
+
+        if self.lattice.tick:
+            self.program.collect_gl_moments_to_texture(
+                self.lattice.queue,
+                self.lattice.grid.size(),
+                self.lattice.layout,
+                self.lattice.memory.cl_pop_b,
+                self.lattice.memory.cl_material,
+                self.cl_gl_moments)
+        else:
+            self.program.collect_gl_moments_to_texture(
+                self.lattice.queue,
+                self.lattice.grid.size(),
+                self.lattice.layout,
+                self.lattice.memory.cl_pop_a,
+                self.lattice.memory.cl_material,
+                self.cl_gl_moments)
+
diff --git a/utility/particles.py b/utility/particles.py
index a7ef777..4b770bd 100644
--- a/utility/particles.py
+++ b/utility/particles.py
@@ -3,16 +3,23 @@ mf = cl.mem_flags
 
 import numpy
 
+from mako.template import Template
+from pathlib import Path
+
 import OpenGL.GL as gl
 from OpenGL.arrays import vbo
 
 class Particles:
-    def __init__(self, context, queue, float_type, grid):
-        self.context = context
+    def __init__(self, lattice, moments, grid):
+        self.lattice = lattice
+        self.context = self.lattice.context
+        self.queue   = self.lattice.queue
+        self.float_type = self.lattice.memory.float_type
+        self.moments = moments
         self.count = len(grid)
 
-        self.np_particles = numpy.ndarray(shape=(self.count, 4), dtype=float_type)
-        self.np_init_particles = numpy.ndarray(shape=(self.count, 4), dtype=float_type)
+        self.np_particles = numpy.ndarray(shape=(self.count, 4), dtype=self.float_type)
+        self.np_init_particles = numpy.ndarray(shape=(self.count, 4), dtype=self.float_type)
 
         if len(grid[0,:]) == 2:
             self.np_particles[:,0:2] = grid
@@ -28,5 +35,36 @@ class Particles:
         self.gl_particles.bind()
         self.cl_gl_particles = cl.GLBuffer(self.context, mf.READ_WRITE, int(self.gl_particles))
 
-        self.cl_init_particles = cl.Buffer(context, mf.READ_ONLY, size=self.count * 4*numpy.float32(0).nbytes)
-        cl.enqueue_copy(queue, self.cl_init_particles, self.np_init_particles).wait();
+        self.cl_init_particles = cl.Buffer(self.context, mf.READ_ONLY, size=self.count * 4*numpy.float32(0).nbytes)
+        cl.enqueue_copy(self.queue, self.cl_init_particles, self.np_init_particles).wait();
+
+        self.build_kernel()
+
+    def build_kernel(self):
+        program_src = Template(filename = str(Path(__file__).parent/'../template/particles.mako')).render(
+            descriptor = self.lattice.descriptor,
+            geometry   = self.lattice.geometry,
+            memory     = self.lattice.memory,
+            float_type = self.float_type,
+        )
+        self.program = cl.Program(self.lattice.context, program_src).build(self.lattice.compiler_args)
+
+    def bind(self):
+        gl.glEnableClientState(gl.GL_VERTEX_ARRAY)
+        self.gl_particles.bind()
+        gl.glVertexPointer(4, gl.GL_FLOAT, 0, self.gl_particles)
+
+    def update(self, aging = False):
+        cl.enqueue_acquire_gl_objects(self.queue, [self.cl_gl_particles])
+
+        if aging:
+            age = numpy.float32(0.000006)
+        else:
+            age = numpy.float32(0.0)
+
+        self.program.update_particles(
+            self.queue, (self.count,1), None,
+            self.moments.cl_gl_moments,
+            self.lattice.memory.cl_material,
+            self.cl_gl_particles, self.cl_init_particles,
+            age)