Saturday, 3 January 2009

Ray Tracing in Clojure

ANSI Common Lisp, by Paul Graham, is a great introduction to Lisp. Lots of compelling examples on run-length encoding (see previous post), poetry generation and ray-tracing.

This post looks at translating the example (from Chapter 9) into Clojure.

Ray-tracing is a very simple technique. From ACL:
"To generate a 3D image, we need to define at least four things: an eye, one or more light sources, a simulated world consisting of one or more surfaces, and a plane (the image plane) that serves as a window onto this world. The image we generate is the projection of the world onto a region of the image plane."

So how do we generate the pictures? That's pretty simple too, all we do is for every pixel in the image plane just trace the ray that the eye would see from there. Done.

We start off by defining some simple maths functions and an object to represent a point in 3D space.

(defn square [x] (* x x))

(defstruct point :x :y :z)

(defn magnitude [p]
(Math/sqrt (+ (square (:x p)) (square (:y p)) (square (:z p)))))

(defn unit-vector [p]
(let [d (magnitude p)]
(struct point (/ (:x p) d) (/ (:y p) d) (/ (:z p) d))))

(defn point-subtract [p1 p2]
(struct point 
(- (:x p1) (:x p2))
(- (:y p1) (:y p2))
(- (:z p1) (:z p2))))

(defn distance [p1 p2]
(magnitude (point-subtract p1 p2)))

(defn minroot [a b c]
(if (zero? a)
(/ (- c) b)
(let [disc (- (square b) (* 4 a c))]
(if (> disc 0)
(let [discroot (Math/sqrt disc)]
(min (/ (+ (- b) discroot) (* 2 a))
(/ (- (- b) discroot) (* 2 a))))))))

The original Lisp code mixed the point structure with individual values. I felt this made the code a bit ugly and hard to read, so in here we try to use the point structure as much as possible. (struct point 1 2 3) feels like clunky syntax, but I was unable to find anything better. Perhaps an alternative is to just use a plain vector / map? Or wait for the future and see if struct support improves?

Anyway, the code above is self explanatory, minroot is the big one and that's just a solver for the quadratic equation. function.

Next we need to define some of the environment. For this we'll fix the image plan between (0,0) and (300,300) and we'll just render spheres. Each surface has a grey-scale colour associated with it (a surface).

(def eye (struct point 150 150 200))

(defstruct surface :color)

(defstruct sphere :color :radius :centre)

(defn defsphere [point r c]
(struct sphere c r point))

(def world [(defsphere (struct point 150 150 -600) 250 0.32)
(defsphere (struct point 175 175 -300) 100 0.64)])

One thing I did find was that Clojure doesn't support the :include for structs that Common Lisp does. For this example, the world is a couple of spheres one smaller than the other and in front (and slightly brighter).

The following functions determine where a sphere gets hit with a ray from a specific source (in this case a point) and the surface normal of a hit.

(defn sphere-normal [s pt]
(let [c (:centre s)]
(unit-vector (point-subtract c pt))))

(defn sphere-intersect [s pt ray]
(let [c (:centre s)
n (minroot (+ (square (:x ray)) (square (:y ray)) (square (:z ray)))
(* 2 (+ (* (- (:x pt) (:x c)) (:x ray))
(* (- (:y pt) (:y c)) (:y ray))
(* (- (:z pt) (:z c)) (:z ray))))
(+ (square (- (:x pt) (:x c)))
(square (- (:y pt) (:y c)))
(square (- (:z pt) (:z c)))
(- (square (:radius s)))))]
(if n
(struct point (+ (:x pt) (* n (:x ray)))
(+ (:y pt) (* n (:y ray)))
(+ (:z pt) (* n (:z ray)))))))

sphere-intersect can return nil if it doesn't hit. Now we define the Lambert function

(defn lambert [s intersection ray]
(let [normal (sphere-normal s intersection)]
(max 0 (+ (* (:x ray) (:x normal))
(* (:y ray) (:y normal))
(* (:z ray) (:z normal))))))

That's it for the machinery to actually generate the image - now we need some UI and something to actually draw it. The original code generated a PPM format image, but since Clojure has a decent UI toolkit with Swing, let's just render something in a Window instead. The UI just uses the "canvas" idiom I used for the bubble sort application.

(def canvas (proxy [JPanel] []
(paintComponent [g]
(proxy-super paintComponent g)    
(.setColor g Color/RED)
(ray-trace world 1 g (.getWidth this) (.getHeight this)))))

(defn raytraceapp []
(let [frame (JFrame. "Ray Tracing")]
(doto frame
(.add canvas)
(.setSize 300 300)
(.setResizable false)
(.setVisible true))))

All that remains is to define the ray-trace function

;; second item = what we hit
;; first item = where we hit
(defn first-hit [pt ray]
(fn [x y]
(let [hx (first x) hy (first y)]
(nil? hx) y
(nil? hy) x
:else (let [d1 (distance hx pt) d2 (distance hy pt)]
(if (< d1 d2) x y)))))
(map (fn [obj]
(let [h (sphere-intersect obj pt ray)]
(if (not (nil? h))
[h obj]))) world)))

(defn send-ray [src ray]
(let [hit (first-hit src ray)]
(if (not (nil? hit))
(let [int (first hit) s (second hit)]
(* (lambert s ray int) (:color s)))

(defn color-at [x y]
(let [ray (unit-vector (point-subtract (struct point x y 0) eye))]
(* (send-ray eye ray) 255)))

(defn ray-trace [world res g w h]
(let [buffered-image (BufferedImage. w h BufferedImage/TYPE_BYTE_GRAY)]
(doseq [x (range 1 w)]
(doseq [y (range 1 h)]
(.setRGB buffered-image x y (color-at x y))))
(.drawImage g buffered-image 0 0 Color/RED nil)))
The only major difference between this and the ACL code, is prefering to use map and reduce instead of the nested do code. This feels more functional to me and also opens up parallelism opportunities which I'll look at for the next post. So what does it look like (not very good, but a background of code looks cool!)? Ray Trace application example picture