Ray Tracer

The following 222-line OCaml program generates and ray traces a sphere flake:

(* Generic maths *)
let delta = sqrt epsilon_float and pi = 4. *. atan 1.
let sqr x = x *. x
let sin t = sin(pi *. t /. 180.) and cos t = cos(pi *. t /. 180.)
let clamp l u x = if x < l then l else if x > u then u else x
type vec = float array
let vec3 x y z = [| x; y; z; 1. |]
let zero = vec3 0. 0. 0.
let ( *| ) s r = [| s *. r.(0); s *. r.(1); s *. r.(2); 1. |]
let ( +| ) a b = [| a.(0) +. b.(0); a.(1) +. b.(1); a.(2) +. b.(2); 1. |]
let ( -| ) a b = [| a.(0) -. b.(0); a.(1) -. b.(1); a.(2) -. b.(2); 1. |]
let dot a b = a.(0) *. b.(0) +. a.(1) *. b.(1) +. a.(2) *. b.(2)
let length2 r = dot r r
let length r = sqrt(length2 r)
let unitise r = (1. /. length r) *| r
let init f = Array.init 4 (fun i -> Array.init 4 (f i))
let translate r =
  init (fun i j -> if i = j then 1. else if j = 3 then r.(i) else 0.)
let identity = translate (vec3 0. 0. 0.)
let sum4 f = f 0 +. f 1 +. f 2 +. f 3
let mat_mul a b = init (fun i j -> sum4 (fun k -> a.(i).(k) *. b.(k).(j)))
let transform m r = Array.init 4 (fun i -> dot m.(i) r +. m.(i).(3))
let rot_x t = let cos, sin = cos t, sin t in [|[|    1.;    0.;    0.; 0. |];
                                               [|    0.;   cos;   sin; 0. |];
                                               [|    0.; -.sin;   cos; 0. |];
                                               [|    0.;    0.;    0.; 1. |]|]
let rot_y t = let cos, sin = cos t, sin t in [|[|   cos;    0.;   sin; 0. |];
                                               [|    0.;    1.;    0.; 0. |];
                                               [| -.sin;    0.;   cos; 0. |];
                                               [|    0.;    0.;    0.; 1. |]|]
let rot_z t = let cos, sin = cos t, sin t in [|[|   cos;   sin;    0.; 0. |];
                                               [| -.sin;   cos;    0.; 0. |];
                                               [|    0.;    0.;    1.; 0. |];
                                               [|    0.;    0.;    0.; 1. |]|]

(* Ray tracing primitives *)
type material = { color: float * float * float; shininess: float }
type sphere = { center: vec; radius: float }
type obj = Sphere of sphere * material | Group of sphere * obj list
type ray = { origin: vec; direction: vec }

(* Ray-sphere intersection *)
let ray_sphere ray sphere =
  let v = sphere.center -| ray.origin in
  let b = dot v ray.direction in
  let disc = sqr b -. length2 v +. sqr sphere.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

(* Find the first intersection of the given ray with the given set of
   objects. *)
let intersect ray scene =
  let rec of_scene ((l, _, _) as first) = function
      Sphere (sphere, material) ->
        let l' = ray_sphere ray sphere in
        if l' >= l then first else (* No nearer *)
          let normal () =
            unitise (ray.origin +| l' *| ray.direction -| sphere.center) in
          l', normal, material (* Replace with nearer intersection *)
    | Group (bound, scenes) ->
        let l' = ray_sphere ray bound in (* Cull if possible *)
        if l' >= l then first else List.fold_left of_scene first scenes in
  let first =
    if ray.direction.(1) < 0. then
      (* Floor *)
      let l = -. ray.origin.(1) /. ray.direction.(1) in
      let r = ray.origin +| l *| ray.direction in
      let x = length [|r.(0); 0.; r.(2); 1.|] in
      let g = 0.5 *. (1. +. cos(180. *. x)) *. exp(-. 0.05 *. x *. x) in
      l, (fun () -> vec3 0. 1. 0.), { color = g, g, g; shininess = 0. }
    else
      infinity, (fun () -> zero), { color = 0., 0.1, 0.2; shininess = 0. } in
  of_scene first scene

(* Trace a single ray *)
let rec ray_trace weight light ray scene = match intersect ray scene with
    lambda, normal, material ->
      if lambda = infinity then 0., 0., 0. else
        let n = normal () in
        let o = ray.origin +| lambda *| ray.direction +| delta *| n in
        (* Recursively examine specular reflections *)
        let r, g, b =
          if weight < 0.5 then 0., 0., 0. else
            let d = ray.direction in
            let ray = { origin = o; direction = d -| (2. *. dot d n) *| n } in
            let s = material.shininess in
            match ray_trace (weight *. s) light ray scene with
              r, g, b -> s *. r, s *. g, s *. b in
        (* Calculate the final color, taking account of shadows, specular and
           diffuse reflection. *)
        let f = max 0. (-. dot n light) in
        let s =
          match intersect { origin = o; direction = zero -| light } scene with
            l, _, _ when l = infinity -> f
          | _ -> 0. in
        match material.color with
          r', g', b' -> (s +. r) *. r', (s +. g) *. g', (s +. b) *. b'

let () =
  (* Get the level of detail from the command line *)
  let level = match Sys.argv with
      [| _; l |] -> int_of_string l
    | _ -> 6 in

  (* Initialise OpenGL and glut *)
  let argv = Glut.init Sys.argv in
  Glut.initDisplayMode ();
  let width = ref 768 and height = ref 768 in
  Glut.initWindowSize !width !height;
  ignore (Glut.createWindow "Ray trace");

  (* Define the scene *)
  let light = unitise (vec3 (-1.) (-3.) 2.) in
  let s = 1. /. 3. in (* Radius ratio of one sphere to the next *)
  let count = ref 0 in
  let scene =
    let rec aux level r m =
      incr count;
      let sphere = { center = transform m zero; radius = r } in
      let material =
        let t = 15. *. float level in
        { color = sin t, sin(t +. 60.), sin(t +. 120.); shininess = 0.5 } in
      let obj = Sphere (sphere, material) in
      if level = 1 then obj else begin
        let objects = ref [] in

        for i = 0 to 2 do
          let m = mat_mul m (rot_y (30. +. 120. *. float i)) in
          let m = mat_mul m (rot_z 45.) in
          let m = mat_mul m (translate (vec3 0. ((1. +. s) *. r) 0.)) in
          objects := aux (level - 1) (s *. r) m :: !objects
        done;

        for i = 0 to 5 do
          let m = mat_mul m (rot_y (60. *. float i)) in
          let m = mat_mul m (rot_z 110.) in
          let m = mat_mul m (translate (vec3 0. ((1. +. s) *. r) 0.)) in
          objects := aux (level - 1) (s *. r) m :: !objects
        done;

        (* Bounding radius *)
        let r' =
          List.fold_left
            (fun r -> function Sphere (sph', _) | Group (sph', _) ->
               max r (length (sph'.center -| sphere.center) +. sph'.radius))
            r !objects in
        
        Group ({ sphere with radius = r' }, obj :: !objects)
      end in
    aux level 1. (translate (vec3 0. 1. 0.)) in
  Printf.printf "Created %d objects\n" !count;

  let x, y = ref 0, ref 0 in
  let pixels = ref [] in
  let last = ref 0. and started = ref 0. in

  (* Render incrementally when idle *)
  let idle () =
    if !y >= !height then begin
      Printf.printf "Took: %fs\n" (Sys.time () -. !started);
      started := Sys.time ();
      Glut.idleFunc None;
      Glut.postRedisplay ()
    end else begin
      let ray =
        let w, h = float !width, float !height in
        let o = vec3 0. 2. (-6.) in
        let x, y = float !x -. 0.5 *. w, float !y -. 0.5 *. h in
        let d = unitise (vec3 x y (max w h)) in
        let d = transform (rot_x (-9.)) d in
        { origin = o ; direction = d } in

      pixels := (!x, !y, ray_trace 1. light ray scene) :: !pixels;

      x := (1 + !x) mod !width;
      if !x = 0 then incr y;

      if Sys.time () > !last +. 0.2 then begin
        last := Sys.time ();
        Glut.postRedisplay ();
      end;
    end in

  (* Rerender when the window resizes *)
  let reshape ~w ~h =
    width := w; height := h;
    GlDraw.viewport 0 0 w h;
    x := 0; y :=0;
    pixels := [];
    Glut.idleFunc (Some idle);
    Glut.postRedisplay () in

  (* Display the current render *)
  let display () =
    GlMat.mode `projection;
    GlMat.load_identity ();
    GluMat.ortho2d ~x:(0., float !width) ~y:(0., float !height);
    GlMat.mode `modelview;
    GlMat.load_identity ();
    GlClear.color (0., 0.1, 0.2);
    GlClear.clear [ `color ];
    GlDraw.begins `points;
    let aux (x, y, (r, g, b)) =
      let f = clamp 0. 1. in
      GlDraw.color (f r, f g, f b);
      GlDraw.vertex2 (float x, float y) in
    List.iter aux !pixels;
    GlDraw.ends ();
    Gl.finish ();
    Glut.swapBuffers ();
    last := Sys.time () in

  Glut.reshapeFunc reshape;
  Glut.displayFunc display;
  Glut.idleFunc (Some idle);
  Glut.keyboardFunc (fun ~key ~x ~y -> if key=27 then exit 0);

  Glut.mainLoop ()

This program can be compiled using:

ocamlopt -I +lablgl lablgl.cmxa lablglut.cmxa ray.ml -o ray

Mac OS X users may need to add -cclib "-framework Foundation" to the compile line.

Executing the program incrementally displays the current rendering using OpenGL:

$ ./ray
Created 66430 objects
Took: 8.120000s

This will display the following result:

By default, 6 levels of spheres are used. This can be controlled by specifing an optional integer command-line argument:

$ ./ray 3
Created 91 objects
Took: 4.770000s

This will display the following result:

Large numbers of objects can be rendered efficiently by this ray tracer due to the use of hierarchical spherical bounding volumes. Specifically, a scene is stored as a tree, each node of which is either a single sphere or a group with a spherical bound and a list of child nodes.

When tracing a ray, large clusters of objects can be culled from the search when the ray fails to intersect the bound of those objects. Consequently, the time taken t to render a scene scales very well with the number of objects n in the scene:

The scene tree is defined and manipulated very efficiently by OCaml, both in terms of the amount of code required and the run-time performance.