raytracer / rays.ml

(* ray.ml: a simple raytracer *)

type vector3 = float * float * float
type color = int * int * int
type normal = Normal of vector3
type light = {light_center:vector3; light_color:color}
type sphere = {center:vector3; radius:float; color:color; sph_id:int}
type plane = {plane_normal:normal; offset:float; plane_color:color;
	      plane_id:int}
type shape = Sphere of sphere | Plane of plane
type intersection = {location:vector3;normal:normal;obj:shape}
type camera = {origin:vector3; out:normal; up:normal; width:int;
	       height:int; xangle:float; yangle:float}

(* color functions and constants *)

let rgb_black : color = (0,0,0)
let rgb_white : color = (255,255,255)
let rgb_red : color = (255,0,0)
let rgb_green : color = (0,255,0)
let rgb_blue : color = (0,0,255)

let rgb_r ((r,g,b):color) = r
let rgb_g ((r,g,b):color) = g
let rgb_b ((r,g,b):color) = b
let rgb_coeff ((r,g,b):color) (coeff:float) = 
  let mult i = int_of_float ((float_of_int i) *. coeff) in
    (mult r, mult g, mult b)
let rgb_mix ((r1,g1,b1):color) ((r2,g2,b2):color) : color =
  let mix a b = (a*b)/255 in
    (mix r1 r2, mix g1 g2, mix b1 b2)
let rgb_sum ((r1,g1,b1):color) ((r2,g2,b2):color) : color =
  (r1+r2,g1+g2,b1+b2)

let fabs (f:float) : float = if f > 0.0 then f else 0.0 -. f
let sq x = x *. x 

(* vector3 functions *)

let v3_origin = (0.0,0.0,0.0)
let v3_l2distsq ((x1,y1,z1):vector3) ((x2,y2,z2):vector3) =
  let xdiff = x1 -. x2 in
  let ydiff = y1 -. y2 in
  let zdiff = z1 -. z2 in
    (sq xdiff) +. (sq ydiff) +. (sq zdiff)
let v3_l2dist p1 p2 = sqrt (v3_l2distsq p1 p2)
let v3_magnitudesq (v3:vector3) : float = v3_l2distsq v3 v3_origin
let v3_magnitude (v3:vector3) : float = sqrt (v3_magnitudesq v3)
let v3_mul ((x,y,z):vector3) (f:float) : vector3 = (x*.f, y*.f, z*.f)
let v3_div (v3:vector3) (f:float) : vector3 =
  let recip = 1.0 /. f in
    v3_mul v3 recip
let v3_neg ((x,y,z):vector3) : vector3 = (-.x, -.y, -.z)
let v3_add ((x1,y1,z1):vector3) ((x2,y2,z2):vector3) : vector3 =
  (x1+.x2, y1+. y2, z1 +. z2)
let v3_sub (v3a:vector3) (v3b:vector3) : vector3 = v3_add v3a (v3_neg v3b)
let v3_dot ((x1,y1,z1):vector3) ((x2,y2,z2):vector3) : float =
  (x1*.x2) +. (y1*.y2) +. (z1*.z2)
let v3_cross ((a1,a2,a3):vector3) ((b1,b2,b3):vector3) : vector3 =
  (a2*.b3 -. a3*.b2, a3*.b1 -. a1*.b3, a1*.b2 -. a2*.b1)  
let v3_norm (v3:vector3) : normal = Normal (v3_div v3 (v3_magnitude v3))

(* shape intersection handlers *)

let compare_shapes (s1:shape) (s2:shape) : bool =
  match (s1,s2) with
    | (Sphere sph1, Sphere sph2) -> sph1.sph_id = sph2.sph_id
    | (Plane p1, Plane p2) -> p1.plane_id = p2.plane_id
    | _ -> false

let shape_color (s:shape) : color =
  match s with
      Sphere sph -> sph.color
    | Plane p -> p.plane_color


let sphere_intersect (sph:sphere) (src:vector3) (dir:normal) :
    intersection option =
  let center = v3_sub sph.center src in
  let l = match dir with Normal n -> n in
  let lcdot = v3_dot l center in
  let lcdot2 = sq lcdot in
  let magc2 = v3_magnitudesq center in
  let r2 = sq sph.radius in
  let radical = lcdot2 -. magc2 +. r2 in
  let solution () =
    let sqrtrad = sqrt radical in
    let plusmag = lcdot +. sqrtrad in
    let minusmag = lcdot -. sqrtrad in
    let plus = v3_add src (v3_mul l plusmag) in
    let minus = v3_add src (v3_mul l minusmag) in
    let plus_closer = (v3_l2distsq src plus) < (v3_l2distsq src minus) in
    let loc = if plus_closer then plus else minus in
    let norm = v3_norm (v3_sub loc sph.center) in
      Some {location=loc;normal=norm;obj=(Sphere sph)}
  in
    if radical >= 0.0 then solution() else None

let plane_intersect (p:plane) (src:vector3) (dir:normal) : intersection option =
  let l = match dir with Normal n -> n in
  let n = match p.plane_normal with Normal n -> n in
  let denom = v3_dot l n in
    if (fabs denom) < 0.000000001 then None else (
      let t = ((-.p.offset) -. (v3_dot src n)) /. denom in
      let loc = v3_add (v3_mul l (t)) src in
      let norm = p.plane_normal in
	if t > 0.0 then
	  (Some {location=loc; normal=norm; obj=Plane p})
	else 
	  None)

let shape_intersect (src:vector3) (dir:normal) (shape:shape) :
    intersection option =
  match shape with
    | Sphere s -> sphere_intersect s src dir
    | Plane p -> plane_intersect p src dir

let closest_intersection_excluded (source_isect:intersection option) 
    (shapes:shape list) (src:vector3) (dir:normal) : intersection option =
  let isects =
    let filter_isect acc s =
      match shape_intersect src dir s with
	  None -> acc
	| Some isect -> (
	    match source_isect with
	      | None -> isect::acc
	      | Some orig_isect -> (
		  if (compare_shapes orig_isect.obj s) &&
		    ((v3_l2distsq orig_isect.location isect.location)
		     < 0.0000001) then acc else isect::acc))
    in
      List.fold_left filter_isect [] shapes
  in
  let choose_closest (prev:intersection*float)
			(next:intersection) =
    let dist = v3_l2distsq src next.location in
    let (best,best_dist) = prev in
      if dist < best_dist then (next,dist) else prev
  in
  let extract_isect (best,best_dist) = best in
    match isects with
      | [] -> None
      | _ -> Some (
	  let hd = List.hd isects in
	  let hd_dist = v3_l2distsq hd.location src in
	  let best_pair = 
	    List.fold_left choose_closest (hd,hd_dist) (List.tl isects)
	  in
	    extract_isect best_pair)

let closest_intersection = closest_intersection_excluded None

(* camera functions *)

let build_rays (cam:camera) : normal list =
  let d_theta_x = cam.xangle /. (float_of_int cam.width) in
  let d_theta_y = cam.yangle /. (float_of_int cam.height) in
  let halfx = cam.xangle /. 2.0 in
  let halfy = cam.yangle /. 2.0 in
  let extract_norm n = match n with Normal v -> v in
  let out = extract_norm cam.out in
  let up = extract_norm cam.up in
  let right = v3_cross out up in
  let rec loop ret x y = 
    match x with
      | 0 -> (
	  match y with
	    | 0 -> ret
	    | _ -> loop ret cam.width (y-1))
      | _ -> (
	  let xtheta = (float_of_int x)*.d_theta_x -. halfx in
	  let ytheta = (float_of_int y)*.d_theta_y -. halfy in
	  let right_comp = v3_mul right (sin xtheta) in
	  let up_comp = v3_mul up (sin ytheta) in
	  let next = v3_norm (v3_add (v3_add out up_comp) right_comp) in
	    loop (next::ret) (x-1) y)
  in
    loop [] cam.width cam.height

let illuminated (isect:intersection) (shapes:shape list) (light:light) : bool =
  let dir = v3_norm (v3_sub light.light_center isect.location) in
  let closest = 
    closest_intersection_excluded (Some isect) shapes isect.location dir
  in
    match closest with
	None -> true
      | Some isect2 -> (
	  let isectd = v3_l2dist isect2.location isect.location in
	  let lightd = v3_l2dist light.light_center isect.location in
	    isectd > lightd)

let handle_illuminated (isect:intersection) (light:light) : color =
  let light_dir = v3_norm (v3_sub light.light_center isect.location) in
  let light_norm = match light_dir with Normal n -> n in
  let isect_norm = match isect.normal with Normal n -> n in
  let lambertian_coeff = max (v3_dot light_norm isect_norm) 0.0 in
  let light_in = rgb_coeff light.light_color lambertian_coeff in
    rgb_mix light_in (shape_color isect.obj)

let handle_ambient (isect:intersection) (ambience:float) : color =
  rgb_coeff (shape_color isect.obj) ambience

let render_ray (src:vector3) (shapes:shape list)
    (lights:light list) (bg:color) (ambience:float) (ray:normal) : color =
  let closest = closest_intersection shapes src ray in
    match closest with
	None -> bg
      | Some isect -> (
	  let handle_light color light =
	    let is_illum = illuminated isect shapes light in
	      if is_illum then (
		rgb_sum color (handle_illuminated isect light))
	      else color
	  in
	    rgb_sum (handle_ambient isect ambience)
	      (List.fold_left handle_light rgb_black lights))
	  
let render_scene (camera:camera) (shapes:shape list) (lights:light list)
    (bg:color) (ambience:float) : color list =
  let rays = build_rays camera in
  let f = render_ray camera.origin shapes lights bg ambience in
  let rec loop ret rem =
    match rem with
      | [] -> List.rev ret
      | hd::rest -> loop ((f hd)::ret) rest
  in
    loop [] rays

let dump_bitmap (out:out_channel) (camera:camera) (colors:color list)  : unit =
  let n_pixels = List.length colors in
  let image_size = n_pixels * 3 in
  let header_size = 54 in
  let total_size = image_size + header_size in
  let out_byte (i:int) = output_byte out i in
  let out_int (i:int) =
    out_byte i;
    out_byte (i asr 8);
    out_byte (i asr 16);
    out_byte (i asr 24)
  in
  let out_short (i:int) =
    out_byte i;
    out_byte (i asr 8);
  in    
  let out_color ((r,g,b):color) =
    let out_c (i:int) = out_byte (min 255 (max 0 i)) in
    out_c b;
    out_c g;
    out_c r;
  in
  let write_header () =
    output_char out 'B';
    output_char out 'M';
    out_int total_size;
    out_int 0; (* used by BMP creating application *)
    out_int 54; (* offset to bitmap data *)
    out_int 40; (* DIB header size -- using Windows V3 version *)
    out_int camera.width;
    out_int camera.height;
    out_short 1; (* number of color planes *)
    out_short 24; (* color bits *)
    out_int 0; (* compression mode - no compression *)
    out_int image_size;
    out_int (1 lsl 24); (* horizontal resolution, arbitrarily chosen *)
    out_int (1 lsl 24); (* vertical resolution, arbitrarily chosen *)
    out_int 0; (* default number of colors in palette *)
    out_int 0; (* number of important colors, 0 = all colors important *)
  in
  let write_data () =
    assert (camera.width mod 4 = 0);
    assert (camera.height mod 4 = 0);
    List.iter out_color colors
  in
    write_header();
    write_data()

let write_bitmap (filename:string) (camera:camera) (colors:color list) : unit =
  let out = open_out_bin filename in
    dump_bitmap out camera colors;
    close_out out

let shapes = [Sphere {center=(10.0,0.0,0.0); radius=2.0; color=rgb_green;
		      sph_id=0};
	      Plane {plane_normal=v3_norm (0.0,1.0,0.0); offset=(5.0);
		    plane_color=rgb_blue; plane_id=1};
	      Plane {plane_normal=v3_norm (-.1.0,0.0,0.5); offset=(15.0);
		     plane_color=(0,255,255); plane_id=2};
	      Plane {plane_normal=v3_norm (0.0,0.0,-1.0); offset=10.0;
		     plane_color=(255,255,0); plane_id=3}]
let lights = [{light_center=(0.0,10.0,0.0); light_color=(200,200,200)};
	      {light_center=(3.0,0.0,-.4.0);light_color=(100,255,100)}]
let camera = {origin=v3_origin; out=v3_norm (1.0,0.0,0.0);
	      up=v3_norm (0.0,1.0,0.0); width=1024; height=768;
	      xangle=2.0; yangle=1.5}

let colors = render_scene camera shapes lights rgb_red 0.07
let out = write_bitmap "rendered.bmp" camera colors

let dtest =
  assert (fabs ((v3_l2dist (3.0, 0.0, 0.0) (0.0, 4.0, 0.0)) -. 5.0) < 0.00001)
Tip: Filter by directory path e.g. /media app.js to search for public/media/app.js.
Tip: Use camelCasing e.g. ProjME to search for ProjectModifiedEvent.java.
Tip: Filter by extension type e.g. /repo .js to search for all .js files in the /repo directory.
Tip: Separate your search with spaces e.g. /ssh pom.xml to search for src/ssh/pom.xml.
Tip: Use ↑ and ↓ arrow keys to navigate and return to view the file.
Tip: You can also navigate files with Ctrl+j (next) and Ctrl+k (previous) and view the file with Ctrl+o.
Tip: You can also navigate files with Alt+j (next) and Alt+k (previous) and view the file with Alt+o.