Raytrace using vector objects

ray-object.lua

local sqrt = math.sqrt
local huge = math.huge

local delta = 1
while delta * delta + 1 ~= 1 do
  delta = delta * 0.5
end


-- vectors -------------------------------------------------------------------

local vector = {}
vector.__index = vector
function vector:new(x, y, z)
  return setmetatable({x=x, y=y, z=z}, self)
end


local function dot_product(a, b)
  return a.x*b.x + a.y*b.y + a.z*b.z
end


function vector:length()
  return sqrt(dot_product(self, self))
end


function vector.__add(a, b)
  -- assume both are vectors
  return vector:new(a.x+b.x, a.y+b.y, a.z+b.z)
end


function vector.__sub(a, b)
  -- assume both are vectors
  return vector:new(a.x-b.x, a.y-b.y, a.z-b.z)
end


function vector.__mul(a, b)
  -- assume either a or b is a number, we don't do a cross product here
  if type(b) == "number" then
    a, b = b, a
  end
  return vector:new(a*b.x, a*b.y, a*b.z)
end


function vector:normalise()
  return (1/self:length()) * self
end


-- rays ----------------------------------------------------------------------

local ray = {}
function ray:new(origin, direction)
  return {origin=origin, direction=direction}
end


-- spheres -------------------------------------------------------------------

local sphere = {}
sphere.__index = sphere
function sphere:new(centre, radius)
  return setmetatable({centre=centre, radius=radius}, self)
end


function sphere:intersection_distance(ray)
  local v = self.centre - ray.origin
  local b = dot_product(v, ray.direction)
  local r = self.radius
  local disc = r*r + b*b - dot_product(v, v)
  if disc < 0 then return huge end
  local d = sqrt(disc)
  local t2 = b + d
  if t2 < 0 then return huge end  -- sphere is behind us
  local t1 = b - d
  return t1 > 0 and t1 or t2
end


function sphere:intersect(ray, best)
  local distance = self:intersection_distance(ray)
  if distance < best.distance then
    best.distance = distance
    best.normal = (ray.origin + distance * ray.direction - self.centre):normalise()
  end
end


-- groups --------------------------------------------------------------------

local group = {}
group.__index = group
function group:new(bound)
  return setmetatable({bound=bound, children={}}, self)
end


function group:add(s)
  self.children[#self.children+1] = s
end


function group:intersect(ray, best)
  local distance = self.bound:intersection_distance(ray)
  if distance < best.distance then
    for _, c in ipairs(self.children) do
      c:intersect(ray, best)
    end
  end
end


-- tracing -------------------------------------------------------------------

local function ray_trace(light, camera, scene)
  local best = { distance = huge }
  scene:intersect(camera, best)
  if best.distance == huge then return 0 end  -- no object intersecting ray
  local g = dot_product(best.normal, light)
  if g >= 0 then return 0 end -- not facing the light

  -- check we're not in shadow
  local just_above = camera.origin + best.distance * camera.direction + delta * best.normal
  best.distance = huge
  scene:intersect(ray:new(just_above, -1 * light), best)
  if best.distance == huge then
    return -g
  else
    return 0  -- there was another object between us and the light
  end
end


-- create the scene ----------------------------------------------------------

local function create(level, centre, radius)
  local s = sphere:new(centre, radius)
  if level == 1 then return s end
  local gr = group:new(sphere:new(centre, 3*radius))
  gr:add(s)
  local rn = 3*radius/sqrt(12)
  for dz = -1,1,2 do
    for dx = -1,1,2 do
      gr:add(create(level-1, centre + rn * vector:new(dx, 1, dz), radius*0.5))
    end
  end
  return gr
end


-- main ----------------------------------------------------------------------

if arg[1] == "--help" then
  io.stderr:write("Usage: ", arg[-1], " ", arg[0], " [scene complexity (5)] [image size (512)] [anti-aliasing (4)] > output.pnm\n" )
  os.exit(0)
end

local level, n, ss = tonumber(arg[1]) or 5, tonumber(arg[2]) or 512, tonumber(arg[3]) or 4
local iss = 1/ss
local gf = 255/(ss*ss)

io.write(("P5\n%d %d\n255\n"):format(n, n))
local light = vector:new(-1, -3, 2):normalise()
local camera = vector:new(0, 0, -4)

local scene = create(level, vector:new(0, -1, 0), 1)

for y = n/2-1, -n/2, -1 do
  for x = -n/2, n/2-1 do
    local g = 0
    for d = y, y+.99, iss do
      for e = x, x+.99, iss do
        g = g + ray_trace(light, ray:new(camera, vector:new(e, d, n):normalise()), scene) 
      end
    end
    io.write(string.char(math.floor(0.5 + g*gf)))
  end
end