Extract indicators, drawing of geometric primitives

6 files changed, 215 insertions, 109 deletions

diff --git a/channel_3d_gl_interop.py b/channel_3d_gl_interop.py
index a023bc3..a9633c4 100644
--- a/channel_3d_gl_interop.py
+++ b/channel_3d_gl_interop.py
@@ -14,38 +14,41 @@ from OpenGL.GL import shaders
 
 from pyrr import matrix44, quaternion
 
+from geometry.sphere   import Sphere
+from geometry.box      import Box
+from geometry.cylinder import Cylinder
+
 lattice_x = 256
 lattice_y = 64
 lattice_z = 64
 
-cylinder_r = 10
-
 updates_per_frame = 10
-particle_count = 50000
+particle_count = 100000
 
-lid_speed = 0.05
+lid_speed = 0.02
 relaxation_time = 0.51
 
-def circle(cx, cy, cz, r):
-    return lambda x, y, z: (x - cx)**2 + (y - cy)**2 + (z - cz)**2 < r*r
-
-def cylinder(cx, cz, r):
-    return lambda x, y, z: (x - cx)**2 + (z - cz)**2 < r*r
-
-def get_cavity_material_map(geometry):
+def get_cavity_material_map(g):
     return [
-        (lambda x, y, z: x > 0 and x < geometry.size_x-1 and
-                         y > 0 and y < geometry.size_y-1 and
-                         z > 0 and z < geometry.size_z-1,            1), # bulk fluid
-        (lambda x, y, z: x == 1 or x == geometry.size_x-2 or
-                         y == 1 or y == geometry.size_y-2 or
-                         z == 1 or z == geometry.size_z-2,           2), # walls
+        (lambda x, y, z: x > 0 and x < g.size_x-1 and
+                         y > 0 and y < g.size_y-1 and
+                         z > 0 and z < g.size_z-1,            1), # bulk fluid
+        (lambda x, y, z: x == 1 or x == g.size_x-2 or
+                         y == 1 or y == g.size_y-2 or
+                         z == 1 or z == g.size_z-2,           2), # walls
         (lambda x, y, z: x == 1,                                     3), # inflow
-        (lambda x, y, z: x == geometry.size_x-2,                     4), # outflow
-        (cylinder(geometry.size_x//6, geometry.size_z//2, cylinder_r), 5), # obstacle
-        (lambda x, y, z: x == 0 or x == geometry.size_x-1 or
-                         y == 0 or y == geometry.size_y-1 or
-                         z == 0 or z == geometry.size_z-1,           0)  # ghost cells
+        (lambda x, y, z: x == g.size_x-2,                     4), # outflow
+
+        (Box(g.size_x//10, 1.5*g.size_x//10, 0,             2*g.size_y//5, 0, g.size_z), 5),
+        (Box(g.size_x//10, 1.5*g.size_x//10, 3*g.size_y//5, g.size_y,      0, g.size_z), 5),
+
+        (Sphere(g.size_x//3, g.size_y//2, g.size_z//2, 16), 5),
+        (Cylinder(g.size_x//3, 0, g.size_z//2, 5, l = g.size_y), 5),
+        (Cylinder(g.size_x//3, g.size_y//2, 0, 5, h = g.size_z), 5),
+
+        (lambda x, y, z: x == 0 or x == g.size_x-1 or
+                         y == 0 or y == g.size_y-1 or
+                         z == 0 or z == g.size_z-1,           0)  # ghost cells
     ]
 
 boundary = Template("""
@@ -120,10 +123,7 @@ particle_shader = shaders.compileShader("""
 #version 430
 
 layout (location=0) in vec4 particles;
-
-out VS_OUT {
-    vec3 color;
-} vs_out;
+                   out vec3 color;
 
 uniform mat4 projection;
 uniform mat4 rotation;
@@ -144,40 +144,9 @@ void main() {
         1.
     );
 
-    vs_out.color = fire(1.0-particles[3]);
+    color = fire(1.0-particles[3]);
 }""", GL_VERTEX_SHADER)
 
-geometry_shader = shaders.compileShader("""
-#version 430
-
-layout (points) in;
-layout (triangle_strip, max_vertices=4) out;
-
-in VS_OUT {
-    vec3 color;
-} gs_in[];
-
-out vec3 color;
-
-void emitSquareAt(vec4 position) {
-    const float size = 0.5;
-
-    gl_Position = position + vec4(-size, -size, 0.0, 0.0);
-    EmitVertex();
-    gl_Position = position + vec4( size, -size, 0.0, 0.0);
-    EmitVertex();
-    gl_Position = position + vec4(-size,  size, 0.0, 0.0);
-    EmitVertex();
-    gl_Position = position + vec4( size,  size, 0.0, 0.0);
-    EmitVertex();
-}
-
-void main() {
-    color = gs_in[0].color;
-    emitSquareAt(gl_in[0].gl_Position);
-    EndPrimitive();
-}""", GL_GEOMETRY_SHADER)
-
 vertex_shader = shaders.compileShader(Template("""
 #version 430
 
@@ -248,7 +217,7 @@ void main(){
     "size_z": lattice_z
 }), GL_FRAGMENT_SHADER)
 
-particle_program = shaders.compileProgram(particle_shader, geometry_shader, fragment_shader)
+particle_program = shaders.compileProgram(particle_shader, fragment_shader)
 projection_id = shaders.glGetUniformLocation(particle_program, 'projection')
 rotation_id   = shaders.glGetUniformLocation(particle_program, 'rotation')
 
@@ -264,8 +233,9 @@ lattice = Lattice(
     opengl       = True
 )
 
-lattice.apply_material_map(
-    get_cavity_material_map(lattice.geometry))
+material_map = get_cavity_material_map(lattice.geometry)
+primitives   = list(map(lambda material: material[0], filter(lambda material: not callable(material[0]), material_map)))
+lattice.apply_material_map(material_map)
 lattice.sync_material()
 
 particles = Particles(
@@ -274,54 +244,14 @@ particles = Particles(
     lattice.memory.float_type,
     numpy.mgrid[
         2*lattice.geometry.size_x//100:4*lattice.geometry.size_x//100:particle_count/10000j,
-        lattice.geometry.size_y//10:9*lattice.geometry.size_y//10:100j,
-        lattice.geometry.size_z//10:9*lattice.geometry.size_z//10:100j,
+        lattice.geometry.size_y//16:15*lattice.geometry.size_y//16:100j,
+        lattice.geometry.size_z//16:15*lattice.geometry.size_z//16:100j,
     ].reshape(3,-1).T)
 
 rotation = Rotation([-lattice_x/2, -lattice_y/2, -lattice_z/2])
 
 cube_vertices, cube_edges = lattice.geometry.wireframe()
 
-def draw_cylinder(cx, cz, height, radius, num_slices):
-    r = radius
-    h = height
-    n = float(num_slices)
-
-    circle_pts = []
-    for i in range(int(n) + 1):
-        angle = 2 * numpy.pi * (i/n)
-        x = cx + r * numpy.cos(angle)
-        z = cz + r * numpy.sin(angle)
-        pt = (x, z)
-        circle_pts.append(pt)
-
-    glBegin(GL_TRIANGLE_FAN)
-    glNormal(0,-1, 0)
-    glVertex(cx, 0, cz)
-    for (x, z) in circle_pts:
-        y = 0
-        glNormal(0,-1, 0)
-        glVertex(x, y, z)
-    glEnd()
-
-    glBegin(GL_TRIANGLE_FAN)
-    glNormal(0, 1, 0)
-    glVertex(cx, h, cz)
-    for (x, z) in circle_pts:
-        y = h
-        glNormal(0, 1, 0)
-        glVertex(x, y, z)
-    glEnd()
-
-    glBegin(GL_TRIANGLE_STRIP)
-    for (x, z) in circle_pts:
-        y = h
-        glNormal(x-cx, 0, z-cz)
-        glVertex(x, 0, z)
-        glNormal(x-cx, 0, z-cz)
-        glVertex(x, h, z)
-    glEnd()
-
 def on_display():
     for i in range(0,updates_per_frame):
         lattice.evolve()
@@ -362,7 +292,9 @@ def on_display():
     shaders.glUseProgram(obstacle_program)
     glUniformMatrix4fv(projection_id, 1, False, numpy.ascontiguousarray(projection))
     glUniformMatrix4fv(rotation_id, 1, False, numpy.ascontiguousarray(rotation.get()))
-    draw_cylinder(lattice.geometry.size_x//6, lattice.geometry.size_z//2, lattice.geometry.size_y, cylinder_r, 32)
+
+    for primitive in primitives:
+        primitive.draw()
 
     glutSwapBuffers()
 
diff --git a/geometry/box.py b/geometry/box.py
new file mode 100644
index 0000000..5538819
--- /dev/null
+++ b/geometry/box.py
@@ -0,0 +1,58 @@
+from OpenGL.GL import *
+
+class Box:
+    def __init__(self, x0, x1, y0, y1, z0, z1):
+        self.x0 = x0
+        self.x1 = x1
+        self.y0 = y0
+        self.y1 = y1
+        self.z0 = z0
+        self.z1 = z1
+
+    def indicator(self):
+        return lambda x, y, z: x >= self.x0 and x <= self.x1 and y >= self.y0 and y <= self.y1 and z >= self.z0 and z <= self.z1
+
+    def draw(self):
+        glBegin(GL_POLYGON)
+        glNormal(-1,0,0)
+        glVertex(self.x0, self.y0, self.z0)
+        glVertex(self.x0, self.y1, self.z0)
+        glVertex(self.x0, self.y1, self.z1)
+        glVertex(self.x0, self.y0, self.z1)
+        glEnd()
+        glBegin(GL_POLYGON)
+        glNormal(1,0,0)
+        glVertex(self.x1, self.y0, self.z0)
+        glVertex(self.x1, self.y1, self.z0)
+        glVertex(self.x1, self.y1, self.z1)
+        glVertex(self.x1, self.y0, self.z1)
+        glEnd()
+        glBegin(GL_POLYGON)
+        glNormal(0,0,1)
+        glVertex(self.x0, self.y0, self.z1)
+        glVertex(self.x1, self.y0, self.z1)
+        glVertex(self.x1, self.y1, self.z1)
+        glVertex(self.x0, self.y1, self.z1)
+        glEnd()
+        glBegin(GL_POLYGON)
+        glNormal(0,0,-1)
+        glVertex(self.x0, self.y0, self.z0)
+        glVertex(self.x1, self.y0, self.z0)
+        glVertex(self.x1, self.y1, self.z0)
+        glVertex(self.x0, self.y1, self.z0)
+        glEnd()
+        glBegin(GL_POLYGON)
+        glNormal(0,-1,0)
+        glVertex(self.x0, self.y0, self.z0)
+        glVertex(self.x1, self.y0, self.z0)
+        glVertex(self.x1, self.y0, self.z1)
+        glVertex(self.x0, self.y0, self.z1)
+        glEnd()
+        glBegin(GL_POLYGON)
+        glNormal(0,1,0)
+        glVertex(self.x0,self.y1,self.z0)
+        glVertex(self.x1,self.y1,self.z0)
+        glVertex(self.x1,self.y1,self.z1)
+        glVertex(self.x0,self.y1,self.z1)
+        glEnd()
+
diff --git a/geometry/cylinder.py b/geometry/cylinder.py
new file mode 100644
index 0000000..b576a62
--- /dev/null
+++ b/geometry/cylinder.py
@@ -0,0 +1,72 @@
+import numpy
+
+from OpenGL.GL import *
+
+class Cylinder:
+    def __init__(self, x, y, z, r, h = 0, l = 0):
+        self.x = x
+        self.y = y
+        self.z = z
+        self.r = r
+        self.h = h
+        self.l = l
+
+        self.circle = []
+        for i in range(33):
+            alpha = 2 * numpy.pi * (i/32)
+            c0 = self.r * numpy.cos(alpha)
+            c1 = self.r * numpy.sin(alpha)
+            self.circle.append((c0, c1))
+
+    def indicator(self):
+        if self.h == 0:
+            return lambda x, y, z: (x - self.x)**2 + (z - self.z)**2 < self.r*self.r and y >= self.y and y <= self.y+self.l
+        else:
+            return lambda x, y, z: (x - self.x)**2 + (y - self.y)**2 < self.r*self.r and z >= self.z and z <= self.z+self.h
+
+    def draw(self):
+        glBegin(GL_TRIANGLE_FAN)
+
+        if self.h == 0:
+            glNormal(0,-1, 0)
+            glVertex(self.x, self.y, self.z)
+            for (c0, c1) in self.circle:
+                glNormal(0, -1, 0)
+                glVertex(self.x+c0, self.y, self.z+c1)
+        else:
+            glNormal(0, 0, -1)
+            glVertex(self.x, self.y, self.z)
+            for (c0, c1) in self.circle:
+                glVertex(self.x+c0, self.y+c1, self.z)
+
+        glEnd()
+
+        glBegin(GL_TRIANGLE_FAN)
+        if self.h == 0:
+            glNormal(0, 1, 0)
+            glVertex(self.x, self.y+self.l, self.z)
+            for (c0, c1) in self.circle:
+                glVertex(self.x+c0, self.y+self.l, self.z+c1)
+        else:
+            glNormal(0, 0, 1)
+            glVertex(self.x, self.y, self.z+self.h)
+            for (c0, c1) in self.circle:
+                glVertex(self.x+c0, self.y+c1, self.z+self.h)
+
+        glEnd()
+
+        glBegin(GL_TRIANGLE_STRIP)
+
+        if self.h == 0:
+            for (c0, c1) in self.circle:
+                glNormal(c0, 0, c1)
+                glVertex(self.x+c0, self.y, self.z+c1)
+                glVertex(self.x+c0, self.y + self.l, self.z+c1)
+        else:
+            for (c0, c1) in self.circle:
+                glNormal(c0, c1, 0)
+                glVertex(self.x+c0, self.y+c1, self.z)
+                glVertex(self.x+c0, self.y+c1, self.z+self.h)
+
+        glEnd()
+
diff --git a/geometry/sphere.py b/geometry/sphere.py
new file mode 100644
index 0000000..9bc245c
--- /dev/null
+++ b/geometry/sphere.py
@@ -0,0 +1,36 @@
+import numpy
+
+from OpenGL.GL import *
+
+class Sphere:
+    def __init__(self, x, y, z, r):
+        self.x = x
+        self.y = y
+        self.z = z
+        self.r = r
+
+    def indicator(self):
+        return lambda x, y, z: (x - self.x)**2 + (y - self.y)**2 + (z - self.z)**2 < self.r*self.r
+
+    def draw(self, resolution = 32):
+        for i in range(0,resolution+1):
+            lat0 = numpy.pi * (-0.5 + (i - 1) / resolution)
+            z0   = numpy.sin(lat0)
+            zr0  = numpy.cos(lat0)
+
+            lat1 = numpy.pi * (-0.5 + i / resolution)
+            z1   = numpy.sin(lat1)
+            zr1  = numpy.cos(lat1)
+
+            glBegin(GL_QUAD_STRIP)
+            for j in range(0,resolution+1):
+                lng = 2 * numpy.pi * (j - 1) / resolution
+                x = numpy.cos(lng)
+                y = numpy.sin(lng)
+
+                glNormal(x * zr0, y * zr0, z0)
+                glVertex(self.x + self.r * x * zr0, self.y + self.r * y * zr0, self.z + self.r * z0)
+                glNormal(x * zr1, y * zr1, z1)
+                glVertex(self.x + self.r * x * zr1, self.y + self.r * y * zr1, self.z + self.r * z1)
+
+            glEnd()
diff --git a/simulation.py b/simulation.py
index 7ddeb45..a274be4 100644
--- a/simulation.py
+++ b/simulation.py
@@ -177,8 +177,12 @@ class Lattice:
         self.material = numpy.ndarray(shape=(self.memory.volume, 1), dtype=numpy.int32)
 
     def apply_material_map(self, material_map):
-        for indicator, material in material_map:
-            self.material[[indicator(*idx) for idx in self.memory.cells()]] = material
+        for primitive, material in material_map:
+            if callable(primitive):
+                self.material[[primitive(*idx) for idx in self.memory.cells()]] = material
+            else:
+                indicator = primitive.indicator()
+                self.material[[indicator(*idx) for idx in self.memory.cells()]] = material
 
     def sync_material(self):
         cl.enqueue_copy(self.queue, self.memory.cl_material, self.material).wait();
diff --git a/template/kernel.mako b/template/kernel.mako
index 529eb30..a57c2ec 100644
--- a/template/kernel.mako
+++ b/template/kernel.mako
@@ -157,18 +157,22 @@ __kernel void update_particles(__global __read_only  float4* moments,
 
   const float4 moment = moments[gid];
 
-  if (material[gid] == 1 && particle.w < 1.0) {
+  if (material[gid] == 1) {
     particle.x += moment.y;
     particle.y += moment.z;
 % if descriptor.d == 2:
     particle.w += min(particle.x, particle.y) * aging;
 % elif descriptor.d == 3:
     particle.z += moment.w;
-    particle.w += min(min(particle.x, particle.y), particle.z) * aging;
+    float dy = (particle.y-${geometry.size_y/2.0});
+    float dz = (particle.z-${geometry.size_z/2.0});
+    dy *= dy;
+    dz *= dz;
+    particle.w = 10.0*sqrt(moment.y*moment.y+moment.z*moment.z+moment.w*moment.w);
 % endif
   } else {
     particle.xyz = init_particles[pid].xyz;
-    particle.w   = particle.w-1.0;
+    particle.w   = 0.0;
   }
 
   particles[pid] = particle;