unix4lyfe git - the-worst-raytracer/blob - raytracer.py

   1 #!/usr/bin/env python
   2 """
   3 The worst raytracer.
   4 Emil Mikulic <emikulic@gmail.com> was here 2012.
   5 """
   6 import numpy as np
   7 from accidental_complexity import show
   8 # this file contains mostly essential complexity
   9
  10 def vec(x, y, z):
  11   return np.array((x, y, z), dtype=np.float)
  12
  13 def color(r, g, b):
  14   return np.reshape(vec(r,g,b), (1,1,3))
  15
  16 def dot(a, b):
  17   return (a * b).sum(axis=2)
  18
  19 def mag(a):
  20   return np.sqrt(np.square(a).sum(axis=2))
  21
  22 def depthwise(a):
  23   return np.expand_dims(a, axis=2)
  24
  25 def norm(a):
  26   return a / depthwise(mag(a))
  27
  28 def reflect(i, n):
  29   c = n * depthwise(dot(n, -i))
  30   return 2*c + i
  31
  32 def diffuse(n, l):
  33   return depthwise(dot(n, l).clip(0, 1))
  34
  35 def specular(r, l, power):
  36   return pow(diffuse(r, l), power)
  37
  38 def lerp(ar, a, b, u, v):
  39   "[u,v] is mapped to [a,b]"
  40   return a + (b-a) * (ar.clip(u,v) - u) / (v-u)
  41
  42 class Scene:
  43   def __init__(self):
  44     self.light1 = vec(20,10,-10)
  45     self.light1_col = color(1,.5,.3)
  46     self.light2 = vec(-5,5,5)
  47     self.light2_col = color(1,1,1) - self.light1_col
  48     self.sky_col = color(.2, .3, .6)
  49     self.sph_center = vec(-.5,0,0)
  50     self.sph_rad = .5
  51     self.sph_col = color(1,1,1)
  52     self.camera = vec(0,0,-5)
  53     self.ground_col1 = color(.4,.4,.4)
  54     self.ground_col2 = color(.8,.8,.8)
  55     self.plane_n = vec(0,1,0)
  56     self.plane_amb = 0.02
  57     self.plane_dif = 0.9
  58     self.plane_height = -0.5
  59     self.fog_near = 0
  60     self.fog_far = 40
  61
  62   def trace_bg(self, camera, ray):
  63     # background
  64     plane_depth = (self.plane_height - camera[:,:,1]) / ray[:,:,1]
  65     plane_p = camera + depthwise(plane_depth) * ray
  66     plane_l1 = norm(self.light1 - plane_p)
  67     plane_l2 = norm(self.light2 - plane_p)
  68     plane_xpos = plane_p[:,:,0] - np.floor(plane_p[:,:,0])
  69     plane_zpos = plane_p[:,:,2] - np.floor(plane_p[:,:,2])
  70     checkers = depthwise(np.logical_xor(plane_xpos < 0.5, plane_zpos < 0.5))
  71
  72     # trace shadows (light 1)
  73     ec = plane_p - self.sph_center
  74     b = 2 * dot(plane_l1, ec)
  75     c = dot(ec, ec) - np.square(self.sph_rad)
  76     det = b*b - 4*c
  77     depth = (-b - np.sqrt(det)) / 2
  78     illum_l1 = np.where(depthwise(depth > 0), 0, self.light1_col)
  79
  80     # light 2
  81     b = 2 * dot(plane_l2, ec)
  82     det = b*b - 4*c
  83     depth = (-b - np.sqrt(det)) / 2
  84     illum_l2 = np.where(depthwise(depth > 0), 0, self.light2_col)
  85
  86     # shade plane
  87     plane_col = np.where(checkers, self.ground_col1, self.ground_col2) * (
  88         self.plane_amb + self.plane_dif * (
  89           diffuse(self.plane_n, plane_l1) * illum_l1 +
  90           diffuse(self.plane_n, plane_l2) * illum_l2))
  91
  92     fog_dist = mag(plane_p - vec(0, self.plane_height, 0))
  93     fog = lerp(fog_dist, 1, 0, self.fog_near, self.fog_far)
  94     plane_col *= depthwise(np.square(fog))
  95
  96     sky = self.sky_col * depthwise(ray[:,:,1])
  97     return np.where(depthwise(plane_depth > 0), plane_col, sky)
  98
  99   def render(self, w, h):
 100     # set up scene
 101     sph_amb = 0.02
 102     sph_diff = 0.05
 103     sph_spec = 1.1
 104     sph_refl = 0.5
 105     sph_shine = 20
 106
 107     # points on the screen plane (z = 0)
 108     aspect = float(w) / float(h)
 109     scr_x = np.linspace(-1 + 1./w, 1 - 1./w, w) * aspect
 110     scr_y = np.linspace(-1 + 1./h, 1 - 1./h, h) * -1
 111     scr = np.zeros((h,w,3), dtype=np.float)
 112     scr[:,:,0] = np.reshape(scr_x, (1, w))
 113     scr[:,:,1] = np.reshape(scr_y, (h, 1))
 114
 115     ray = norm(scr - self.camera)
 116
 117     # sphere
 118     ec = self.camera - self.sph_center
 119     b = 2 * dot(ray, ec)
 120     c = ec.dot(ec) - np.square(self.sph_rad)
 121     det = b*b - 4*c
 122     sph_depth = (-b - np.sqrt(det)) / 2
 123
 124     # sphere color
 125     sph_p = self.camera + depthwise(sph_depth) * ray
 126     sph_l1 = norm(self.light1 - sph_p)
 127     sph_l2 = norm(self.light2 - sph_p)
 128     sph_n = (sph_p - self.sph_center) / self.sph_rad
 129     sph_r = norm(reflect(ray, sph_n))
 130     r_col = self.trace_bg(sph_p, sph_r)
 131
 132     sph = self.sph_col * (sph_amb +
 133         sph_diff * diffuse(sph_n, sph_l1) * self.light1_col +
 134         sph_diff * diffuse(sph_n, sph_l2) * self.light2_col +
 135         sph_spec * specular(sph_r, sph_l1, sph_shine) * self.light1_col +
 136         sph_spec * specular(sph_r, sph_l2, sph_shine) * self.light2_col +
 137         r_col * sph_refl)
 138
 139     bg = self.trace_bg(np.reshape(self.camera, (1,1,3)), ray)
 140     return np.where(depthwise(sph_depth > 0), sph, bg)
 141
 142 show(Scene().render(800, 600))
 143 # vim:set ts=2 sw=2 et: