Gather interop moments in a more generic manner

i.e. return unshifted moments in a implicitly ordered float4 array.
Cell positions are reconstructed by a vertex shaded analogously to
how it is done in compustream.

3 files changed, 71 insertions, 23 deletions

diff --git a/ldc_2d_gl_interop.py b/ldc_2d_gl_interop.py
index 6ff10c7..a10664d 100644
--- a/ldc_2d_gl_interop.py
+++ b/ldc_2d_gl_interop.py
@@ -1,5 +1,6 @@
 import numpy
 import time
+from string import Template
 
 from simulation         import Lattice, Geometry
 from symbolic.generator import LBM
@@ -9,6 +10,8 @@ import symbolic.D2Q9 as D2Q9
 from OpenGL.GL   import *
 from OpenGL.GLUT import *
 
+from OpenGL.GL import shaders
+
 screen_x = 1920
 screen_y = 1200
 
@@ -31,6 +34,12 @@ boundary = """
     }
 """
 
+def get_projection():
+    scale = numpy.diag([8.0/screen_x, 8.0/screen_y, 1.0, 1.0])
+    translation        = numpy.matrix(numpy.identity(4))
+    translation[3,0:3] = [-1.0, -1.0, 0.0]
+    return scale * translation;
+
 def glut_window(fullscreen = False):
     glutInit(sys.argv)
     glutInitDisplayMode(GLUT_RGBA | GLUT_DOUBLE | GLUT_DEPTH)
@@ -42,48 +51,80 @@ def glut_window(fullscreen = False):
         glutInitWindowPosition(0, 0)
         window = glutCreateWindow("LBM")
 
-    glMatrixMode(GL_PROJECTION)
-    glLoadIdentity()
-
-    glOrtho(
-        0, screen_x,
-        0, screen_y,
-        0.1, 100.0
-    )
-
     return window
 
 lbm = LBM(D2Q9)
 
 window = glut_window(fullscreen = False)
 
+vertex_shader = shaders.compileShader(Template("""
+#version 430
+
+layout (location=0) in vec4 CellMoments;
+
+uniform mat4 projection;
+
+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 + 500.*CellMoments[1],
+        idx.y + 500.*CellMoments[2],
+        0.,
+        1.
+    );
+}""").substitute({'size_x': screen_x//4}), GL_VERTEX_SHADER)
+
+fragment_shader = shaders.compileShader("""
+#version 120
+void main(){
+    gl_FragColor = vec4(1,1,1,1);
+}""", GL_FRAGMENT_SHADER)
+
+
+
+shader_program = shaders.compileProgram(vertex_shader, fragment_shader)
+projection_id = shaders.glGetUniformLocation(shader_program, 'projection')
+
 lattice = Lattice(
     descriptor   = D2Q9,
     geometry     = Geometry(screen_x//4, screen_y//4),
-    moments      = lbm.moments(optimize = False),
-    collide      = lbm.bgk(f_eq = lbm.equilibrium(), tau = 0.52),
+    moments      = lbm.moments(optimize = True),
+    collide      = lbm.bgk(f_eq = lbm.equilibrium(), tau = 0.515),
     boundary_src = boundary,
     opengl       = True
 )
 
 lattice.setup_geometry(cavity)
 
+projection = get_projection()
+
 def on_display():
-    for i in range(0,100):
+    for i in range(0,60):
         lattice.evolve()
 
     lattice.sync_gl_moments()
 
     glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
-    glMatrixMode(GL_MODELVIEW)
-    glLoadIdentity()
-
-    glTranslatef(0., 0., -1.)
 
     lattice.gl_moments.bind()
     glEnableClientState(GL_VERTEX_ARRAY)
+
+    shaders.glUseProgram(shader_program)
+    glUniformMatrix4fv(projection_id, 1, False, numpy.asfortranarray(projection))
+
     glVertexPointer(4, GL_FLOAT, 0, lattice.gl_moments)
+
     glDrawArrays(GL_POINTS, 0, lattice.geometry.volume)
+
     glDisableClientState(GL_VERTEX_ARRAY)
 
     glutSwapBuffers()
diff --git a/symbolic/generator.py b/symbolic/generator.py
index 853bbf9..f94031f 100644
--- a/symbolic/generator.py
+++ b/symbolic/generator.py
@@ -47,7 +47,7 @@ class LBM:
         exprs = [ self.f_curr[i] + 1/tau * (f_eq_i - self.f_curr[i]) for i, f_eq_i in enumerate(f_eq) ]
 
         if optimize:
-            helper, f = cse(exprs, optimizations=optimizations.custom)
-            return (helper, assign(self.f_next, f))
+            subexprs, f = cse(exprs, optimizations=optimizations.custom)
+            return (subexprs, assign(self.f_next, f))
         else:
-            return ([], assign(self.f_next, f))
+            return ([], assign(self.f_next, exprs))
diff --git a/template/kernel.mako b/template/kernel.mako
index d9b1c80..417851a 100644
--- a/template/kernel.mako
+++ b/template/kernel.mako
@@ -115,10 +115,17 @@ __kernel void collect_gl_moments(__global __read_only  float* f,
 
     float4 data;
 
-    data.x = 4.0*((float)(get_global_id(0))) + ${ccode(2000*moments_assignment[1].rhs)};
-    data.y = 4.0*((float)(get_global_id(1))) + ${ccode(2000*moments_assignment[2].rhs)};
-    data.z = 0.0;
-    data.w = 1.0;
+% 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);
+% 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)};
+% endif
 
     moments[gid] = data;
 }