Ray Tracing in Haskell

I’m giving a talk at a user group meeting next week on how to program in Haskell. Rather than just listing off language constructs, I wanted to motivate the whole thing with an actual application. So that’s how I ended up spending yesterday afternoon writing a ray tracer in Haskell. I’m certainly not the first to do so, nor the best, but I am writing this blog post to explain the process I went through. I’ll try making this a literate Haskell post, so you can copy and paste it into a source file (extension .lhs) and run it. The actual code is split into three (very short) modules, but this should work.

Some imports to start us off:

> import Data.Maybe (mapMaybe)
> import Data.Function (on)
> import Data.List (minimumBy)
> import Data.Ix (range)
> import Control.Monad (forM_)
> import Graphics.UI.Gtk hiding (Color, Point, Object)
> import qualified Graphics.UI.Gtk as GTK

First, I built a couple basics of vector math. Out of a desire to demonstrate some list manipulation functions in the presentation, my Vector type is simply [Double].

> type Vector = [Double]
> 
> (<+>)       = zipWith (+)
> (<->)       = zipWith (-)
> (<*>)       = zipWith (*)
> vlength u   = sqrt (sum (map (^2) u))
> normalize u = map (/ vlength u) u
> u .*. v     = sum (u <*> v)
> k #*  v     = map (*k) v

Since vectors don’t really even approximate a ring, it seems quite inappropriate to call them Num. Therefore, I’ve invented my own operators. The operators in angle brackets are component-wise, while the normal dot and scalar products are given their own operators, .*. and #*, respectively. I found this to be quite usable, though it would be equally reasonable to eschew operators, and define these as named functions to be used in infix notation.

I then define a few precedences, to make these easier to use.

> infixl 6 <+>
> infixl 6 <->
> infixl 7 <*>
> infixl 7 .*.
> infixl 7 #*

I’ll treat points as displacement vectors from an origin, and colors as vectors in a color space, but it helps to use separate names for them when writing type signatures:

> type Point = Vector
> type Color = Vector

Now for a few real data types. The idea is that an object has some general-purpose properties, plus a shape. The shape could be arbitrary, but is currently limited to only spheres and infinite planes. (Hey, this talk is only an hour and a half long!) So I’ve separated it out. The other properties I’ve got are material properties: how shiny (that is, reflective) is the surface, and what is its color. These are defined as functions from a point in space to the desired values, but in practice I’m just going to use const.

> data Object = Obj          { shape       :: Shape,
>                              objectColor :: Point -> Color,
>                              objectShine :: Point -> Double }
> 
> data Shape  = Sphere       { center  :: Point, radius      :: Double }
>             | Plane        { point   :: Point, normal      :: Vector }

And lights are important to. Right now, I’ve got point lights and ambient lights. Adding directed lights (e.g. spot lights) wouldn’t be too difficult, but it is neither of critical importance nor necessary to demonstrate a Haskell language idea, so I’ve left it out.

> data Light  = PointLight   { lightPt :: Point, lightColor  :: Color  }
>             | AmbientLight { lightColor :: Color                     }

The two things we need to know about any shape are how to compute the distance to it along a given ray (specified by origin and a unit direction vector), and how to find the normal (another unit vector) at a given point on the surface of the object. These next two functions do this. They are defined for spheres and planes, but can of course be extended to new variants of Shape via additional pattern matching.

> shapeDistance :: Shape -> Point -> Vector -> Maybe Double
> 
> shapeDistance (Sphere c r) orig dir =
>     let p    = orig <+> (c <-> orig) .*. dir #* dir
>         d_cp = vlength (p <-> c)
>         d    = vlength (p <-> orig) - sqrt (r^2 - d_cp^2)
>         ans | d_cp >= r                 = Nothing
>             | (p <-> orig) .*. dir <= 0 = Nothing
>             | otherwise                 = Just d
>     in  ans
> 
> shapeDistance (Plane p n) orig dir
>     | dir .*. n >= 0  = Nothing
>     | otherwise       = Just (((p <-> orig) .*. n) / (dir .*. n))
> 
> shapeNormal :: Shape -> Point -> Vector
> shapeNormal (Sphere c r) pt = normalize (pt <-> c)
> shapeNormal (Plane p n)  pt  = n

The math there can be derived without a whole lot of trouble, provided you’re relatively familiar with the basic ideas of vector operations. It surprises me, though, how many people only understand dot products in terms of their implementation (the sum of the component-wise products, or a b cos(theta)), and not in terms of what they mean. In particular, when one of the vectors in a dot product is a unit vector, the dot product is the component of the other vector along the direction of the unit vector. That’s a pretty powerful tool to have.

Next come the lights. For the lights, the general question is how much (in each color component) the given light illuminates a particular point with a particular normal. The object list is needed here as well, since lights can be occluded. This is all captured in a function applyLight, which is implemented for each type of light. It’s trivial for ambient lights, since by definition they ignore all the specifics of the situation; but it’s non-trivial for point lights. I also fence all light components to eliminate “negative” lights.

> fence a b x | x < a     = a
>             | x > b     = b
>             | otherwise = x
> 
> applyLight :: [Object] -> Point -> Vector -> Light -> Vector
> applyLight objs pt n (PointLight lpt color) =
>     let dir   = normalize (lpt <-> pt)
>         dist  = vlength   (lpt <-> pt)
>         lval  = (n .*. dir) / dist^2 #* color
>         final = map (fence 0 9999) lval
>     in  case findNearest objs pt dir of
>             Just (_,opt)
>                 | vlength (opt <-> pt) < dist -> [0,0,0]
>                 | otherwise                   -> final
>             Nothing                           -> final
> 
> applyLight objs pt n (AmbientLight color) = color

Now for the common stuff. A very common thing we want is to take our list of objects and find the nearest point of intersection for a given ray. This function does that, by calling objectDistance for each object in turn, and then choosing the minimum one. It gives back Nothing if there is no intersection, or Just (obj,point) if there is.

> findNearest :: [Object] -> Point -> Vector -> Maybe (Object, Point)
> findNearest objs orig dir =
>         let consider obj = case shapeDistance (shape obj) orig dir of
>                                 Nothing -> Nothing
>                                 Just d  -> Just (obj, d)
>             result = mapMaybe consider objs
>         in  case result of
>                 [] -> Nothing
>                 _  -> let (obj,d) = minimumBy (compare `on` snd) result
>                       in  Just (obj, orig <+> d #* dir)

And finally, the ray tracing itself. I’ve also thrown in defaultColor, which is the color assigned to rays that have either reflected back and forth beyond the recursion limit, or that proceed into infinity. For now, this is always black.

> defaultColor :: Vector -> Color
> defaultColor _ = [0, 0, 0]
> 
> trace :: [Object] -> [Light] -> Point -> Vector -> Int -> Color
> trace objs lights orig dir 0     = defaultColor dir
> trace objs lights orig dir limit =
>     case findNearest objs orig dir of
>         Nothing        -> defaultColor dir
>         Just (obj, pt) ->
>             let n        = shapeNormal (shape obj) pt
>                 ndir     = dir <-> 2 * (n .*. dir) #* n
>                 refl     = trace objs lights pt ndir (limit - 1)
>                 color    = objectColor obj pt
>                 shine    = objectShine obj pt
>                 lvals    = map (applyLight objs pt n) lights
>                 lighting = (foldl (<+>) [0,0,0] lvals)
>                 in_light = shine #* refl <+> (1 - shine) #* lighting
>             in map (fence 0 1) (color <*> in_light)

Finally, there’s the GUI stuff. I won’t discuss this much, except to say that it’s horridly bad. I’m doing the actual ray tracing in an onExpose event handler. If I had even a modicum of common sense, I’d be rendering to an image or something, so that I wouldn’t have to redo all the ray tracing just because you switched over to a different window. Oh well. I’ll improve this before the talk on Tuesday.

> objs = [
>     Obj (Sphere [-50, -50, 100] 100)
>         (const [1.0, 0.3, 1.0])
>         (const 0.1),
>     Obj (Sphere [ 50,  50, 100] 100)
>         (const [0.2, 1.0, 0.2])
>         (const 0.1),
>     Obj (Plane [-100,100,0] (normalize [1,-0.2,-0.2]))
>         (const [1.0,1.0,1.0])
>         (const 0.0)
>     ]
> 
> lights = [ PointLight   [175, -200, 10] [40000.0, 20000, 20000],
>            AmbientLight                [0.0, 0.05, 0.1]
>          ]
> 
> drawScene d ev = do
>     dw    <- widgetGetDrawWindow d
>     (w,h) <- widgetGetSize d
>     gc    <- gcNew dw
> 
>     let eye = [0, 0, -500]
> 
>     forM_ (range ((0,0), (w-1, h-1))) $ \(x,y) -> do
>         let (x', y') = (fromIntegral x - fromIntegral w / 2,
>                         fromIntegral y - fromIntegral h / 2)
>         let [r,g,b]  = trace objs lights eye
>                              (normalize ([x',y',0] <-> eye)) 5
>         let fg       = GTK.Color (round (65535 * r))
>                                  (round (65535 * g))
>                                  (round (65535 * b))
>         gcSetValues gc $ newGCValues { foreground = fg }
>         drawPoint dw gc (x,y)
>     return True
> 
> main = do
>     initGUI
>     w <- windowNew
>     d <- drawingAreaNew
>     windowSetTitle w "Ray Tracer"
>     containerAdd   w d
>     onExpose  d (drawScene d)
>     onDestroy w mainQuit
>     widgetShowAll w
>     mainGUI

And there you go: a ray tracer! Of course, there are many things that ought to be improved. They range from the trivial (render to an off-screen image, perhaps save to an image file, and read from a file instead of coding the scene as an immediate data structure) to the very complex (implement diffuse effects, more complex interpolated curves, etc.). But as a quick demo goes, this one didn’t turn out so bad.