Use lattice-internal moments buffer for particles

5 files changed, 77 insertions, 58 deletions

diff --git a/channel_2d_gl_interop.py b/channel_2d_gl_interop.py
index cae2f5b..eaf2c46 100644
--- a/channel_2d_gl_interop.py
+++ b/channel_2d_gl_interop.py
@@ -2,7 +2,7 @@ import numpy
 from string import Template
 
 from simulation         import Lattice, Geometry
-from utility.opengl     import MomentsVertexBuffer
+from utility.opengl     import MomentsTexture
 from utility.particles  import Particles
 from symbolic.generator import LBM
 
@@ -85,15 +85,32 @@ lbm = LBM(D2Q9)
 
 window = glut_window(fullscreen = False)
 
-vertex_shader = shaders.compileShader(Template("""
+vertex_shader = shaders.compileShader("""
 #version 430
 
-layout (location=0) in vec4 CellMoments;
-
-out vec3 color;
+layout (location=0) in vec4 vertex;
+                   out vec2 frag_pos;
 
 uniform mat4 projection;
 
+void main() {
+    gl_Position = projection * vertex;
+    frag_pos    = vertex.xy;
+}""", GL_VERTEX_SHADER)
+
+fragment_shader = shaders.compileShader(Template("""
+#version 430
+
+in vec2 frag_pos;
+
+uniform sampler2D moments;
+
+out vec4 result;
+
+vec2 unit(vec2 v) {
+    return vec2(v[0] / $size_x, v[1] / $size_y);
+}
+
 vec3 blueRedPalette(float x) {
     return mix(
         vec3(0.0, 0.0, 1.0),
@@ -102,35 +119,23 @@ vec3 blueRedPalette(float x) {
     );
 }
 
-vec2 fluidVertexAtIndex(uint i) {
-    const float y = floor(float(i) / $size_x);
-    return vec2(
-        i - $size_x*y,
-        y
-    );
-}
-
-void main() {
-    const vec2 idx = fluidVertexAtIndex(gl_VertexID);
-
-    gl_Position = projection * vec4(
-        idx.x,
-        idx.y,
-        0.,
-        1.
-    );
-
-    if (CellMoments[3] > 0.0) {
-        color = blueRedPalette(CellMoments[3] / 0.2);
+void main(){
+    const vec2 sample_pos = unit(frag_pos);
+    const vec4 data = texture(moments, sample_pos);
+    if (data.w < 0.0) {
+        result.a   = 1.0;
+        result.rgb = vec3(0.4);
     } else {
-        color = vec3(0.4,0.4,0.4);
+        result.a   = 1.0;
+        result.rgb = blueRedPalette(data[3] / 0.2);
     }
-}""").substitute({
-    'size_x': lattice_x,
-    'inflow': inflow
-}), GL_VERTEX_SHADER)
+}
+""").substitute({
+    "size_x": lattice_x,
+    "size_y": lattice_y
+}), GL_FRAGMENT_SHADER)
 
-fragment_shader = shaders.compileShader("""
+particle_fragment_shader = shaders.compileShader("""
 #version 430
 
 in vec3 color;
@@ -139,7 +144,7 @@ void main(){
     gl_FragColor = vec4(color.xyz, 0.0);
 }""", GL_FRAGMENT_SHADER)
 
-particle_shader = shaders.compileShader(Template("""
+particle_vertex_shader = shaders.compileShader(Template("""
 #version 430
 
 layout (location=0) in vec4 Particles;
@@ -160,9 +165,11 @@ void main() {
 }""").substitute({}), GL_VERTEX_SHADER)
 
 shader_program = shaders.compileProgram(vertex_shader, fragment_shader)
-particle_program = shaders.compileProgram(particle_shader, fragment_shader)
 projection_id = shaders.glGetUniformLocation(shader_program, 'projection')
 
+particle_program = shaders.compileProgram(particle_vertex_shader, particle_fragment_shader)
+particle_projection_id = shaders.glGetUniformLocation(particle_program, 'projection')
+
 lattice = Lattice(
     descriptor   = D2Q9,
     geometry     = Geometry(lattice_x, lattice_y),
@@ -176,11 +183,10 @@ lattice.apply_material_map(
     get_channel_material_map(lattice.geometry))
 lattice.sync_material()
 
-moments_vbo = MomentsVertexBuffer(lattice)
+moments_texture = MomentsTexture(lattice)
 
 particles = Particles(
     lattice,
-    moments_vbo,
     numpy.mgrid[
         lattice.geometry.size_x//20:2*lattice.geometry.size_x//20:100j,
         4*lattice.geometry.size_y//9:5*lattice.geometry.size_y//9:100000/100j
@@ -190,7 +196,8 @@ def on_display():
     for i in range(0,updates_per_frame):
         lattice.evolve()
 
-    moments_vbo.collect()
+    lattice.update_moments()
+    moments_texture.collect()
 
     for i in range(0,updates_per_frame):
         particles.update(aging = False)
@@ -200,13 +207,18 @@ def on_display():
     glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
 
     shaders.glUseProgram(shader_program)
-    glUniformMatrix4fv(projection_id, 1, False, numpy.ascontiguousarray(projection))
-    moments_vbo.bind()
-    glPointSize(point_size)
-    glDrawArrays(GL_POINTS, 0, lattice.geometry.volume)
+    glUniformMatrix4fv(projection_id, 1, False, numpy.asfortranarray(projection))
+    moments_texture.bind()
+
+    glBegin(GL_POLYGON)
+    glVertex(0,0,0)
+    glVertex(lattice.geometry.size_x,0,0)
+    glVertex(lattice.geometry.size_x,lattice.geometry.size_y,0)
+    glVertex(0,lattice.geometry.size_y,0)
+    glEnd()
 
     shaders.glUseProgram(particle_program)
-    glUniformMatrix4fv(projection_id, 1, False, numpy.asfortranarray(projection))
+    glUniformMatrix4fv(particle_projection_id, 1, False, numpy.asfortranarray(projection))
     particles.bind()
     glPointSize(point_size)
     glDrawArrays(GL_POINTS, 0, particles.count)
diff --git a/template/opengl.mako b/template/opengl.mako
index 47264e8..2029623 100644
--- a/template/opengl.mako
+++ b/template/opengl.mako
@@ -65,6 +65,7 @@ def neighbor_offset(c_i):
 %>
 
 __kernel void collect_gl_moments_to_texture(__global ${float_type}* f,
+                                            __global int* material,
 % if descriptor.d == 2:
                                             __write_only image2d_t moments)
 % elif descriptor.d == 3:
@@ -85,17 +86,24 @@ __kernel void collect_gl_moments_to_texture(__global ${float_type}* f,
 
     float4 data;
 
+    if (material[gid] == 1) {
 % if descriptor.d == 2:
-    data.x = ${ccode(moments_assignment[0].rhs)};
-    data.y = ${ccode(moments_assignment[1].rhs)};
-    data.z = ${ccode(moments_assignment[2].rhs)};
-    data.w = sqrt(data.y*data.y + data.z*data.z);
+      data.x = ${ccode(moments_assignment[0].rhs)};
+      data.y = ${ccode(moments_assignment[1].rhs)};
+      data.z = ${ccode(moments_assignment[2].rhs)};
+      data.w = sqrt(data.y*data.y + data.z*data.z);
 % elif descriptor.d == 3:
-    data.x = ${ccode(moments_assignment[0].rhs)};
-    data.y = ${ccode(moments_assignment[1].rhs)};
-    data.z = ${ccode(moments_assignment[2].rhs)};
-    data.w = ${ccode(moments_assignment[3].rhs)};
+      data.x = ${ccode(moments_assignment[0].rhs)};
+      data.y = ${ccode(moments_assignment[1].rhs)};
+      data.z = ${ccode(moments_assignment[2].rhs)};
+      data.w = ${ccode(moments_assignment[3].rhs)};
 % endif
+    } else {
+      data.x = 0.0;
+      data.y = 0.0;
+      data.z = 0.0;
+      data.w = -material[gid];
+    }
 
     write_imagef(moments, ${moments_cell()}, data);
 }
diff --git a/template/particles.mako b/template/particles.mako
index 48191d9..0b9a67e 100644
--- a/template/particles.mako
+++ b/template/particles.mako
@@ -1,4 +1,4 @@
-__kernel void update_particles(__global float4* moments,
+__kernel void update_particles(__global float*  moments,
                                __global int*    material,
                                __global float4* particles,
                                __global float4* init_particles,
@@ -14,13 +14,11 @@ __kernel void update_particles(__global float4* moments,
   const unsigned int gid = floor(particle.z)*${memory.size_x*memory.size_y} + floor(particle.y)*${memory.size_x} + floor(particle.x);
 % endif
 
-  const float4 moment = moments[gid];
-
   if (material[gid] == 1 && particle.w < 1.0) {
-    particle.x += moment.y;
-    particle.y += moment.z;
+    particle.x += moments[${1*memory.volume}+gid];
+    particle.y += moments[${2*memory.volume}+gid];
 % if descriptor.d == 3:
-    particle.z += moment.w;
+    particle.z += moments[${3*memory.volume}+gid];
 % endif
     particle.w += min(particle.x, particle.y) * aging;
   } else {
diff --git a/utility/opengl.py b/utility/opengl.py
index e201e29..1037cbd 100644
--- a/utility/opengl.py
+++ b/utility/opengl.py
@@ -123,6 +123,7 @@ class MomentsTexture:
                 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(
@@ -130,5 +131,6 @@ class MomentsTexture:
                 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 4b770bd..75d2245 100644
--- a/utility/particles.py
+++ b/utility/particles.py
@@ -10,12 +10,11 @@ import OpenGL.GL as gl
 from OpenGL.arrays import vbo
 
 class Particles:
-    def __init__(self, lattice, moments, grid):
+    def __init__(self, lattice, 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=self.float_type)
@@ -64,7 +63,7 @@ class Particles:
 
         self.program.update_particles(
             self.queue, (self.count,1), None,
-            self.moments.cl_gl_moments,
+            self.lattice.memory.cl_moments,
             self.lattice.memory.cl_material,
             self.cl_gl_particles, self.cl_init_particles,
             age)