Sample ray trace using python

ray-trace.py

import numpy as np
import pygame
import time

# Define screen dimensions
WIDTH, HEIGHT = 1024, 768

# Define colors
LIGHT_COLOR = np.array([1, 1, 1])

def create_background():
    background = np.zeros((HEIGHT, WIDTH, 3))
    
    # Sky gradient
    for y in range(HEIGHT):
        t = y / HEIGHT
        background[y, :, :] = np.array([0.5, 0.7, 0.9]) * (1 - t) + np.array([0.1, 0.5, 0.2]) * t
    
    # Hills
    for x in range(WIDTH):
        hill_y = int(HEIGHT * (0.5 + 0.1 * np.sin(x * 0.02)))
        background[hill_y:, x, :] = np.array([0.1, 0.5, 0.2])
    
    # Clouds
    for _ in range(10):
        cx, cy = np.random.randint(0, WIDTH), np.random.randint(0, HEIGHT // 2)
        r = np.random.randint(20, 50)
        for x in range(cx - r, cx + r):
            for y in range(cy - r, cy + r):
                if 0 <= x < WIDTH and 0 <= y < HEIGHT:
                    if (x - cx) ** 2 + (y - cy) ** 2 < r ** 2:
                        background[y, x, :] = np.array([0.9, 0.9, 0.9])
    
    return background

# Define scene objects
class Sphere:
    def __init__(self, center, radius, color, reflection=0.5, refraction=0.9):
        self.center = np.array(center)
        self.radius = radius
        self.color = np.array(color)
        self.reflection = reflection
        self.refraction = refraction

    def intersect(self, ray_origin, ray_direction):
        oc = ray_origin - self.center
        a = np.dot(ray_direction, ray_direction)
        b = 2 * np.dot(oc, ray_direction)
        c = np.dot(oc, oc) - self.radius ** 2
        discriminant = b ** 2 - 4 * a * c

        if discriminant < 0:
            return None

        t1 = (-b - np.sqrt(discriminant)) / (2 * a)
        t2 = (-b + np.sqrt(discriminant)) / (2 * a)
        return t1 if t1 > 0 else t2 if t2 > 0 else None

def normalize(v):
    return v / np.linalg.norm(v)

def trace(ray_origin, ray_direction, scene, background, depth=6):
    y_idx = int(np.clip((ray_direction[1] + 1) * HEIGHT // 2, 0, HEIGHT - 1))
    x_idx = int(np.clip((ray_direction[0] + 1) * WIDTH // 2, 0, WIDTH - 1))
    background_color = background[y_idx, x_idx]

    if depth == 0:
        return background_color

    closest_t = float('inf')
    closest_sphere = None
    for sphere in scene:
        t = sphere.intersect(ray_origin, ray_direction)
        if t and t < closest_t:
            closest_t = t
            closest_sphere = sphere

    if closest_sphere is None:
        return background_color

    hit_point = ray_origin + closest_t * ray_direction
    normal = normalize(hit_point - closest_sphere.center)
    view_dir = -ray_direction
    reflection_dir = normalize(ray_direction - 2 * np.dot(ray_direction, normal) * normal)
    reflection_color = trace(hit_point, reflection_dir, scene, background, depth - 1)

    color = closest_sphere.color * (1 - closest_sphere.reflection) + reflection_color * closest_sphere.reflection
    return np.clip(color, 0, 1)

def render(scene, glass_sphere):
    aspect_ratio = WIDTH / HEIGHT
    screen = pygame.display.set_mode((WIDTH, HEIGHT))
    pygame.display.set_caption("Ray Tracing Glass Sphere")

    camera_origin = np.array([0, 0, -3])
    background = create_background()
    running = True
    needs_redraw = True

    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False

        keys = pygame.key.get_pressed()
        if keys[pygame.K_w]:
            glass_sphere.center[1] -= 1
            needs_redraw = True
        if keys[pygame.K_s]:
            glass_sphere.center[1] += 1
            needs_redraw = True
        if keys[pygame.K_a]:
            glass_sphere.center[0] -= 1
            needs_redraw = True
        if keys[pygame.K_d]:
            glass_sphere.center[0] += 1
            needs_redraw = True

        if needs_redraw:
            print("Ray tracing started...")
            start_time = time.time()
            framebuffer = np.zeros((HEIGHT, WIDTH, 3))
            for y in range(HEIGHT):
                for x in range(WIDTH):
                    u = (x / WIDTH) * 2 - 1
                    v = (y / HEIGHT) * 2 - 1
                    u *= aspect_ratio
                    ray_direction = normalize(np.array([u, v, 1]))
                    framebuffer[y, x] = trace(camera_origin, ray_direction, scene, background)
            
            framebuffer = (framebuffer * 255).astype(np.uint8)
            pygame.surfarray.blit_array(screen, np.transpose(framebuffer, (1, 0, 2)))
            pygame.display.flip()
            needs_redraw = False
            print(f"Rendering time: {time.time() - start_time:.2f} seconds")
    
    pygame.quit()

if __name__ == "__main__":
    glass_sphere = Sphere(center=[0, 0, 3], radius=1, color=[0.6, 0.7, 1.0], reflection=0.5)
    scene = [glass_sphere]
    render(scene, glass_sphere)