Tuesday, 22 March 2011

Sphere Tree Construction and Traversal in Haskell

Finally. I've finally finished writing and debugging my sphere tree construction code in Haskell.

It should have been so easy. Tree structures are very easily recursively defined which is ideal for functional languages like Haskell. Take your data, split it up, and build the children by recursion. Simple. Only it was anything but...

The logic of the task was simple enough. The main problem was debugging. This is the sort of task that is prone to simple mistakes that cause time-consuming debugging problems. When working in an imperative language, you can easily put down a breakpoint and inspect your data, or if you recur too much the stack blows up, or you can step through your code and see why you're getting stuck in loops. These strategies don't map well to Haskell.

What does work in Haskell is to start small, build up gradually and constantly verify your results. You end up with a development cycle that is fundamentally very similar to test driven development. You want to prove your code one little function at a time by feeding it controlled inputs and inspecting the outputs. In Haskell, your primary debugger is your brain - you need to think very carefully through your code and how your express your algorithms. Functional languages reduce the scope of creating simple bugs due to side effects, and Haskell in particular eliminates a lot of plumbing and housekeeping code. What's left is purely your algorithm - inspect it carefully, that is where your problem lies.

Bounding Volume Hierarchies

I considered a few approaches for my bounding volume hierarchy, and for now, I've settled on a fairly generic implementation of bounding volume hierarchies that currently is configured as a sphere tree.

Each tree node can optionally (Maybe) hold an object. Each node has a sphere, but that could later be generalised. The children are defined as a list, allowing me to switch using a quadtree, kdtree or octree type approach.

(Note: The code below is best copied and viewed in a real text editor. Also, I am not yet a Haskell expert, and there may be many superior ways to achieve these results. My purpose in writing is to share my experiences as a learner with other learners.)

data SphereTreeNode = SphereTreeNode { object :: Maybe Object, children :: [SphereTreeNode], boundingRadius :: Float, boundingCentre :: Vector } deriving (Show, Read)

Sphere Tree Construction

The tree itself is constructed using a fairly simple piece of Haskell:

-- Build up a sphere tree
buildSphereTree :: ([Object] -> [[Object]]) -> [Object] -> SphereTreeNode
buildSphereTree _ (obj : []) = SphereTreeNode (Just obj) [] nodeRadius nodeCentre
      nodeCentre = calculateMeanPosition (obj:[])
      nodeRadius = calculateBoundingRadius (obj:[]) nodeCentre
buildSphereTree builder (obj:objs)
    | length (obj:objs) == 1 = error "Should have been handled by a different pattern"
    | length (obj:objs) == 0 = error "Should not have zero objects"
    | otherwise = SphereTreeNode Nothing nodeChildren nodeRadius nodeCentre
      nodeCentre = calculateMeanPosition (obj:objs)
      nodeRadius = calculateBoundingRadius (obj:objs) nodeCentre
      nodeChildren = map (buildSphereTree builder) (builder (obj:objs))
buildSphereTree _ [] = error "Should not hit this pattern for buildSphereTree" 

-- Find the mean of a collection of objects
calculateMeanPosition' :: [Object] -> Vector -> Vector
calculateMeanPosition' (obj : objects) acc = calculateMeanPosition' objects acc + (getCentre obj)
calculateMeanPosition' [] acc = acc

calculateMeanPosition :: [Object] -> Vector
calculateMeanPosition objects = setWTo1 ((calculateMeanPosition' objects zeroVector) fromIntegral (length objects))

-- Find the overall bounding radius of a list of objects
calculateBoundingRadius :: [Object] -> Vector -> Float
calculateBoundingRadius objs centre = foldr Prelude.max 0 (map (\obj -> shapeBoundingRadius (shape obj) (transform obj) centre) objs)

This object simply recursively builds a tree until it encounters an object list of size 1. 

I've deliberately exercised a couple of Haskell idioms here. I've used another tail-recursive loop to calculate the mean centre of the objects, and a map/fold pair to calculate the sphere's bounding radius.

The most interesting part is the first parameter to buildSphereTree. I pass in a function of type ([Object] -> [[Object]]). This function is responsible for dividing the object list into a number of lists, each of which is the object list used to build a new child node. This trick allows me to abstract out the specific algorithm for building the tree into a user-supplied function.

In my current code, I'm using a simple KD-type approach:

-- Generate a plane to split the objects along
makeSplittingPlane :: AABB -> (Vector, Float)
makeSplittingPlane (boxMin, boxMax) = case largestAxis (boxMax - boxMin) of
                                        0 -> (xaxis, -(vecX midPoint))
                                        1 -> (yaxis, -(vecY midPoint))
                                        2 -> (zaxis, -(vecZ midPoint))
                                        _ -> error "Undefined value"
      midPoint = (boxMin + boxMax) <*> 0.5

-- Make children using a kd tree
generateSceneGraphUsingKDTree :: [Object] -> [[Object]]
generateSceneGraphUsingKDTree objs = leftObjects : rightObjects : []
      objBox = objectListBoundingBox objs
      (planeNormal, planeDist) = makeSplittingPlane objBox
      (leftObjects, rightObjects) = partition (\obj -> planeDist + (dot3 planeNormal $ getCentre obj) > 0) objs

We simply make a splitting plane and split the list of objects according to what side of the plane their centre lies: [[Object], [Object]].

I could easily replace this function with an alternative strategy. Or even, make a function that returns a function based on the characteristics of the presented list...

Bounding Volume Hierarchy Intersection

Now this is the real meat of the implementation. The function again takes the form of a tail-recursive loop. The code is essentially similar to implementing a recursive function using a software stack in C. The tail recursion is used to make a simple loop-type recursion rather than a true state-accumulation call-type recursion.

The code is somewhat complex as several optimisations have been added.

The code is passed a list of nodes to process. We process the first node in the list, and recur with the remainder of the list - and possibly, some newly added nodes. The code seeks to find and maintain the closest-intersecting object.

The code constructs a tuple of three values - a result for this object, a list of additional nodes to process and a ray. If an object or bounding volume is hit, this tuples holds new values. If no intersection is found, we simply continue with the existing state passed into the function.

The code first considers whether the given node has any children; if it has none, it only makes sense to test the object that the node should hold. We report an error if the node has no object (a node without children or objects is useless!).
If the node does have children, it holds a subtree we wish to accept or reject. We first test the bounding volume. If this is rejected, we simply retain the current results of the loop with no additional nodes to process. If the bounding volume is intersected, then we will at least recur to the children nodes. If the node contains an object and it is intersected, we update our current intersection results and shorten our ray.

Finally, we pattern match the empty array case to simply return the current results.

Here is the code:

intersectSphereTree :: [SphereTreeNode] -> Ray -> Maybe (Object, Float, Int) -> Maybe (Object, Float, Int)
intersectSphereTree (node:nodes) ray currentHit = intersectSphereTree (newNodeList ++ nodes) newRay thisResult
      -- Intersect the ray with the bounding volume of this node
      (thisResult, newNodeList, newRay) = case children node of
                                            -- If the node has no children, don't bother with it's bounding volume and just check the object (if it has one)
                                            [] -> case object node of
                                                             Nothing -> error "A node with no children should hold an object"
                                                             Just obj -> case shapeIntersect (shape obj) ray (transform obj) of
                                                                           -- Didn't hit the object. Retain the current hit, and continue with remaining nodes on the list
                                                                           Nothing -> (currentHit, [], ray)
                                                                           -- We did hit this object. Update the intersection, and continue with remaining nodes on the list
                                                                           Just (objHitDistance, objHitId) -> (Just (obj, objHitDistance, objHitId), [], shortenRay ray objHitDistance)
                                            -- We have children. In this case it makes sense to test our bounding volume
                                            nodeChildren -> case shapeIntersect (Sphere (boundingRadius node)) ray (translationMatrix' (boundingCentre node)) of -- (make a sphere centred at the object's transform matrix with given radius)
                                                              -- If we do not find an intersection, we do not update the results and we offer no further nodes to be traversed, thus skipping this subtree
                                                              Nothing -> (currentHit, [], ray)
                                                              -- If we do find an intersection against the bounding volume, then we try again against the actual object (if present)
                                                              Just (_, _) -> case object node of
                                                                               Nothing -> (currentHit, nodeChildren, ray) -- No object; just pass to the children
                                                                               Just obj -> case shapeIntersect (shape obj) ray (transform obj) of
                                                                                             -- Didn't hit the object. Retain the current hit, but offer up the children of the node as we hit the bounding volume
                                                                                             Nothing -> (currentHit, nodeChildren, ray)
                                                                                             -- We did hit this object. Update the intersection, and continue with the bounding volume's children
                                                                                             Just (objHitDistance, objHitId) -> (Just (obj, objHitDistance, objHitId), nodeChildren, shortenRay ray objHitDistance)
intersectSphereTree [] _ currentHit = currentHit

The striking thing about this code is that it steadily marches to the right of the screen, with each indentation being a new test of success or failure. Ideally, this code could be rewritten as a monad.

What About Infinite Objects?

They're simply "filter"ed out at the time of the sphere tree construction, and added to a separate list. It makes little sense to include an infinite object into any bounding volume hierarchy.


The code continues to render my test scene as expected. Disappointingly, it is currently a little slower; however, this is not due to the scene graph code. I have recently rewritten my vector classes to be instances of the Num typeclass, so that I can simply write a + b, rather than coining new operators such as a <+> b. I haven't yet fully optimised my new code.

Other Observations

As I get further into Haskell, I'm finding myself increasingly thinking in terms of Lambda calculus. I no longer find myself thinking of what iterative steps I need to do to complete an operation, rather I find myself thinking in terms of the transformations of the underlying data and accompanying functions.

I find this new perspective very helpful when returning to other languages. It gives you a somewhat stream-processing-like perspective and thought process when facing standard C++ code.

What's Next?

The main focus now is to diagnose, understand and resolve the issues with my new typeclass-based fundamental data types such as Vector and Colour. I want to get it below the ~70 seconds it was previously taking. Following that, it is now time to build a parser and load some standard test scenes that will allow me to develop and optimise the ray tracer further.