Implement basic diffusive lighting for 3d obstacle

1 files changed, 63 insertions, 6 deletions

diff --git a/channel_3d_gl_interop.py b/channel_3d_gl_interop.py
index 695dbfe..a023bc3 100644
--- a/channel_3d_gl_interop.py
+++ b/channel_3d_gl_interop.py
@@ -189,9 +189,28 @@ uniform mat4 rotation;
 
 void main() {
     gl_Position = projection * rotation * vertex;
-    color = vec3(0.5,0.5,0.5);
+    color = vec3(1.0,1.0,1.0);
 }""").substitute({}), GL_VERTEX_SHADER)
 
+lighting_vertex_shader = shaders.compileShader("""
+#version 430
+
+layout (location=0) in vec4 vertex;
+layout (location=2) in vec4 normal;
+                   out vec3 color;
+                   out vec3 frag_pos;
+                   out vec3 frag_normal;
+
+uniform mat4 projection;
+uniform mat4 rotation;
+
+void main() {
+    gl_Position = projection * rotation * vertex;
+    frag_pos    = vertex.xyz;
+    frag_normal = normalize(normal.xyz);
+    color = vec3(0.6,0.6,0.6);
+}""", GL_VERTEX_SHADER)
+
 fragment_shader = shaders.compileShader("""
 #version 430
 
@@ -201,11 +220,40 @@ void main(){
     gl_FragColor = vec4(color.xyz, 0.0);
 }""", GL_FRAGMENT_SHADER)
 
+lighting_fragment_shader = shaders.compileShader(Template("""
+#version 430
+
+in vec3 color;
+in vec3 frag_pos;
+in vec3 frag_normal;
+
+uniform mat4 projection;
+uniform mat4 rotation;
+
+out vec4 result;
+
+void main(){
+    const vec4 light_pos = rotation * vec4($size_x/2,-$size_x,$size_z/2,1);
+    const vec3 light_color = vec3(1.0,1.0,1.0);
+
+    const vec3 ray = light_pos.xyz - frag_pos;
+    float brightness = dot(frag_normal, ray) / length(ray);
+    brightness = clamp(brightness, 0, 1);
+
+    result = vec4(max(0.4,brightness) * light_color * color.rgb, 1.0);
+}
+""").substitute({
+    "size_x": lattice_x,
+    "size_y": lattice_y,
+    "size_z": lattice_z
+}), GL_FRAGMENT_SHADER)
+
 particle_program = shaders.compileProgram(particle_shader, geometry_shader, fragment_shader)
 projection_id = shaders.glGetUniformLocation(particle_program, 'projection')
 rotation_id   = shaders.glGetUniformLocation(particle_program, 'rotation')
 
-geometry_program = shaders.compileProgram(vertex_shader, fragment_shader)
+domain_program   = shaders.compileProgram(vertex_shader, fragment_shader)
+obstacle_program = shaders.compileProgram(lighting_vertex_shader, lighting_fragment_shader)
 
 lattice = Lattice(
     descriptor   = D3Q27,
@@ -248,23 +296,29 @@ def draw_cylinder(cx, cz, height, radius, num_slices):
         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()
 
@@ -295,16 +349,19 @@ def on_display():
     glPointSize(point_size)
     glDrawArrays(GL_POINTS, 0, particles.count)
 
-    shaders.glUseProgram(geometry_program)
+    shaders.glUseProgram(domain_program)
     glUniformMatrix4fv(projection_id, 1, False, numpy.ascontiguousarray(projection))
     glUniformMatrix4fv(rotation_id, 1, False, numpy.ascontiguousarray(rotation.get()))
-    glLineWidth(2*point_size)
+    glLineWidth(point_size)
     glBegin(GL_LINES)
     for i, j in cube_edges:
-        glVertex3fv(cube_vertices[i])
-        glVertex3fv(cube_vertices[j])
+        glVertex(cube_vertices[i])
+        glVertex(cube_vertices[j])
     glEnd()
 
+    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)
 
     glutSwapBuffers()