Add SDF-based grid fin example

5 files changed, 502 insertions, 4 deletions

diff --git a/channel_3d_sdf_grid_fin_volumetric_rendering_gl_interop.py b/channel_3d_sdf_grid_fin_volumetric_rendering_gl_interop.py
new file mode 100644
index 0000000..166c36d
--- /dev/null
+++ b/channel_3d_sdf_grid_fin_volumetric_rendering_gl_interop.py
@@ -0,0 +1,352 @@
+import numpy
+
+from mako.template import Template
+from mako.lookup import TemplateLookup
+
+from pathlib import Path
+
+from simulation         import Lattice, Geometry
+from utility.particles  import Particles
+from symbolic.generator import LBM
+
+import symbolic.D3Q19 as D3Q19
+
+from OpenGL.GL   import *
+from OpenGL.GLUT import *
+
+from OpenGL.GL import shaders
+
+from geometry.box import Box
+
+from utility.projection import Projection, Rotation
+from utility.opengl     import MomentsTexture
+from utility.mouse      import MouseDragMonitor, MouseScrollMonitor
+
+lattice_x = 200
+lattice_y = 90
+lattice_z = 90
+
+updates_per_frame = 5
+
+inflow = 0.075
+relaxation_time = 0.51
+
+lbm = LBM(D3Q19)
+
+boundary = Template("""
+    if ( m == 2 ) {
+        u_0 = 0.0;
+        u_1 = 0.0;
+        u_2 = 0.0;
+    }
+    if ( m == 3 ) {
+        u_0 = min(time/5000.0 * ${inflow}, ${inflow});
+        u_1 = 0.0;
+        u_2 = 0.0;
+    }
+    if ( m == 4 ) {
+        rho = 1.0;
+    }
+""").render(
+    inflow = inflow
+)
+
+grid_fin = """
+v = rotate_z(translate(v, v3(center.x/2.5, center.y, center.z)), -0.5);
+float width = 1;
+float rotat = 0.64;
+
+return rounded(unify(
+	sub(
+		box(v, v3(5, 30, 40)),
+		box(v, v3(6, 28, 38))
+	),
+	intersect(
+		box(v, v3(5, 30, 40)),
+		unify(
+			unify(
+				box(rotate_x(v, rotat), v3(10, width, 100)),
+				box(rotate_x(v, -rotat), v3(10, width, 100))
+			),
+			unify(
+				unify(
+					unify(
+						box(rotate_x(translate(v, v3(0,0,25)), rotat), v3(10, width, 100)),
+						box(rotate_x(translate(v, v3(0,0,25)), -rotat), v3(10, width, 100))
+					),
+					unify(
+						box(rotate_x(translate(v, v3(0,0,-25)), rotat), v3(10, width, 100)),
+						box(rotate_x(translate(v, v3(0,0,-25)), -rotat), v3(10, width, 100))
+					)
+				),
+				unify(
+					unify(
+						box(rotate_x(translate(v, v3(0,0,50)), rotat), v3(10, width, 100)),
+						box(rotate_x(translate(v, v3(0,0,50)), -rotat), v3(10, width, 100))
+					),
+					unify(
+						box(rotate_x(translate(v, v3(0,0,-50)), rotat), v3(10, width, 100)),
+						box(rotate_x(translate(v, v3(0,0,-50)), -rotat), v3(10, width, 100))
+					)
+				)
+			)
+		)
+	)
+), 0.4);
+"""
+
+def glut_window(fullscreen = False):
+    glutInit(sys.argv)
+    glutSetOption(GLUT_MULTISAMPLE, 8)
+    glutInitDisplayMode(GLUT_RGBA | GLUT_DOUBLE | GLUT_MULTISAMPLE)
+
+    if fullscreen:
+        window = glutEnterGameMode()
+    else:
+        glutInitWindowSize(800, 500)
+        glutInitWindowPosition(0, 0)
+        window = glutCreateWindow("LBM")
+
+    return window
+
+window = glut_window(fullscreen = False)
+
+vertex_shader = shaders.compileShader("""
+#version 430
+
+layout (location=0) in vec4 vertex;
+
+uniform mat4 projection;
+uniform mat4 rotation;
+
+void main() {
+    gl_Position = projection * rotation * vertex;
+}""", GL_VERTEX_SHADER)
+
+fragment_shader = shaders.compileShader("""
+#version 430
+
+in vec3 color;
+
+void main(){
+    gl_FragColor = vec4(vec3(0.5), 1.0);
+}""", GL_FRAGMENT_SHADER)
+
+raycast_vertex_shader = shaders.compileShader("""
+#version 430
+
+layout (location=0) in vec4 vertex;
+                   out vec3 frag_pos;
+
+uniform mat4 projection;
+uniform mat4 rotation;
+
+void main() {
+    gl_Position = projection * rotation * vertex;
+    frag_pos    = vertex.xyz;
+}""", GL_VERTEX_SHADER)
+
+raycast_fragment_shader = shaders.compileShader(Template("""
+#version 430
+
+in vec3 frag_pos;
+
+uniform vec4 camera_pos;
+
+uniform sampler3D moments;
+
+out vec4 result;
+
+vec3 unit(vec3 v) {
+    return vec3(v[0] / ${size_x}, v[1] / ${size_y}, v[2] / ${size_z});
+}
+
+const vec3 center = vec3(${size_x/2.5}, ${size_y/2}, ${size_z/2});
+
+#define EPSILON 1e-1
+#define RAYMARCH_STEPS 64
+#define OBSTACLE_STEPS 16
+
+vec3 v3(float x, float y, float z) {
+	return vec3(x,y,z);
+}
+
+vec2 v2(float x, float y) {
+	return vec2(x,y);
+}
+
+vec3 fabs(vec3 x) {
+    return abs(x);
+}
+
+<%include file="template/sdf.lib.glsl.mako"/>
+
+float sdf(vec3 v) {
+    ${sdf_source}
+}
+
+vec3 sdf_normal(vec3 v) {
+    return normalize(vec3(
+        sdf(vec3(v.x + EPSILON, v.y, v.z)) - sdf(vec3(v.x - EPSILON, v.y, v.z)),
+        sdf(vec3(v.x, v.y + EPSILON, v.z)) - sdf(vec3(v.x, v.y - EPSILON, v.z)),
+        sdf(vec3(v.x, v.y, v.z + EPSILON)) - sdf(vec3(v.x, v.y, v.z - EPSILON))
+    ));
+}
+
+vec3 palette(float x) {
+    return mix(
+        vec3(0.0, 0.0, 0.0),
+        vec3(1.0, 0.0, 0.0),
+        x
+    );
+}
+
+float distanceToLattice(vec3 v) {
+  return box(v, vec3(${size_x},${size_y},${size_z}));
+}
+
+float maxRayLength(vec3 origin, vec3 ray) {
+    return max(1.0, ${max_ray_length} - distanceToLattice(origin + ${max_ray_length}*ray));
+}
+
+vec4 trace_obstacle(vec3 pos, vec3 ray) {
+    bool hit = false;
+    vec3 color  = vec3(0.05);
+    vec3 sample_pos = pos;
+
+    for (int i = 0; i < OBSTACLE_STEPS; ++i) {
+        float dist = sdf(sample_pos);
+        if (abs(dist) < EPSILON) {
+            vec3 n = normalize(sdf_normal(sample_pos));
+            float brightness = abs(dot(n, ray));
+            color += brightness;
+            hit = true;
+            break;
+        } else {
+            sample_pos += dist*ray;
+        }
+    }
+
+    if (hit) {
+        return vec4(color, 1.0);
+    } else {
+        return vec4(color, 0.0);
+    }
+}
+
+vec3 trace(vec3 pos, vec3 ray) {
+    const float delta = maxRayLength(pos, ray) / RAYMARCH_STEPS;
+
+    vec3 color = vec3(0.0);
+
+    for (int i=0; i < RAYMARCH_STEPS; ++i) {
+        const vec3 pos = pos + i*delta*ray;
+        const vec4 data = texture(moments, unit(pos));
+        if (sdf(pos) > delta) {
+            color += 1./RAYMARCH_STEPS * palette(length(data.yzw) / ${inflow});
+        } else {
+            vec4 obstacle_color = trace_obstacle(pos, ray);
+            if (obstacle_color.w == 1.0) {
+                color += obstacle_color.xyz;
+                break;
+            } else {
+                color += 1./RAYMARCH_STEPS * palette(length(data.yzw) / ${inflow});
+            }
+        }
+    }
+
+    return color;
+}
+
+void main(){
+    const vec3 ray = normalize(frag_pos - camera_pos.xyz);
+
+    result = vec4(trace(frag_pos, ray), 1.0);
+}
+""", lookup = TemplateLookup(directories = [
+    Path(__file__).parent
+])).render(
+    size_x = lattice_x,
+    size_y = lattice_y,
+    size_z = lattice_z,
+    inflow = inflow,
+    max_ray_length = max(lattice_x,lattice_y,lattice_z)**3,
+    sdf_source = grid_fin
+), GL_FRAGMENT_SHADER)
+
+domain_program   = shaders.compileProgram(vertex_shader, fragment_shader)
+domain_projection_id = shaders.glGetUniformLocation(domain_program, 'projection')
+domain_rotation_id   = shaders.glGetUniformLocation(domain_program, 'rotation')
+
+raycast_program = shaders.compileProgram(raycast_vertex_shader, raycast_fragment_shader)
+raycast_projection_id       = shaders.glGetUniformLocation(raycast_program, 'projection')
+raycast_rotation_id         = shaders.glGetUniformLocation(raycast_program, 'rotation')
+raycast_camera_pos_id = shaders.glGetUniformLocation(raycast_program, 'camera_pos')
+
+lattice = Lattice(
+    descriptor   = D3Q19,
+    geometry     = Geometry(lattice_x, lattice_y, lattice_z),
+    moments      = lbm.moments(optimize = True),
+    collide      = lbm.bgk(f_eq = lbm.equilibrium(), tau = relaxation_time, optimize = True),
+    boundary_src = boundary,
+    opengl       = True
+)
+
+lattice.setup_channel_with_sdf_obstacle(grid_fin)
+
+moments_texture = MomentsTexture(lattice)
+
+projection = Projection(distance = 2*lattice_x)
+rotation = Rotation([-0.5*lattice_x, -0.5*lattice_y, -0.5*lattice_z])
+
+cube_vertices, cube_edges = lattice.geometry.wireframe()
+
+def on_display():
+    for i in range(0,updates_per_frame):
+        lattice.evolve()
+
+    moments_texture.collect()
+
+    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
+
+    glEnable(GL_DEPTH_TEST)
+    glDepthFunc(GL_LESS)
+
+    camera_pos = numpy.matmul([0,-projection.distance,0,1], rotation.get_inverse())
+
+    shaders.glUseProgram(domain_program)
+    glUniformMatrix4fv(domain_projection_id, 1, False, numpy.ascontiguousarray(projection.get()))
+    glUniformMatrix4fv(domain_rotation_id,   1, False, numpy.ascontiguousarray(rotation.get()))
+    glLineWidth(3)
+    glBegin(GL_LINES)
+    for i, j in cube_edges:
+        glVertex(cube_vertices[i])
+        glVertex(cube_vertices[j])
+    glEnd()
+
+    shaders.glUseProgram(raycast_program)
+    glUniformMatrix4fv(raycast_projection_id, 1, False, numpy.ascontiguousarray(projection.get()))
+    glUniformMatrix4fv(raycast_rotation_id,   1, False, numpy.ascontiguousarray(rotation.get()))
+    glUniform4fv(raycast_camera_pos_id, 1, camera_pos)
+    moments_texture.bind()
+    Box(0,lattice.geometry.size_x,0,lattice.geometry.size_y,0,lattice.geometry.size_z).draw()
+
+    glutSwapBuffers()
+
+mouse_monitors = [
+    MouseDragMonitor(GLUT_LEFT_BUTTON,  lambda dx, dy: rotation.update(0.005*dy, 0.005*dx)),
+    MouseDragMonitor(GLUT_RIGHT_BUTTON, lambda dx, dy: rotation.shift(0.25*dx, 0.25*dy)),
+    MouseScrollMonitor(lambda zoom: projection.update_distance(5*zoom))
+]
+
+def on_timer(t):
+    glutTimerFunc(t, on_timer, t)
+    glutPostRedisplay()
+
+glutDisplayFunc(on_display)
+glutReshapeFunc(lambda w, h: projection.update_ratio(w, h))
+glutMouseFunc(lambda *args: list(map(lambda m: m.on_mouse(*args), mouse_monitors)))
+glutMotionFunc(lambda *args: list(map(lambda m: m.on_mouse_move(*args), mouse_monitors)))
+glutTimerFunc(30, on_timer, 30)
+
+glutMainLoop()
diff --git a/simulation.py b/simulation.py
index de30e64..535310f 100644
--- a/simulation.py
+++ b/simulation.py
@@ -7,6 +7,8 @@ from utility.ndindex import ndindex
 import sympy
 
 from mako.template import Template
+from mako.lookup import TemplateLookup
+
 from pathlib import Path
 
 from pyopencl.tools import get_gl_sharing_context_properties
@@ -113,7 +115,7 @@ class Memory:
         self.cl_pop_b = cl.Buffer(self.context, mf.READ_WRITE, size=self.pop_size)
 
         self.cl_moments  = cl.Buffer(self.context, mf.WRITE_ONLY, size=self.moments_size)
-        self.cl_material = cl.Buffer(self.context, mf.READ_ONLY, size=self.volume * numpy.int32(0).nbytes)
+        self.cl_material = cl.Buffer(self.context, mf.READ_WRITE, size=self.volume * numpy.int32(0).nbytes)
 
     def gid(self, x, y, z = 0):
         return z * (self.size_x*self.size_y) + y * self.size_x + x;
@@ -131,7 +133,7 @@ class Lattice:
     def __init__(self,
         descriptor, geometry, moments, collide,
         pop_eq_src = '', boundary_src = '',
-        platform = 0, precision = 'single', layout = None, padding = None, align = True, opengl = False
+        platform = 0, precision = 'single', layout = None, padding = None, align = False, opengl = False
     ):
         self.descriptor = descriptor
         self.geometry   = geometry
@@ -144,6 +146,10 @@ class Lattice:
             'double': (numpy.float64, 'double'),
         }.get(precision, None)
 
+        self.mako_lookup = TemplateLookup(directories = [
+            Path(__file__).parent
+        ])
+
         self.platform = cl.get_platforms()[platform]
 
         if opengl:
@@ -186,8 +192,21 @@ class Lattice:
                 indicator = primitive.indicator()
                 self.material[[indicator(*idx) for idx in self.memory.cells()]] = material
 
+    def setup_channel_with_sdf_obstacle(self, sdf_src):
+        sdf_kernel_src = Template(
+            filename = 'template/sdf.cl.mako',
+            lookup   = self.mako_lookup
+        ).render(
+            geometry = self.memory,
+            sdf_src  = sdf_src
+        )
+
+        sdf_program = cl.Program(self.context, sdf_kernel_src).build(self.compiler_args)
+        sdf_program.setup_channel_with_sdf_obstacle(self.queue, self.memory.size(), None, self.memory.cl_material)
+        cl.enqueue_copy(self.queue, self.material, self.memory.cl_material).wait()
+
     def sync_material(self):
-        cl.enqueue_copy(self.queue, self.memory.cl_material, self.material).wait();
+        cl.enqueue_copy(self.queue, self.memory.cl_material, self.material).wait()
 
     def build_kernel(self):
         program_src = Template(filename = str(Path(__file__).parent/'template/kernel.mako')).render(
diff --git a/template/sdf.cl.mako b/template/sdf.cl.mako
new file mode 100644
index 0000000..b98b35f
--- /dev/null
+++ b/template/sdf.cl.mako
@@ -0,0 +1,56 @@
+typedef float3 vec3;
+typedef float2 vec2;
+
+float3 v3(float x, float y, float z) {
+	return (float3)(x,y,z);
+}
+
+float2 v2(float x, float y) {
+	return (float2)(x,y);
+}
+
+__constant float3 center = (float3)(${geometry.size_x/2.5}, ${geometry.size_y/2}, ${geometry.size_z/2});
+
+<%include file="sdf.lib.glsl.mako"/>
+
+float sdf(vec3 v) {
+    ${sdf_src}
+}
+
+__kernel void setup_channel_with_sdf_obstacle(__global int* material) {
+    const unsigned x = get_global_id(0);
+    const unsigned y = get_global_id(1);
+    const unsigned z = get_global_id(2);
+
+    const unsigned gid = z*${geometry.size_x*geometry.size_y} + y*${geometry.size_x} + x;
+
+    if (x == 0 || x == ${geometry.size_x-1} ||
+        y == 0 || y == ${geometry.size_y-1} ||
+        z == 0 || z == ${geometry.size_z-1}) {
+        material[gid] = 0;
+        return;
+    }
+
+    if (x == 1) {
+        material[gid] = 3;
+        return;
+    }
+
+    if (x == ${geometry.size_x-2}) {
+        material[gid] = 4;
+        return;
+    }
+
+    if (y == 1 || y == ${geometry.size_y-2} ||
+        z == 1 || z == ${geometry.size_z-2}) {
+        material[gid] = 2;
+        return;
+    }
+
+    if (sdf((float3)(x,y,z)) < 0.0) {
+        material[gid] = 2;
+        return;
+    }
+
+    material[gid] = 1;
+}
diff --git a/template/sdf.lib.glsl.mako b/template/sdf.lib.glsl.mako
new file mode 100644
index 0000000..515f959
--- /dev/null
+++ b/template/sdf.lib.glsl.mako
@@ -0,0 +1,71 @@
+float sphere(vec3 v, float r) {
+	return length(v) - r;
+}
+
+float torus(vec3 v, vec2 t) {
+	vec2 q = v2(length(v.xz)-t.x,v.y);
+	return length(q)-t.y;
+}
+
+float box(vec3 v, vec3 b) {
+  vec3 q = fabs(v) - b;
+  return length(max(q,0.0)) + min(max(q.x,max(q.y,q.z)),0.0);
+}
+
+vec3 flip_xy(vec3 v) {
+	return v3(v.y, v.x, v.z);
+}
+
+vec3 flip_yz(vec3 v) {
+	return v3(v.x, v.z, v.y);
+}
+
+vec3 rotate_x(vec3 v, float r) {
+	v.yz = cos(r)*v.yz + sin(r)*v2(v.z, -v.y);
+	return v;
+}
+
+vec3 rotate_y(vec3 v, float r) {
+	v.xz = cos(r)*v.xz + sin(r)*v2(v.z, -v.x);
+	return v;
+}
+
+vec3 rotate_z(vec3 v, float r) {
+	v.xy = cos(r)*v.xy + sin(r)*v2(v.y, -v.x);
+	return v;
+}
+
+vec3 translate(vec3 v, vec3 w) {
+	return v - w;
+}
+
+float rounded(float a, float r) {
+	return a - r;
+}
+
+float sunify(float a, float b, float k) {
+	float h = clamp(0.5 + 0.5 * (b - a) / k, 0.0, 1.0);
+	return mix(b, a, h) - k * h * (1 - h);
+}
+
+float ssub(float b, float a, float k) {
+	float h = clamp(0.5 - 0.5*(b+a)/k, 0.0, 1.0);
+	return mix(b, -a, h) + k*h*(1.0-h);
+}
+
+float sintersect(float a, float b, float k) {
+	float h = clamp(0.5 - 0.5*(b-a)/k, 0.0, 1.0);
+	return mix(b, a, h) + k*h*(1.0-h);
+}
+
+float sub(float a, float b) {
+	return max(-b, a);
+}
+
+float unify(float a, float b) {
+	return min(a, b);
+}
+
+float intersect(float a, float b) {
+	return max(a, b);
+}
diff --git a/utility/projection.py b/utility/projection.py
index 1750fac..76791d9 100644
--- a/utility/projection.py
+++ b/utility/projection.py
@@ -11,7 +11,7 @@ class Projection:
         self.update()
 
     def update(self):
-        projection = matrix44.create_perspective_projection(20.0, self.ratio, 0.1, 1000.0)
+        projection = matrix44.create_perspective_projection(20.0, self.ratio, 0.1, 5000.0)
         look = matrix44.create_look_at(
             eye    = [0, -self.distance, 0],
             target = [0, 0, 0],