Very slow ray tracer

tinyraytracer.py

import numpy as np
import numba as nb

jit = nb.njit(parallel=True, fastmath=True)


@jit
def v(x):
    return np.array(x, dtype=np.float64)


@nb.jitclass((
        ('position', nb.float64[:]),
        ('intensity', nb.float64)))
class Light:
    def __init__(self, position, intensity):
        self.position = position
        self.intensity = intensity


@nb.jitclass((
        ('refractive_index', nb.float64),
        ('albedo', nb.float64[:]),
        ('diffuse_color', nb.float64[:]),
        ('specular_exponent', nb.float64)))
class Material:
    def __init__(self, refractive_index, albedo, diffuse_color, specular_exponent):
        self.refractive_index = refractive_index
        self.albedo = albedo
        self.diffuse_color = diffuse_color
        self.specular_exponent = specular_exponent


material_type = nb.deferred_type()
material_type.define(Material.class_type.instance_type)


@nb.jitclass((
        ('center', nb.float64[:]),
        ('radius', nb.float64),
        ('material', material_type)))
class Sphere:
    def __init__(self, center, radius, material):
        self.center = center
        self.radius = radius
        self.material = material

    def ray_intersect(self, orig, dir):
        L = self.center - orig
        tca = L @ dir
        d2 = L @ L - tca ** 2
        radius2 = self.radius ** 2
        if d2 > radius2:
            return False, 0.0
        thc = np.sqrt(radius2 - d2)
        t0 = tca - thc
        t1 = tca + thc
        if t0 < 0:
            t0 = t1
        if t0 < 0:
            return False, 0.0
        return True, t0


@jit
def reflect(I, N):
    return I - N * 2 * (I @ N)


@jit
def refract(I, N, refractive_index):
    """Snell's law"""
    cosi = -max(-1.0, min(1.0, I @ N))
    etai = 1
    etat = refractive_index
    n = N
    # if the ray is inside the object, swap the indices and invert the normal to get the correct result
    if cosi < 0:
        cosi = -cosi
        etai, etat = etat, etai
        n = -N
    eta = etai / etat
    k = 1 - eta * eta * (1 - cosi * cosi)
    if k < 0:
        return v([0.0, 0.0, 0.0])
    else:
        return I * eta + n * (eta * cosi - np.sqrt(k))


@jit
def normalize(x):
    return x / np.linalg.norm(x)


@jit
def scene_intersect(orig, dir, spheres):
    hit = v([0.0, 0.0, 0.0])
    N = v([0.0, 0.0, 0.0])
    material = Material(1.0,
                        v([1.0, 0.0, 0.0, 0.0]),
                        v([0.0, 0.0, 0.0]),
                        0.0)

    spheres_dist = np.Inf
    for sphere in spheres:
        intersects, dist_i = sphere.ray_intersect(orig, dir)
        if intersects and dist_i < spheres_dist:
            spheres_dist = dist_i
            hit = orig + dir * dist_i
            N = normalize(hit - sphere.center)
            material = sphere.material

    checkerboard_dist = np.Inf
    if np.abs(dir[1]) > 1e-3:
        # the checkerboard plane has equation y = -4
        d = -(orig[1] + 4) / dir[1]
        pt = orig + dir * d
        if 0 < d < spheres_dist and np.abs(pt[0]) < 10 and -30 < pt[2] < -10:
            checkerboard_dist = d
            hit = pt
            N = v([0.0, 1.0, 0.0])
            if (int(0.5 * hit[0] + 1000) + int(0.5 * hit[2])) & 1 != 0:
                material.diffuse_color = v([0.3, 0.3, 0.3])
            else:
                material.diffuse_color = v([0.3, 0.2, 0.1])
    return min(spheres_dist, checkerboard_dist) < 1000, hit, N, material


@jit
def cast_ray(orig, dir, spheres, lights, depth=0):
    intersects, point, N, material = scene_intersect(orig, dir, spheres)

    if depth > 4 or not intersects:
        return v([0.2, 0.7, 0.8])  # background color

    reflect_dir = normalize(reflect(dir, N))
    refract_dir = normalize(refract(dir, N, material.refractive_index))

    # offset the original point to avoid occlusion by the object itself
    reflect_orig = point - N * 1e-3 if reflect_dir @ N < 0 else point + N * 1e-3
    refract_orig = point - N * 1e-3 if refract_dir @ N < 0 else point + N * 1e-3

    reflect_color = cast_ray(reflect_orig, reflect_dir, spheres, lights, depth + 1)
    refract_color = cast_ray(refract_orig, refract_dir, spheres, lights, depth + 1)

    diffuse_light_intensity = 0.0
    specular_light_intensity = 0.0
    for light in lights:
        light_dir = normalize(light.position - point)
        light_distance = np.linalg.norm(light.position - point)

        # checking if the point lies in the shadow of the light
        shadow_orig = point - N * 1e-3 if light_dir @ N < 0 else point + N * 1e-3

        intersects, shadow_pt, shadow_N, tmpmaterial = scene_intersect(shadow_orig, light_dir, spheres)
        if intersects and np.linalg.norm(shadow_pt - shadow_orig) < light_distance:
            continue

        diffuse_light_intensity += light.intensity * max(0.0, light_dir @ N)
        specular_light_intensity += (max(0.0,
                                         -reflect(-light_dir, N) @ dir) ** material.specular_exponent) * light.intensity

    return material.diffuse_color * diffuse_light_intensity * material.albedo[0] + \
           v([1.0, 1.0, 1.0]) * specular_light_intensity * material.albedo[1] + \
           reflect_color * material.albedo[2] + refract_color * material.albedo[3]


def render(spheres, lights):
    width = 1024
    height = 768
    fov = np.pi / 2.0
    framebuffer = []

    for j in nb.prange(height):
        for i in nb.prange(width):
            x = (2 * (i + 0.5) / width - 1) * np.tan(fov / 2.0) * width / height
            y = -(2 * (j + 0.5) / height - 1) * np.tan(fov / 2.0)
            dir = normalize(v([x, y, -1.0]))
            framebuffer.append(cast_ray(v([0.0, 0.0, 0.0]), dir, spheres, lights))

    with open('out.ppm', 'wb') as ofs:
        ofs.write(b'P6\n')
        ofs.write(str(width).encode())
        ofs.write(b' ')
        ofs.write(str(height).encode())
        ofs.write(b'\n255\n')
        for c in framebuffer:
            mx = max(c[0], c[1], c[2])
            if mx > 1:
                c /= mx
            ofs.write(bytes([int(255 * max(0.0, min(1.0, x))) for x in c]))


def main():
    ivory = Material(1.0, v([0.6, 0.3, 0.1, 0.0]), v([0.4, 0.4, 0.3]), 50.0)
    glass = Material(1.5, v([0.0, 0.5, 0.1, 0.8]), v([0.6, 0.7, 0.8]), 125.0)
    red_rubber = Material(1.0, v([0.9, 0.1, 0.0, 0.0]), v([0.3, 0.1, 0.1]), 10.0)
    mirror = Material(1.0, v([0.0, 10.0, 0.8, 0.0]), v([1.0, 1.0, 1.0]), 1425.0)

    spheres = [
        Sphere(v([-3.0, 0.0, -16.0]), 2.0, ivory),
        Sphere(v([-1.0, -1.5, -12.0]), 2.0, glass),
        Sphere(v([1.5, -0.5, -18.0]), 3.0, red_rubber),
        Sphere(v([7.0, 5.0, -18.0]), 4.0, mirror),
    ]

    lights = [
        Light(v([-20.0, 20.0, 20.0]), 1.5),
        Light(v([30.0, 50.0, -25.0]), 1.8),
        Light(v([30.0, 20.0, 30.0]), 1.7),
    ]

    render(spheres, lights)


if __name__ == '__main__':
    main()