Prototype "ink" particles visualization

4 files changed, 124 insertions, 38 deletions

diff --git a/channel_2d_gl_interop.py b/channel_2d_gl_interop.py
index 7c78af5..1f18cd5 100644
--- a/channel_2d_gl_interop.py
+++ b/channel_2d_gl_interop.py
@@ -1,7 +1,7 @@
 import numpy
 from string import Template
 
-from simulation         import Lattice, Geometry
+from simulation         import Lattice, Geometry, Particles
 from symbolic.generator import LBM
 
 import symbolic.D2Q9 as D2Q9
@@ -14,10 +14,13 @@ from OpenGL.GL import shaders
 screen_x = 1920
 screen_y = 1200
 pixels_per_cell   = 4
-updates_per_frame = 80
+updates_per_frame = 100
 
-inflow = 0.02
-relaxation_time = 0.51
+inflow = 0.01
+relaxation_time = 0.52
+
+def circle(cx, cy, r):
+    return lambda x, y: (x - cx)**2 + (y - cy)**2 < r*r
 
 def get_channel_material_map(geometry):
     return [
@@ -26,19 +29,13 @@ def get_channel_material_map(geometry):
         (lambda x, y: x == geometry.size_x-2, 4), # outflow
         (lambda x, y: y == 1,                 2), # bottom
         (lambda x, y: y == geometry.size_y-2, 2), # top
+        (lambda x, y: x > geometry.size_x//20 and x < 2*geometry.size_x//20 and y < 4*geometry.size_y//9, 2),
+        (lambda x, y: x > geometry.size_x//20 and x < 2*geometry.size_x//20 and y > 5*geometry.size_y//9, 2),
+        (circle(geometry.size_x//4,    geometry.size_y//2, 50), 2),
+        (circle(geometry.size_x//4-25, geometry.size_y//2, 50), 1),
         (lambda x, y: x == 0 or x == geometry.size_x-1 or y == 0 or y == geometry.size_y-1, 0) # ghost cells
     ]
 
-def get_obstacles(geometry):
-    ys = numpy.linspace(geometry.size_y//50, geometry.size_y-geometry.size_y//50, num = 20)
-    xs = [ 50 for i, y in enumerate(ys) ]
-    rs = [ geometry.size_x//100 for i, y in enumerate(ys) ]
-    return list(zip(xs, ys, rs))
-
-def get_obstacles_material_map(obstacles):
-    indicator = lambda ox, oy, r: (lambda x, y: (ox - x)**2 + (oy - y)**2 < r*r)
-    return [ (indicator(ox, oy, r), 2) for ox, oy, r in obstacles ]
-
 boundary = Template("""
     if ( m == 2 ) {
         u_0 = 0.0;
@@ -113,7 +110,11 @@ void main() {
         1.
     );
 
-    color = blueRedPalette(CellMoments[3] / 0.08);
+    if (CellMoments[3] >= 0.0) {
+        color = blueRedPalette(CellMoments[3] / 0.2);
+    } else {
+        color = vec3(0.0,0.0,0.0);
+    }
 }""").substitute({
     'size_x': screen_x//pixels_per_cell,
     'inflow': inflow
@@ -128,8 +129,28 @@ void main(){
     gl_FragColor = vec4(color.xyz, 0.0);
 }""", GL_FRAGMENT_SHADER)
 
+particle_shader = shaders.compileShader(Template("""
+#version 430
+
+layout (location=0) in vec4 Particles;
+
+out vec3 color;
+
+uniform mat4 projection;
+
+void main() {
+    gl_Position = projection * vec4(
+        Particles[0],
+        Particles[1],
+        0.,
+        1.
+    );
+
+    color = vec3(0.0,1.0,0.0);
+}""").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')
 
 lattice = Lattice(
@@ -143,18 +164,18 @@ lattice = Lattice(
 
 lattice.apply_material_map(
     get_channel_material_map(lattice.geometry))
-lattice.apply_material_map(
-    get_obstacles_material_map(get_obstacles(lattice.geometry)))
-
 lattice.sync_material()
 
 projection = get_projection()
 
+particles = Particles(lattice.context, lattice.memory.float_type, lattice.geometry, 100000)
+
 def on_display():
     for i in range(0,updates_per_frame):
         lattice.evolve()
+        lattice.update_gl_particles(particles)
 
-    lattice.sync_gl_moments()
+    lattice.sync()
 
     glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
 
@@ -169,6 +190,17 @@ def on_display():
     glPointSize(pixels_per_cell)
     glDrawArrays(GL_POINTS, 0, lattice.geometry.volume)
 
+    particles.gl_particles.bind()
+    glEnableClientState(GL_VERTEX_ARRAY)
+
+    shaders.glUseProgram(particle_program)
+    glUniformMatrix4fv(projection_id, 1, False, numpy.asfortranarray(projection))
+
+    glVertexPointer(4, GL_FLOAT, 0, particles.gl_particles)
+
+    glPointSize(2)
+    glDrawArrays(GL_POINTS, 0, particles.count)
+
     glDisableClientState(GL_VERTEX_ARRAY)
 
     glutSwapBuffers()
diff --git a/ldc_2d_gl_interop.py b/ldc_2d_gl_interop.py
index ec94a75..45c9bf1 100644
--- a/ldc_2d_gl_interop.py
+++ b/ldc_2d_gl_interop.py
@@ -13,7 +13,7 @@ from OpenGL.GL import shaders
 
 screen_x = 1920
 screen_y = 1200
-cells_per_pixel   = 4
+pixels_per_cell   = 4
 updates_per_frame = 200
 
 lid_speed = 0.1
@@ -41,7 +41,7 @@ boundary = Template("""
 })
 
 def get_projection():
-    scale = numpy.diag([(2.0*cells_per_pixel)/screen_x, (2.0*cells_per_pixel)/screen_y, 1.0, 1.0])
+    scale = numpy.diag([(2.0*pixels_per_cell)/screen_x, (2.0*pixels_per_cell)/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;
@@ -100,7 +100,7 @@ void main() {
 
     color = blueRedPalette(CellMoments[3] / $lid_speed);
 }""").substitute({
-    'size_x'   : screen_x//cells_per_pixel,
+    'size_x'   : screen_x//pixels_per_cell,
     'lid_speed': lid_speed
 }), GL_VERTEX_SHADER)
 
@@ -119,7 +119,7 @@ projection_id = shaders.glGetUniformLocation(shader_program, 'projection')
 
 lattice = Lattice(
     descriptor   = D2Q9,
-    geometry     = Geometry(screen_x//cells_per_pixel, screen_y//cells_per_pixel),
+    geometry     = Geometry(screen_x//pixels_per_cell, screen_y//pixels_per_cell),
     moments      = lbm.moments(optimize = True),
     collide      = lbm.bgk(f_eq = lbm.equilibrium(), tau = relaxation_time),
     boundary_src = boundary,
@@ -136,7 +136,7 @@ def on_display():
     for i in range(0,updates_per_frame):
         lattice.evolve()
 
-    lattice.sync_gl_moments()
+    lattice.collect_gl_moments()
 
     glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
 
@@ -148,7 +148,7 @@ def on_display():
 
     glVertexPointer(4, GL_FLOAT, 0, lattice.memory.gl_moments)
 
-    glPointSize(cells_per_pixel)
+    glPointSize(pixels_per_cell)
     glDrawArrays(GL_POINTS, 0, lattice.geometry.volume)
 
     glDisableClientState(GL_VERTEX_ARRAY)
diff --git a/simulation.py b/simulation.py
index 1f9a6c1..0f7e904 100644
--- a/simulation.py
+++ b/simulation.py
@@ -37,6 +37,7 @@ class Geometry:
             return (self.size_x-2, self.size_y-2, self.size_z-2)
 
 
+
 def pad(n, m):
     return (n // m + min(1,n % m)) * m
 
@@ -109,6 +110,24 @@ class Memory:
     def cells(self):
         return ndindex(self.size(), order='F')
 
+class Particles:
+    def __init__(self, context, float_type, geometry, n):
+        self.context = context
+        self.count = n
+
+        self.np_particles = numpy.ndarray(shape=(self.count, 4), dtype=float_type)
+
+        self.np_particles[:,0:2] = numpy.mgrid[
+            geometry.size_x//20:2*geometry.size_x//20:100j,
+            4*geometry.size_y//9:5*geometry.size_y//9:n/100j
+        ].reshape(2,-1).T
+        self.np_particles[:,2] = 0.0
+        self.np_particles[:,3] = 0.0
+
+        self.gl_particles = vbo.VBO(data=self.np_particles, usage=gl.GL_DYNAMIC_DRAW, target=gl.GL_ARRAY_BUFFER)
+        self.gl_particles.bind()
+        self.cl_gl_particles = cl.GLBuffer(self.context, mf.READ_WRITE, int(self.gl_particles))
+
 class Lattice:
     def __init__(self,
         descriptor, geometry, moments, collide,
@@ -225,15 +244,25 @@ class Lattice:
 
         return moments
 
-    def sync_gl_moments(self):
+    def collect_gl_moments(self):
         cl.enqueue_acquire_gl_objects(self.queue, [self.memory.cl_gl_moments])
 
         if self.tick:
             self.program.collect_gl_moments(
-                self.queue, self.grid.size(), self.layout, self.memory.cl_pop_b, self.memory.cl_gl_moments)
+                self.queue, self.grid.size(), self.layout, self.memory.cl_pop_b, self.memory.cl_material, self.memory.cl_gl_moments)
+        else:
+            self.program.collect_gl_moments(
+                self.queue, self.grid.size(), self.layout, self.memory.cl_pop_a, self.memory.cl_material, self.memory.cl_gl_moments)
+
+    def update_gl_particles(self, particles):
+        cl.enqueue_acquire_gl_objects(self.queue, [self.memory.cl_gl_moments, particles.cl_gl_particles])
+
+        if self.tick:
+            self.program.collect_gl_moments(
+                self.queue, self.grid.size(), self.layout, self.memory.cl_pop_b, self.memory.cl_material, self.memory.cl_gl_moments)
         else:
             self.program.collect_gl_moments(
-                self.queue, self.grid.size(), self.layout, self.memory.cl_pop_a, self.memory.cl_gl_moments)
+                self.queue, self.grid.size(), self.layout, self.memory.cl_pop_a, self.memory.cl_material, self.memory.cl_gl_moments)
 
-        #cl.enqueue_release_gl_objects(self.queue, [self.cl_gl_moments])
-        self.sync()
+        self.program.update_particles(
+            self.queue, (particles.count,1), None, self.memory.cl_gl_moments, particles.cl_gl_particles)
diff --git a/template/kernel.mako b/template/kernel.mako
index 00a4345..f6bac7f 100644
--- a/template/kernel.mako
+++ b/template/kernel.mako
@@ -100,6 +100,7 @@ __kernel void collect_moments(__global __read_only  ${float_type}* f,
 }
 
 __kernel void collect_gl_moments(__global __read_only  ${float_type}* f,
+                                 __global __read_only  int* material,
                                  __global __write_only float4* moments)
 {
     const unsigned int gid = ${gid()};
@@ -116,17 +117,41 @@ __kernel void collect_gl_moments(__global __read_only  ${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];
+    }
 
     moments[gid] = data;
 }
+
+__kernel void update_particles(__global __read_only  float4* moments,
+                               __global __write_only float4* particles)
+{
+  const unsigned int pid = get_global_id(0);
+
+  float4 particle = particles[pid];
+
+  const unsigned int gid = floor(particle.y)*${memory.size_x} + floor(particle.x);
+
+  float4 moment = moments[gid];
+
+  particle.x += moment.y;
+  particle.y += moment.z;
+
+  particles[pid] = particle;
+}