BVH and BVHLayer classes with unittests.

These classes will be the data types used by the new BVH generation code.

2 files changed, 389 insertions, 0 deletions

diff --git a/chroma/bvh.py b/chroma/bvh.py
new file mode 100644
index 0000000..d13d285
--- /dev/null
+++ b/chroma/bvh.py
@@ -0,0 +1,230 @@
+'''chroma.bvh: Bounding Volume Hierarchy generation and manipulation.'''
+
+import numpy as np
+from pycuda.gpuarray import vec
+
+uint4 = vec.uint4 # pylint: disable-msg=C0103,E1101
+
+def unpack_nodes(nodes):
+    '''Creates a numpy record array with the contents of nodes
+    array unpacked into separate fields.
+
+      ``nodes``: ndarray(shape=n, dtype=uint4)
+         BVH node array in the packed x,y,z,w format.
+
+    Returns ndarray(shape=n, dtype=[('xlo', np.uint16), ('xhi', np.uint16),
+                                    ('ylo', np.uint16), ('yhi', np.uint16),
+                                    ('zlo', np.uint16), ('zhi', np.uint16),
+                                    ('child', np.uint32), ('leaf', np.bool)])
+    '''
+    unpacked_dtype = np.dtype([('xlo', np.uint16), ('xhi', np.uint16),
+                               ('ylo', np.uint16), ('yhi', np.uint16),
+                               ('zlo', np.uint16), ('zhi', np.uint16),
+                               ('child', np.uint32), ('leaf', np.bool)])
+    unpacked = np.empty(shape=len(nodes), dtype=unpacked_dtype)
+    
+    for axis in ['x', 'y', 'z']:
+        unpacked[axis+'lo'] = nodes[axis] & 0xFFFF
+        unpacked[axis+'hi'] = nodes[axis] >> 16
+    unpacked['child'] = nodes['w'] & 0x7FFFFFFF
+    unpacked['leaf']  = nodes['w'] & 0x80000000 > 0
+
+    return unpacked
+
+class OutOfRangeError(Exception):
+    '''The world coordinates cannot be transformed into fixed point
+    coordinates because they exceed the range of an unsigned 16-bit
+    fixed point number.'''
+    def __init__(self, msg):
+        Exception.__init__(self, msg)
+
+class WorldCoords(object):
+    '''A simple helper object that represents the transformation
+    between floating point world coordinates and unsigned 16-bit fixed
+    point coordinates.
+
+    This class has a ``world_origin`` (3-vector), and a ``world_scale``
+    (a scalar) property such that:
+    
+       \vec{r} = ``world_scale`` * (\vec{f} + ``world_origin``)
+
+    Where \vec{r} is a vector in the real coordinate system and \vec{f}
+    is a vector in the fixed point coordinate system.
+    '''
+    MAX_INT = 2**16 - 1
+
+    def __init__(self, world_origin, world_scale):
+        '''Transformation from fixed point to world coordinates.
+
+          ``world_origin``: ndarray(shape=3)
+              World x,y,z coordinates of (0,0,0) in fixed point coordinates.
+
+          ``world_scale``: float
+              Multiplicative scale factor to go from fixed point distances
+              to real world distances.
+        '''
+        self.world_origin = np.array(world_origin, dtype=np.float32)
+        self.world_scale = np.float32(world_scale)
+
+    def world_to_fixed(self, world):
+        '''Convert a vector, or array of vectors to the fixed point
+        representation.
+        
+          ``world``: ndarray(shape=3) or ndarray(shape=(n,3))
+             Vector or array of vectors in real world coordinates to convert to
+             fixed point.
+
+          .. warning: Conversion to fixed point rounds to nearest integer.
+
+          Returns ndarray(shape=3, dtype=np.uint16) or 
+                  ndarray(shape=(n,3), dtype=np.uint16)
+        '''
+        fixed = ((np.asfarray(world) - self.world_origin)
+                 / self.world_scale).round()
+        if int(fixed.max()) > WorldCoords.MAX_INT or fixed.min() < 0:
+            raise OutOfRangeError('range = (%f, %f)' 
+                                  % (fixed.min(), fixed.max()))
+        else:
+            return fixed.astype(np.uint16)
+
+    def fixed_to_world(self, fixed):
+        '''Convert a vector, or array of vectors to the world representation.
+        
+          ``fixed``: ndarray(shape=3, dtype=np.uint16) or 
+                   ndarray(shape=(n,3), dtype=np.uint16)
+             Vector or array of vectors in unsigned fixed point
+             coordinates to convert to world coordinates.
+
+          Returns ndarray(shape=3) or ndarray(shape=(n,3))    
+        '''
+        return np.asarray(fixed) * self.world_scale + self.world_origin
+
+
+class BVH(object):
+    '''A bounding volume hierarchy for a triangle mesh.
+
+    For the purposes of Chroma, a BVH is a tree with the following properties:
+
+      * Each node consists of an axis-aligned bounding box, a child ID
+        number, and a boolean flag indicating whether the node is a
+        leaf.  The bounding box is represented as a lower and upper
+        bound for each Cartesian axis.
+
+      * All nodes are stored in a 1D array with the root node first.
+
+      * A node with a bounding box that has no surface area (upper and
+        lower bounds equal for all axes) is a dummy node that should
+        be ignored.  Dummy nodes are used to pad the tree to satisfy
+        the fixed degree requirement described below, and have no
+        children.
+
+      * If the node is a leaf, then the child ID number refers to the
+        ID number of the triangle this node contains.
+
+      * If the node is not a leaf (an "inner" node), then the child ID
+        number indicates the offset in the node array of the first
+        child.  The other children of this node will be stored
+        immediately after the first child.
+
+      * All inner nodes have the same number of children, called the
+        "degree" (technically the "out-degree") of the tree.  This
+        avoid the requirement to save the degree with the node.
+
+      * For simplicity, we also require nodes at the same depth
+        in the tree to be contiguous, and the layers to be in order
+        of increasing depth.
+
+      * All nodes satisfy the bounding volume hierarchy constraint:
+        their bounding boxes contain the bounding boxes of all their
+        children.
+
+    For space reasons, the BVH bounds are internally represented using
+    16-bit unsigned fixed point coordinates.  Normally, we would want
+    to hide that from you, but we would like to avoid rounding issues
+    and high memory usage caused by converting back and forth between
+    floating point and fixed point representations.  For similar
+    reasons, the node array is stored in a packed record format that
+    can be directly mapped to the GPU.  In general, you will not need
+    to manipulate the contents of the BVH node array directly.
+    '''
+
+    def __init__(self, degree, world_coords, nodes, layer_offsets):
+        '''Create a BVH object with the given properties.
+
+          ``degree``: int
+              Number of child nodes per parent
+              
+           ``world_coords``: chroma.bvh.WorldCoords
+              Transformation from fixed point to world coordinates.
+              
+           ``nodes``: ndarray(shape=n, dtype=chroma.bvh.uint4) 
+              List of nodes.  x,y,z attributes of array are the lower
+              and upper limits of the bounding box (lower is the least
+              significant 16 bits), and the w attribute is the child
+              ID with the leaf state set in bit 31.  First node is root
+              node.
+
+           ``layer_offsets``: list of ints
+              Offset in node array for the start of each layer.  The first
+              entry must be 0, since the root node is first, and the 
+              second entry must be 1, for the first child of the root node,
+              unless the root node is also a leaf node.
+        '''
+        self.degree = degree
+        self.world_coords = world_coords
+        self.nodes = nodes
+        self.layer_offsets = layer_offsets
+
+        # for convenience when slicing in get_layer
+        self.layer_bounds = list(layer_offsets) + [len(nodes)]
+
+    def get_layer(self, layer_number):
+        '''Returns a BVHLayerSlice object corresponding to the given layer
+        number in this BVH, with the root node at layer 0.
+        '''
+        layer_slice = slice(self.layer_bounds[layer_number],
+                            self.layer_bounds[layer_number+1])
+        return BVHLayerSlice(degree=self.degree,
+                             world_coords=self.world_coords,
+                             nodes=self.nodes[layer_slice])
+
+    def layer_count(self):
+        '''Returns the number of layers in this BVH'''
+        return len(self.layer_offsets)
+
+    def __len__(self):
+        '''Returns the number of nodes in this BVH'''
+        return len(self.nodes)
+
+class BVHLayerSlice(object):
+    '''A single layer in a bounding volume hierarchy represented as a slice
+    in the node array of the parent.  
+
+    .. warning: Modifications to the node values in the BVHLayerSlice
+    object affect the parent storage!
+
+    A BVHLayerSlice has the same properties as the ``BVH`` class,
+    except no ``layer_offsets`` list.
+    '''
+
+    def __init__(self, degree, world_coords, nodes):
+        self.degree = degree
+        self.world_coords = world_coords
+        self.nodes = nodes
+
+    def __len__(self):
+        '''Return number of nodes in this layer'''
+        return len(self.nodes)
+
+    def area(self):
+        '''Return the surface area of all the nodes in this layer in world
+        units.'''
+        unpacked = unpack_nodes(self.nodes)
+        delta = np.empty(shape=len(self.nodes), 
+                         dtype=[('x', float), ('y', float), ('z', float)])
+        for axis in ['x', 'y', 'z']:
+            delta[axis] = unpacked[axis+'hi'] - unpacked[axis+'lo']
+
+        half_area = delta['x']*delta['y'] + delta['y']*delta['z'] \
+            + delta['z']*delta['x']
+        return 2.0 * half_area.sum() * self.world_coords.world_scale**2
diff --git a/test/test_bvh.py b/test/test_bvh.py
new file mode 100644
index 0000000..5372f6e
--- /dev/null
+++ b/test/test_bvh.py
@@ -0,0 +1,159 @@
+import unittest
+from chroma.bvh import BVH, BVHLayerSlice, WorldCoords, uint4, \
+    OutOfRangeError, unpack_nodes
+import numpy as np
+from numpy.testing import assert_array_max_ulp, assert_array_equal, \
+    assert_approx_equal
+
+def lslice(layer_bounds, layer_number):
+    '''Returns a slice object for retrieving a particular layer in an
+    array of nodes.'''
+    return slice(layer_bounds[layer_number], layer_bounds[layer_number+1])
+
+class TestWorldCoords(unittest.TestCase):
+    def setUp(self):
+        self.coords = WorldCoords([-1,-1,-1], 0.1)
+
+    def test_fixed_to_world(self):
+        f = [0, 1, 100]
+        assert_array_max_ulp(self.coords.fixed_to_world(f), 
+                             [-1.0, -0.9, 9.0], dtype=np.float32)
+    
+    def test_world_to_fixed(self):
+        w = [-1.0, -0.9, 9.0]
+        assert_array_equal(self.coords.world_to_fixed(w), 
+                           [0, 1, 100])
+
+
+    def test_fixed_array_to_world(self):
+        f = [[0, 1, 100],
+             [20, 40, 60],
+             [210, 310, 410]]
+
+        assert_array_max_ulp(self.coords.fixed_to_world(f), 
+                             [[-1.0, -0.9, 9.0], 
+                              [1.0, 3.0, 5.0],
+                              [20.0, 30.0, 40.0]],
+                             dtype=np.float32)
+
+    def test_world_array_to_fixed(self):
+        w = [[-1.0, -0.9, 9.0], 
+             [1.0, 3.0, 5.0],
+             [20.0, 30.0, 40.0]]
+        assert_array_equal(self.coords.world_to_fixed(w), 
+                           [[0, 1, 100],
+                            [20, 40, 60],
+                            [210, 310, 410]])
+
+    def test_out_of_range(self):
+        with self.assertRaises(OutOfRangeError):
+            self.coords.world_to_fixed([-2.0, 0.0, 0.0])
+        with self.assertRaises(OutOfRangeError):
+            self.coords.world_to_fixed([0.0, 1e9, 0.0])
+
+def create_bvh():
+    # 3 layer binary tree
+    degree = 2
+    world_coords = WorldCoords(np.array([-1.0,-1.0,-1.0]), 0.1)
+    layer_bounds = [0, 1, 3, 7]
+    nodes = np.empty(shape=layer_bounds[-1], dtype=uint4)
+
+    # bottom layer
+    layer = lslice(layer_bounds, 2)
+    nodes['x'][layer] = [ 0x00010000, 0x00020001,
+                          0x00010000, 0x00010000]
+    nodes['y'][layer] = [ 0x00010000, 0x00010000,
+                          0x00020001, 0x00010000]
+    nodes['z'][layer] = [ 0x00010000, 0x00010000,
+                          0x00010000, 0x00020001]
+    nodes['w'][layer] = 0x80000000 # leaf nodes
+
+    # middle layer
+    layer = lslice(layer_bounds, 1)
+    nodes['x'][layer] = [ 0x00020000, 0x00010000 ]
+    nodes['y'][layer] = [ 0x00010000, 0x00020000 ]
+    nodes['z'][layer] = [ 0x00010000, 0x00020000 ]
+    nodes['w'][layer] = [ 0x00000003, 0x00000005 ]
+
+    # top layer
+    layer = lslice(layer_bounds, 0)
+    nodes['x'][layer] = [ 0x00020000 ]
+    nodes['y'][layer] = [ 0x00020000 ]
+    nodes['z'][layer] = [ 0x00020000 ]
+    nodes['w'][layer] = [ 0x00000001 ]
+    
+    layer_offsets = list(layer_bounds[:-1]) # trim last entry
+    bvh = BVH(degree=degree, world_coords=world_coords,
+              nodes=nodes, layer_offsets=layer_offsets)
+
+    return bvh
+
+def test_unpack_nodes():
+    bvh = create_bvh()
+
+    layer = bvh.get_layer(2)
+    unpack = unpack_nodes(layer.nodes)
+    assert_array_equal(unpack['xlo'], [0, 1, 0, 0])
+    assert_array_equal(unpack['xhi'], [1, 2, 1, 1])
+    assert_array_equal(unpack['ylo'], [0, 0, 1, 0])
+    assert_array_equal(unpack['yhi'], [1, 1, 2, 1])
+    assert_array_equal(unpack['zlo'], [0, 0, 0, 1])
+    assert_array_equal(unpack['zhi'], [1, 1, 1, 2])
+    assert unpack['leaf'].all()
+
+    layer = bvh.get_layer(1)
+    unpack = unpack_nodes(layer.nodes)
+    assert_array_equal(unpack['xlo'], [0, 0])
+    assert_array_equal(unpack['xhi'], [2, 1])
+    assert_array_equal(unpack['ylo'], [0, 0])
+    assert_array_equal(unpack['yhi'], [1, 2])
+    assert_array_equal(unpack['zlo'], [0, 0])
+    assert_array_equal(unpack['zhi'], [1, 2])
+    assert_array_equal(unpack['child'], [3, 5])
+    assert not unpack['leaf'].any()
+
+    layer = bvh.get_layer(0)
+    unpack = unpack_nodes(layer.nodes)
+    assert_array_equal(unpack['xlo'], [0])
+    assert_array_equal(unpack['xhi'], [2])
+    assert_array_equal(unpack['ylo'], [0])
+    assert_array_equal(unpack['yhi'], [2])
+    assert_array_equal(unpack['zlo'], [0])
+    assert_array_equal(unpack['zhi'], [2])
+    assert_array_equal(unpack['child'], [1])
+    assert not unpack['leaf'].any()
+
+
+class TestBVH(unittest.TestCase):
+    def setUp(self):
+        self.bvh = create_bvh()
+
+    def test_len(self):
+        self.assertEqual(len(self.bvh), 7)
+
+    def test_layer_count(self):
+        self.assertEqual(self.bvh.layer_count(), 3)
+
+    def test_get_layer(self):
+        layer = self.bvh.get_layer(0)
+        self.assertEqual(len(layer), 1)
+
+        layer = self.bvh.get_layer(1)
+        self.assertEqual(len(layer), 2)
+
+        layer = self.bvh.get_layer(2)
+        self.assertEqual(len(layer), 4)
+
+class TestBVHLayer(unittest.TestCase):
+    def setUp(self):
+        self.bvh = create_bvh()
+
+    def test_area(self):
+        layer = self.bvh.get_layer(2)
+        assert_approx_equal(layer.area(), 4*6*0.1**2)
+
+        layer = self.bvh.get_layer(1)
+        assert_approx_equal(layer.area(), (4*2+2 + 4*2+2*4)*0.1**2)
+
+        layer = self.bvh.get_layer(0)
+        assert_approx_equal(layer.area(), (6*2*2)*0.1**2)