added documentation

1 files changed, 131 insertions, 31 deletions

diff --git a/geometry.py b/geometry.py
index cf92d53..2c136a3 100644
--- a/geometry.py
+++ b/geometry.py
@@ -2,9 +2,33 @@ import numpy as np
 from itertools import chain
 
 def compress(data, selectors):
+    """
+    Make an iterator that filters elements from `data` returning only those
+    that have a corresponding element in `selectors` that evaluates to True.
+    Stops when either the `data` or `selectors` iterables has been
+    exhausted.
+
+    From Python 2.7 itertools.
+    """
     return (d for d, s in zip(data, selectors) if s)
 
 class Leaf(object):
+    """
+    Leaf object represents the last layer in the bounding volume hierarchy
+    which points to a subset of the triangle mesh.
+
+    Args:
+        - mesh: array,
+            A subset of the triangle mesh.
+        - mesh_idx: int,
+            The index of the first triangle in the global mesh.
+        - zvlaue: int, *optional*
+            The zvalue associated with this leaf.
+
+    .. note::
+        The `mesh` passed in the constructor is not actually stored in the
+        leaf. It is simply used to construct the leaf's bounding box.
+    """
     def __init__(self, mesh, mesh_idx, zvalue=None):
         self.zvalue = zvalue
         self.mesh_idx = mesh_idx
@@ -15,6 +39,16 @@ class Leaf(object):
         self.size = mesh.shape[0]
 
 class Node(object):
+    """
+    Node object represents a node in the bounding volume hierarchy which
+    contains a list of child nodes.
+
+    Args:
+        - children: list,
+            List of child nodes/leafs.
+        - zvalue: int, *optional*
+            The zvalue associated with this node.
+    """
     def __init__(self, children, zvalue=None):
         self.zvalue = zvalue
 
@@ -29,15 +63,29 @@ class Node(object):
         self.children = children
 
 def interleave(arr):
+    """
+    Interleave the bits of quantized three-dimensional points in space.
+
+    Example
+        >>> interleave(np.identity(3, dtpe=np.int))
+        array([4, 2, 1], dtype=uint64)
+    """
     if len(arr.shape) != 2 or arr.shape[1] != 3:
         raise Exception('shape mismatch')
 
     z = np.zeros(arr.shape[0], dtype=np.uint64)
     for i in range(arr.dtype.itemsize*8):
-        z |= (arr[:,2] & 1 << i) << (2*i) | (arr[:,1] & 1 << i) << (2*i+1) | (arr[:,0] & 1 << i) << (2*i+2)
+        z |= (arr[:,2] & 1 << i) << (2*i) | \
+             (arr[:,1] & 1 << i) << (2*i+1) | \
+             (arr[:,0] & 1 << i) << (2*i+2)
     return z
 
-def morton_order(mesh, bits=3):
+def morton_order(mesh, bits):
+    """
+    Return a list of zvalues for triangles in `mesh` by interleaving the
+    bits of the quantized center coordinates of each triangle. Each coordinate
+    axis is quantized into 2**bits bins.
+    """
     lower_bound = np.array([np.min(mesh[:,:,0]), np.min(mesh[:,:,1]), np.min(mesh[:,:,2])])
     upper_bound = np.array([np.max(mesh[:,:,0]), np.max(mesh[:,:,1]), np.max(mesh[:,:,2])])
 
@@ -51,63 +99,115 @@ def morton_order(mesh, bits=3):
     return interleave(mean_positions)
 
 class Solid(object):
-    def __init__(self, mesh, inside, outside):
+    """
+    Object which stores a closed triangle mesh associated with a physically
+    distinct object.
+
+    Args:
+        - mesh, array
+            A closed triangle mesh.
+        - material1, Material
+            The material inside within the mesh.
+        - material2, Material
+            The material outside the mesh.
+        - surface1, Surface,
+            The surface on the inside of the mesh.
+        - surface2, Surface,
+            The surface on the outside of the mesh.
+
+    .. warning::
+        It is possible to define logically inconsistent geometries unless
+        you are careful. For example, solid A may define its inside material
+        as water but contain solid B which defines its outside material as air.
+        In this case, a photon traveling out of solid B will reflect/refract
+        assuming it's going into air, but upon reaching solid A's boundary will
+        calculate attenuation factors assuming it just traveled through water.
+    """
+    def __init__(self, mesh, material1, material2, \
+                     surface1=None, surface2=None):
         if len(mesh.shape) != 3 or mesh.shape[1] != 3 or mesh.shape[2] != 3:
-            raise Exception('shape mismatch')
+            raise Exception('shape mismatch; mesh must be a triangle mesh')
 
         self.mesh = mesh
-        self.inside = inside
-        self.outside = outside
+        self.material1 = material1
+        self.material2 = material2
+        self.surface1 = surface1
+        self.surface2 = surface2
 
     def __len__(self):
         return self.mesh.shape[0]
 
 class Geometry(object):
-    """
-    Geometry object.
-    """
+    """Object which stores the global mesh for a geometry."""
 
     def __init__(self):
         self.solids = []
         self.materials = []
+        self.surfaces = []
 
     def add_solid(self, solid):
+        """Add a solid to the geometry."""
         self.solids.append(solid)
 
-        if solid.inside not in self.materials:
-            self.materials.append(solid.inside)
+        if solid.material1 not in self.materials:
+            self.materials.append(solid.material1)
 
-        if solid.outside not in self.materials:
-            self.materials.append(solid.outside)
+        if solid.material2 not in self.materials:
+            self.materials.append(solid.material2)
 
-    def build(self, bits=3):
-        self.mesh = np.concatenate([solid.mesh for solid in self.solids])
-        self.solid_index = np.concatenate([np.tile(self.solids.index(solid), len(solid)) for solid in self.solids])
-        self.inside_index = np.concatenate([np.tile(self.materials.index(solid.outside), len(solid)) for solid in self.solids])
-        self.outside_index = np.concatenate([np.tile(self.materials.index(solid.outside), len(solid)) for solid in self.solids])
+        if solid.surface1 not in self.surfaces:
+            self.surfaces.append(solid.surface1)
 
-        zvalues = morton_order(self.mesh, bits)
+        if solid.surface2 not in self.surfaces:
+            self.surfaces.append(solid.surface1)
 
+    def build(self, bits=3):
+        """Build the bounding volume hierarchy of the geometry."""
+        mesh = np.concatenate([solid.mesh for solid in self.solids])
+
+        # lookup solid/material/surface index from triangle index
+        solid_index = \
+            np.concatenate([np.tile(self.solids.index(solid), \
+                                        len(solid)) for solid in self.solids])
+        material1_index = \
+            np.concatenate([np.tile(self.materials.index(solid.material1), \
+                                        len(solid)) for solid in self.solids])
+        material2_index = \
+            np.concatenate([np.tile(self.materials.index(solid.material2), \
+                                        len(solid)) for solid in self.solids])
+        surface1_index = \
+            np.concatenate([np.tile(self.surfaces.index(solid.surface1), \
+                                        len(solid)) for solid in self.solids])
+        surface2_index = \
+            np.concatenate([np.tile(self.surfaces.index(solid.surface2), \
+                                        len(solid)) for solid in self.solids])
+
+        # array of zvalue for each triangle in mesh
+        zvalues = morton_order(mesh, bits)
+
+        # mesh indices in order of increasing zvalue
         order = np.array(zip(*sorted(zip(zvalues, range(zvalues.size))))[-1])
 
-        zvalues = zvalues[order]
-        self.mesh = self.mesh[order]
-        self.solid_index = self.solid_index[order]
-        self.inside_index = self.inside_index[order]
-        self.outside_index = self.outside_index[order]
+        # reorder all arrays ordered by triangle index
+        self.zvalues = zvalues[order]
+        self.mesh = mesh[order]
+        self.solid_index = solid_index[order]
+        self.material1_index = material1_index[order]
+        self.material2_index = material2_index[order]
+        self.surface1_index = surface1_index[order]
+        self.surface2_index = surface2_index[order]
 
         leafs = []
-        for z in sorted(set(zvalues)):
-            mask = (zvalues == z)
+        for z in sorted(set(self.zvalues)):
+            mask = (self.zvalues == z)
             leafs.append(Leaf(self.mesh[mask], np.where(mask)[0][0], z))
 
         layers = []
         layers.append(leafs)
 
         while True:
-            zvalues = np.array([node.zvalue for node in layers[-1]])
-
-            bit_shifted_zvalues = zvalues >> 1
+            bit_shifted_zvalues = \
+                np.array([node.zvalue for node in layers[-1]]) >> 1
 
             nodes = []
             for z in sorted(set(bit_shifted_zvalues)):
@@ -131,7 +231,7 @@ class Geometry(object):
         self.upper_bound = np.empty((len(nodes),3))
         self.child_map = np.empty(len(nodes))
         self.child_len = np.empty(len(nodes))
-        self.first_leaf = -1
+        self.first_leaf = None
 
         for i, node in enumerate(nodes):
             self.lower_bound[i] = node.lower_bound
@@ -148,5 +248,5 @@ class Geometry(object):
             if isinstance(node, Leaf):
                 self.child_map[i] = node.mesh_idx
 
-                if self.first_leaf == -1:
+                if self.first_leaf is None:
                     self.first_leaf = i