feat: add primitive raytracing

.python-version

@@ -0,0 +1 @@
+3.11.11

code/main.py

@@ -1,87 +1,158 @@
-from OpenGL.GL import *
-from OpenGL.GLUT import *
-from OpenGL.GL.shaders import compileProgram, compileShader
 import numpy as np
+import matplotlib.pyplot as plt
 
-# Vertex shader
+w = 400
-VERTEX_SHADER = """
+h = 300
-#version 330
-in vec3 position;
-in vec3 color;
-out vec3 vertexColor;
-void main() {
-    gl_Position = vec4(position, 1.0);
-    vertexColor = color;
-}
-"""
 
-# Fragment shader
+def normalize(x):
-FRAGMENT_SHADER = """
+    x /= np.linalg.norm(x)
-#version 330
+    return x
-in vec3 vertexColor;
-out vec4 fragColor;
-void main() {
-    fragColor = vec4(vertexColor, 1.0);
-}
-"""
 
-# Define vertices and colors
+def intersect_plane(O, D, P, N):
-vertices = np.array([
+    # Return the distance from O to the intersection of the ray (O, D) with the 
-    # Positions       # Colors
+    # plane (P, N), or +inf if there is no intersection.
-    0.0,  0.5, 0.0,    1.0, 0.0, 0.0,  # Top (red)
+    # O and P are 3D points, D and N (normal) are normalized vectors.
-   -0.5, -0.5, 0.0,    0.0, 1.0, 0.0,  # Bottom-left (green)
+    denom = np.dot(D, N)
-    0.5, -0.5, 0.0,    0.0, 0.0, 1.0   # Bottom-right (blue)
+    if np.abs(denom) < 1e-6:
-], dtype=np.float32)
+        return np.inf
+    d = np.dot(P - O, N) / denom
+    if d < 0:
+        return np.inf
+    return d
 
-# Initialize OpenGL
+def intersect_sphere(O, D, S, R):
-def init():
+    # Return the distance from O to the intersection of the ray (O, D) with the 
-    global shader, VAO
+    # sphere (S, R), or +inf if there is no intersection.
+    # O and S are 3D points, D (direction) is a normalized vector, R is a scalar.
+    a = np.dot(D, D)
+    OS = O - S
+    b = 2 * np.dot(D, OS)
+    c = np.dot(OS, OS) - R * R
+    disc = b * b - 4 * a * c
+    if disc > 0:
+        distSqrt = np.sqrt(disc)
+        q = (-b - distSqrt) / 2.0 if b < 0 else (-b + distSqrt) / 2.0
+        t0 = q / a
+        t1 = c / q
+        t0, t1 = min(t0, t1), max(t0, t1)
+        if t1 >= 0:
+            return t1 if t0 < 0 else t0
+    return np.inf
 
-    # Compile shaders
+def intersect(O, D, obj):
-    shader = compileProgram(
+    if obj['type'] == 'plane':
-        compileShader(VERTEX_SHADER, GL_VERTEX_SHADER),
+        return intersect_plane(O, D, obj['position'], obj['normal'])
-        compileShader(FRAGMENT_SHADER, GL_FRAGMENT_SHADER)
+    elif obj['type'] == 'sphere':
-    )
+        return intersect_sphere(O, D, obj['position'], obj['radius'])
 
-    # Generate VAO and VBO
+def get_normal(obj, M):
-    VAO = glGenVertexArrays(1)
+    # Find normal.
-    VBO = glGenBuffers(1)
+    if obj['type'] == 'sphere':
+        N = normalize(M - obj['position'])
+    elif obj['type'] == 'plane':
+        N = obj['normal']
+    return N
+    
+def get_color(obj, M):
+    color = obj['color']
+    if not hasattr(color, '__len__'):
+        color = color(M)
+    return color
 
-    glBindVertexArray(VAO)
+def trace_ray(rayO, rayD):
+    # Find first point of intersection with the scene.
+    t = np.inf
+    for i, obj in enumerate(scene):
+        t_obj = intersect(rayO, rayD, obj)
+        if t_obj < t:
+            t, obj_idx = t_obj, i
+    # Return None if the ray does not intersect any object.
+    if t == np.inf:
+        return
+    # Find the object.
+    obj = scene[obj_idx]
+    # Find the point of intersection on the object.
+    M = rayO + rayD * t
+    # Find properties of the object.
+    N = get_normal(obj, M)
+    color = get_color(obj, M)
+    toL = normalize(L - M)
+    toO = normalize(O - M)
+    # Shadow: find if the point is shadowed or not.
+    l = [intersect(M + N * .0001, toL, obj_sh) 
+            for k, obj_sh in enumerate(scene) if k != obj_idx]
+    if l and min(l) < np.inf:
+        return
+    # Start computing the color.
+    col_ray = ambient
+    # Lambert shading (diffuse).
+    col_ray += obj.get('diffuse_c', diffuse_c) * max(np.dot(N, toL), 0) * color
+    # Blinn-Phong shading (specular).
+    col_ray += obj.get('specular_c', specular_c) * max(np.dot(N, normalize(toL + toO)), 0) ** specular_k * color_light
+    return obj, M, N, col_ray
 
-    glBindBuffer(GL_ARRAY_BUFFER, VBO)
+def add_sphere(position, radius, color):
-    glBufferData(GL_ARRAY_BUFFER, vertices.nbytes, vertices, GL_STATIC_DRAW)
+    return dict(type='sphere', position=np.array(position), 
+        radius=np.array(radius), color=np.array(color), reflection=.5)
+    
+def add_plane(position, normal):
+    return dict(type='plane', position=np.array(position), 
+        normal=np.array(normal),
+        color=lambda M: (color_plane0 
+            if (int(M[0] * 2) % 2) == (int(M[2] * 2) % 2) else color_plane1),
+        diffuse_c=.75, specular_c=.5, reflection=.25)
+    
+# List of objects.
+color_plane0 = 1. * np.ones(3)
+color_plane1 = 0. * np.ones(3)
+scene = [add_sphere([.75, .1, 1.], .6, [1., 0., 0.]),
+         add_sphere([-.75, .1, 2.25], .6, [0., 1., 0.]),
+         add_sphere([-2.75, .1, 3.5], .6, [0., 0., 1.]),
+         add_plane([0., -.5, 0.], [0., 1., 0.]),
+    ]
 
-    # Position attribute
+# Light position and color.
-    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 6 * vertices.itemsize, None)
+L = np.array([5., 5., -10.])
-    glEnableVertexAttribArray(0)
+color_light = np.ones(3)
 
-    # Color attribute
+# Default light and material parameters.
-    glVertexAttribPointer(1, 3, GL_FLOAT, GL_FALSE, 6 * vertices.itemsize, ctypes.c_void_p(3 * vertices.itemsize))
+ambient = .05
-    glEnableVertexAttribArray(1)
+diffuse_c = 1.
+specular_c = 1.
+specular_k = 50
 
-    glBindBuffer(GL_ARRAY_BUFFER, 0)
+depth_max = 5  # Maximum number of light reflections.
-    glBindVertexArray(0)
+col = np.zeros(3)  # Current color.
+O = np.array([0., 0.35, -1.])  # Camera.
+Q = np.array([0., 0., 0.])  # Camera pointing to.
+img = np.zeros((h, w, 3))
 
-# Render function
+r = float(w) / h
-def display():
+# Screen coordinates: x0, y0, x1, y1.
-    glClear(GL_COLOR_BUFFER_BIT)
+S = (-1., -1. / r + .25, 1., 1. / r + .25)
-    glUseProgram(shader)
-    glBindVertexArray(VAO)
-    glDrawArrays(GL_TRIANGLES, 0, 3)
-    glBindVertexArray(0)
-    glUseProgram(0)
-    glutSwapBuffers()
 
-# Main function
+# Loop through all pixels.
-def main():
+for i, x in enumerate(np.linspace(S[0], S[2], w)):
-    glutInit()
+    if i % 10 == 0:
-    glutInitDisplayMode(GLUT_DOUBLE | GLUT_RGB)
+        print(i / float(w) * 100, "%")
-    glutCreateWindow(b"PyOpenGL Triangle")
+    for j, y in enumerate(np.linspace(S[1], S[3], h)):
-    glutInitWindowSize(800, 600)
+        col[:] = 0
-    glutDisplayFunc(display)
+        Q[:2] = (x, y)
-    init()
+        D = normalize(Q - O)
-    glutMainLoop()
+        depth = 0
+        rayO, rayD = O, D
+        reflection = 1.
+        # Loop through initial and secondary rays.
+        while depth < depth_max:
+            traced = trace_ray(rayO, rayD)
+            if not traced:
+                break
+            obj, M, N, col_ray = traced
+            # Reflection: create a new ray.
+            rayO, rayD = M + N * .0001, normalize(rayD - 2 * np.dot(rayD, N) * N)
+            depth += 1
+            col += reflection * col_ray
+            reflection *= obj.get('reflection', 1.)
+        img[h - j - 1, i, :] = np.clip(col, 0, 1)
 
-if __name__ == "__main__":
+plt.imsave('fig.png', img)
-    main()

code/requirements.txt

@@ -1,4 +1,5 @@
 glfw
 numpy
 pyrr
-PyOpenGL
+PyOpenGL
+matplotlib