Use OpenCL buffer to access moments in streamline impl

5 files changed, 13 insertions, 17 deletions

diff --git a/channel_2d_streamlines_gl_interop.py b/channel_2d_streamlines_gl_interop.py
index c1aa1e8..7c370a2 100644
--- a/channel_2d_streamlines_gl_interop.py
+++ b/channel_2d_streamlines_gl_interop.py
@@ -137,16 +137,15 @@ lattice.apply_material_map(
     get_channel_material_map(lattice.geometry))
 lattice.sync_material()
 
-moments_vbo = MomentsVertexBuffer(lattice)
 streamline_texture = Streamlines(
-    lattice, moments_vbo,
+    lattice,
     list(map(lambda y: [2, y*lattice.geometry.size_y//48], range(1,48))))
 
 def on_display():
     for i in range(0,updates_per_frame):
         lattice.evolve()
 
-    moments_vbo.collect()
+    lattice.update_moments()
     streamline_texture.update()
 
     glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
diff --git a/simulation.py b/simulation.py
index 5ebe9f6..e2174ee 100644
--- a/simulation.py
+++ b/simulation.py
@@ -261,9 +261,7 @@ class Lattice:
     def sync(self):
         self.queue.finish()
 
-    def get_moments(self):
-        moments = numpy.ndarray(shape=(self.descriptor.d+1, self.memory.volume), dtype=self.float_type[0])
-
+    def update_moments(self):
         if self.tick:
             self.program.collect_moments(
                 self.queue, self.grid.size(), self.layout, self.memory.cl_pop_b, self.memory.cl_moments)
@@ -271,6 +269,8 @@ class Lattice:
             self.program.collect_moments(
                 self.queue, self.grid.size(), self.layout, self.memory.cl_pop_a, self.memory.cl_moments)
 
+    def get_moments(self):
+        moments = numpy.ndarray(shape=(self.descriptor.d+1, self.memory.volume), dtype=self.float_type[0])
+        self.update_moments()
         cl.enqueue_copy(self.queue, moments, self.memory.cl_moments).wait();
-
         return moments
diff --git a/template/kernel.mako b/template/kernel.mako
index dd8eaee..2e26124 100644
--- a/template/kernel.mako
+++ b/template/kernel.mako
@@ -105,4 +105,4 @@ __kernel void collect_moments(__global ${float_type}* f,
 % for i, expr in enumerate(moments_assignment):
     moments[${pop_offset(i)} + gid] = ${ccode(expr.rhs)};
 % endfor
-}
+}
\ No newline at end of file
diff --git a/template/streamline.mako b/template/streamline.mako
index 73c80e2..adf5efd 100644
--- a/template/streamline.mako
+++ b/template/streamline.mako
@@ -31,7 +31,7 @@ float3 blueRedPalette(float x) {
     );
 }
 
-__kernel void draw_streamline(__global float4* moments,
+__kernel void draw_streamline(__global float*  moments,
                               __global int*    material,
                               __global float2* origins,
                               __read_write image2d_t streamlines)
@@ -40,14 +40,12 @@ __kernel void draw_streamline(__global float4* moments,
 
     for (int i = 0; i < ${2*memory.size_x}; ++i) {
     const unsigned int gid = round(particle.y)*${memory.size_x} + round(particle.x);
-        const float4 moment = moments[gid];
-
         if (material[gid] != 1) {
             break;
         }
 
-        particle.x += 0.5 * moment.y / 0.01;
-        particle.y += 0.5 * moment.z / 0.01;
+        particle.x += 0.5 * moments[${1*memory.volume}+gid] / 0.01;
+        particle.y += 0.5 * moments[${2*memory.volume}+gid] / 0.01;
 
         const int2 pos = (int2)(round(particle.x), round(particle.y));
 
@@ -56,4 +54,4 @@ __kernel void draw_streamline(__global float4* moments,
 
         write_imagef(streamlines, pos, color);
     }
-}
+}
\ No newline at end of file
diff --git a/utility/streamline.py b/utility/streamline.py
index 5111fc2..9f0fb55 100644
--- a/utility/streamline.py
+++ b/utility/streamline.py
@@ -10,12 +10,11 @@ from OpenGL.GL import *
 from OpenGL.arrays import vbo
 
 class Streamlines:
-    def __init__(self, lattice, moments, origins):
+    def __init__(self, lattice, origins):
         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(origins)
         self.np_origins = numpy.ndarray(shape=(self.count, 2), dtype=self.float_type)
@@ -66,7 +65,7 @@ class Streamlines:
 
         self.program.draw_streamline(
             self.queue, (self.count,1), None,
-            self.moments.cl_gl_moments,
+            self.lattice.memory.cl_moments,
             self.lattice.memory.cl_material,
             self.cl_origins,
             self.cl_gl_streamlines)