Reset stuck particles to starting position

4 files changed, 45 insertions, 28 deletions

diff --git a/channel_2d_gl_interop.py b/channel_2d_gl_interop.py
index 1f18cd5..69af67b 100644
--- a/channel_2d_gl_interop.py
+++ b/channel_2d_gl_interop.py
@@ -1,7 +1,8 @@
 import numpy
 from string import Template
 
-from simulation         import Lattice, Geometry, Particles
+from simulation         import Lattice, Geometry
+from utility.particles  import Particles
 from symbolic.generator import LBM
 
 import symbolic.D2Q9 as D2Q9
@@ -110,7 +111,7 @@ void main() {
         1.
     );
 
-    if (CellMoments[3] >= 0.0) {
+    if (CellMoments[3] > 0.0) {
         color = blueRedPalette(CellMoments[3] / 0.2);
     } else {
         color = vec3(0.0,0.0,0.0);
@@ -146,7 +147,7 @@ void main() {
         1.
     );
 
-    color = vec3(0.0,1.0,0.0);
+    color = vec3(0.0,0.8,0.0);
 }""").substitute({}), GL_VERTEX_SHADER)
 
 shader_program = shaders.compileProgram(vertex_shader, fragment_shader)
@@ -168,7 +169,13 @@ lattice.sync_material()
 
 projection = get_projection()
 
-particles = Particles(lattice.context, lattice.memory.float_type, lattice.geometry, 100000)
+particles = Particles(
+    lattice.context,
+    lattice.memory.float_type,
+    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
+    ].reshape(2,-1).T)
 
 def on_display():
     for i in range(0,updates_per_frame):
@@ -187,6 +194,7 @@ def on_display():
 
     glVertexPointer(4, GL_FLOAT, 0, lattice.memory.gl_moments)
 
+    glDisable(GL_POINT_SMOOTH)
     glPointSize(pixels_per_cell)
     glDrawArrays(GL_POINTS, 0, lattice.geometry.volume)
 
@@ -198,7 +206,8 @@ def on_display():
 
     glVertexPointer(4, GL_FLOAT, 0, particles.gl_particles)
 
-    glPointSize(2)
+    glEnable(GL_POINT_SMOOTH)
+    glPointSize(1)
     glDrawArrays(GL_POINTS, 0, particles.count)
 
     glDisableClientState(GL_VERTEX_ARRAY)
diff --git a/simulation.py b/simulation.py
index 0f7e904..b9df3f0 100644
--- a/simulation.py
+++ b/simulation.py
@@ -36,8 +36,6 @@ class Geometry:
         else:
             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
 
@@ -110,24 +108,6 @@ 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,
@@ -265,4 +245,4 @@ class Lattice:
                 self.queue, self.grid.size(), self.layout, self.memory.cl_pop_a, self.memory.cl_material, self.memory.cl_gl_moments)
 
         self.program.update_particles(
-            self.queue, (particles.count,1), None, self.memory.cl_gl_moments, particles.cl_gl_particles)
+            self.queue, (particles.count,1), None, self.memory.cl_gl_moments, self.memory.cl_material, particles.cl_gl_particles)
diff --git a/template/kernel.mako b/template/kernel.mako
index f6bac7f..1ee39cd 100644
--- a/template/kernel.mako
+++ b/template/kernel.mako
@@ -140,6 +140,7 @@ __kernel void collect_gl_moments(__global __read_only  ${float_type}* f,
 }
 
 __kernel void update_particles(__global __read_only  float4* moments,
+                               __global __read_only  int*    material,
                                __global __write_only float4* particles)
 {
   const unsigned int pid = get_global_id(0);
@@ -150,8 +151,13 @@ __kernel void update_particles(__global __read_only  float4* moments,
 
   float4 moment = moments[gid];
 
-  particle.x += moment.y;
-  particle.y += moment.z;
+  if (material[gid] == 1) {
+    particle.x += moment.y;
+    particle.y += moment.z;
+  } else {
+    particle.x = particle.z;
+    particle.y = particle.w;
+  }
 
   particles[pid] = particle;
 }
diff --git a/utility/particles.py b/utility/particles.py
new file mode 100644
index 0000000..b838790
--- /dev/null
+++ b/utility/particles.py
@@ -0,0 +1,22 @@
+import pyopencl as cl
+mf = cl.mem_flags
+
+import numpy
+
+import OpenGL.GL as gl
+from OpenGL.arrays import vbo
+
+class Particles:
+    def __init__(self, context, float_type, grid):
+        self.context = context
+        self.count = len(grid)
+
+        self.np_particles = numpy.ndarray(shape=(self.count, 4), dtype=float_type)
+
+        self.np_particles[:,0:2] = grid
+        self.np_particles[:,2:4] = self.np_particles[:,0:2]
+
+        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))
+