Ray tracing in OCaml (revisited)

When I posted a simple ray tracer written in OCaml to this newgroup a few
months back, Thierry Berger-Perrin responded with a C++ ray tracer. The
C++
implementation used skip arrays to improve cache coherence but remained
concise (167 LOC or 4,708 non-whitespace bytes):

  http://ompf.org/ray/sphereflake/

I always wondered how the performance of my OCaml compared to the
optimised
C++ but was never able to compare them because Thierry's code didn't do
the
same thing. So I just spent a couple of hours translating it back into
OCaml and the results are very interesting.

The resulting program is 108 LOC (2,557 non-whitespace bytes) of OCaml and
takes 6.12s to rendering a 1024x1024 image of a scene containing 66430
spheres with 2x2 supersampling compared to 3.83s for the C++:

Spheres  C++ OCaml
     10 1.1s 2.6s
     91 2.0s 3.8s
    820 2.7s 4.7s
  7,381 3.4s 5.5s
 66,430 3.9s 6.1s

This OCaml is really medium performance and medium brevity. It could be
made
more concise by removing the specialised shadow ray intersection,
the "exists" function implementation (e.g. using the built-in List.exists
instead) and so on. Performance could be improved by inlining the vector
arithmetic in ray_sphere and reducing unboxing to alleviate the GC.

let pi = 4. *. atan 1. and delta = sqrt epsilon_float

type vec = {x: float; y: float; z: float}
let ( *| ) s r = {x = s *. r.x; y = s *. r.y; z = s *. r.z}
let ( +| ) a b = {x = a.x +. b.x; y = a.y +. b.y; z = a.z +. b.z}
let ( -| ) a b = {x = a.x -. b.x; y = a.y -. b.y; z = a.z -. b.z}
let dot a b = a.x *. b.x +. a.y *. b.y +. a.z *. b.z
let unitise r = 1. /. sqrt(dot r r) *| r
let cross u v = { x = u.y *. v.z -. u.z *. v.y;
                  y = u.z *. v.x -. u.x *. v.z;
                  z = u.x *. v.y -. u.y *. v.x }

let ray_sphere dir v radius =
  let b = dot v dir in
  let disc = b *. b -. dot v v +. radius *. radius in
  if disc < 0. then infinity else
    let disc = sqrt disc in
    let t2 = b +. disc in
    if t2 < 0. then infinity else
      let t1 = b -. disc in
      if t1 > 0. then t1 else t2

let ray_sphere' orig dir center radius =
  let v = center -| orig in
  let b = dot v dir in
  let disc = b *. b -. dot v v +. radius *. radius in
  if disc < 0. then false else b +. sqrt disc >= 0.

let intersect dir =
  let rec aux (l, _ as first) (bc, br, sc, sr, scenes) =
    let l' = ray_sphere dir bc br in
    if l' >= l then first else
      let l' = ray_sphere dir sc sr in
      let first = if l' >= l then first else l', unitise(l' *| dir -| sc)
in
      Array.fold_left aux first scenes in
  aux (infinity, {x=0.; y=0.; z=0.})

let exists f a =
  try
    for i=0 to Array.length a - 1 do
      if f a.(i) then raise Exit
    done;
    false
  with Exit -> true

let intersect' orig dir =
  let rec aux (bc, br, sc, sr, scenes) =
    ray_sphere' orig dir bc br &&
      (ray_sphere' orig dir sc sr || exists aux scenes) in
  aux

let rec ray_trace light dir scene =
  let lambda, normal = intersect dir scene in
  if lambda = infinity then 0. else
    let g = dot normal light in
    if g >= 0. then 0. else
      let p = lambda *| dir +| delta *| normal in
      if intersect' p (-1. *| light) scene then 0. else -. g

let basis v =
  let n = unitise v in
  let b1 =
    if n.x *. n.x <> 1. && n.y *. n.y <> 1. && n.z *. n.z <> 1. then
      if n.y *. n.y > n.x *. n.x then
        if n.y *. n.y > n.z *. n.z then
          {n with y = -. n.y}
        else
          {n with z = -. n.z}
      else
        if n.z *. n.z > n.x *. n.x then
        {n with z = -. n.z}
      else
        {n with x = -. n.x}
      else
        {x=n.z; y=n.x; z=n.y} in
  let b2 = cross n b1 in
  n, cross n b2, b2

let rec create level c d r =
  let n = c, 2. *. r, c, r, [||] in
  if level <= 1 then n else
    let up, b1, b2 = basis d in
    let nr = r /. 3. and daL = 2. *. pi /. 6. and daU = 2. *. pi /. 3. in
    let a = ref 0. and cs = Array.make 9 n in
    for i=0 to 5 do
      let ndir = unitise(-0.2 *| d +| sin !a *| b1 +| cos !a *| b2) in
      cs.(i) <- create (level - 1) (c +| (r +. nr) *| ndir) ndir nr;
      a := !a +. daL;
    done;
    a := !a -. daL /. 3.;
    for i=0 to 2 do
      let ndir = unitise(0.6 *| d +| sin !a *| b1 +| cos !a *| b2) in
      cs.(6+i) <- create (level - 1) (c +| (r +. nr) *| ndir) ndir nr;
      a := !a +. daU;
    done;
    c, 2. *. r, c, r, cs

let level, n = match Sys.argv with
    [|_; level; n|] -> int_of_string level, int_of_string n
  | _ -> 6, 1024

let light = unitise {x = -0.5; y = -0.65; z = 0.9} and ss = 2
let scene =
  create level {x=0.; y=0.; z=4.5} (unitise{x=0.25; y=1.; z= -0.5}) 1.;;
Printf.printf "P5\n%d %d\n255\n" n n;;
for y = n - 1 downto 0 do
  for x = 0 to n - 1 do
    let g = ref 0. in
    for dx = 0 to ss - 1 do
      for dy = 0 to ss - 1 do
        let aux x d = float x -. float n /. 2. +. float d /. float ss in
        let dir = unitise {x = aux x dx; y = aux y dy; z = float n} in
        g := !g +. ray_trace light dir scene
      done;
    done;
    let g = 0.5 +. 255. *. !g /. float (ss*ss) in
    Printf.printf "%c" (char_of_int (int_of_float g))
  done;
done

-- 
Dr Jon D Harrop, Flying Frog Consultancy
Objective CAML for Scientists
http://www.ffconsultancy.com/products/ocaml_for_scientists