update geometry.py() documentation

1 files changed, 65 insertions, 38 deletions

diff --git a/geometry.py b/geometry.py
index b08de54..856ded3 100644
--- a/geometry.py
+++ b/geometry.py
@@ -2,7 +2,7 @@ import sys
 import os
 import numpy as np
 from itertoolset import *
-from tools import timeit
+from tools import timeit, profile_if_possible
 from hashlib import md5
 import cPickle as pickle
 import gzip
@@ -36,14 +36,33 @@ class Mesh(object):
             self.remove_duplicate_vertices()
 
     def get_bounds(self):
+        "Return the lower and upper bounds for the mesh as a tuple."
         return np.min(self.vertices, axis=0), np.max(self.vertices, axis=0)
 
     def remove_duplicate_vertices(self):
+        "Remove any duplicate vertices in the mesh."
+        # view the vertices as a structured array in order to identify unique
+        # rows, i.e. unique vertices
         unique_vertices, inverse = np.unique(self.vertices.view([('', self.vertices.dtype)]*self.vertices.shape[1]), return_inverse=True)
+        # turn the structured vertex array back into a normal array
         self.vertices = unique_vertices.view(self.vertices.dtype).reshape((unique_vertices.shape[0], self.vertices.shape[1]))
+        # remap the triangles to the unique vertices
         self.triangles = np.vectorize(lambda i: inverse[i])(self.triangles)
 
     def assemble(self, key=slice(None), group=True):
+        """
+        Return an assembled triangle mesh; i.e. return the vertex positions
+        of every triangle. If `group` is True, the array returned will have
+        an extra axis denoting the triangle number; i.e. if the mesh contains
+        N triangles, the returned array will have the shape (N,3,3). If `group`
+        is False, return just the vertex positions without any grouping; in
+        this case the grouping is implied (i.e. elements [0:3] are the first
+        triangle, [3:6] the second, and so on.
+
+        The `key` argument is a slice object if you just want to assemble
+        a certain range of the triangles, i.e. to get the assembled mesh for
+        triangles 3 through 6, key = slice(3,7).
+        """
         if group:
             vertex_indices = self.triangles[key]
         else:
@@ -51,9 +70,6 @@ class Mesh(object):
 
         return self.vertices[vertex_indices]
 
-    def __len__(self):
-        return len(self.triangles)
-
     def __add__(self, other):
         return Mesh(np.concatenate((self.vertices, other.vertices)), np.concatenate((self.triangles, other.triangles + len(self.vertices))))
 
@@ -62,41 +78,38 @@ class Solid(object):
         self.mesh = mesh
 
         if np.iterable(material1):
-            if len(material1) != len(mesh):
+            if len(material1) != len(mesh.triangles):
                 raise ValueError('shape mismatch')
             self.material1 = np.array(material1, dtype=np.object)
         else:
-            self.material1 = np.tile(material1, len(self.mesh))
+            self.material1 = np.tile(material1, len(self.mesh.triangles))
 
         if np.iterable(material2):
-            if len(material2) != len(mesh):
+            if len(material2) != len(mesh.triangles):
                 raise ValueError('shape mismatch')
             self.material2 = np.array(material2, dtype=np.object)
         else:
-            self.material2 = np.tile(material2, len(self.mesh))
+            self.material2 = np.tile(material2, len(self.mesh.triangles))
 
         if np.iterable(surface):
-            if len(surface) != len(mesh):
+            if len(surface) != len(mesh.triangles):
                 raise ValueError('shape mismatch')
             self.surface = np.array(surface, dtype=np.object)
         else:
-            self.surface = np.tile(surface, len(self.mesh))
+            self.surface = np.tile(surface, len(self.mesh.triangles))
 
         if np.iterable(color):
-            if len(color) != len(mesh):
+            if len(color) != len(mesh.triangles):
                 raise ValueError('shape mismatch')
             self.color = np.array(color, dtype=np.uint32)
         else:
-            self.color = np.tile(color, len(self.mesh)).astype(np.uint32)
+            self.color = np.tile(color, len(self.mesh.triangles)).astype(np.uint32)
 
         self.unique_materials = \
             np.unique(np.concatenate([self.material1, self.material2]))
 
         self.unique_surfaces = np.unique(self.surface)
 
-    def __len__(self):
-        return len(self.mesh)
-
     def __add__(self, other):
         return Solid(self.mesh + other.mesh, np.concatenate((self.material1, other.material1)), np.concatenate((self.material2, other.material2)), np.concatenate((self.surface, other.surface)), np.concatenate((self.color, other.color)))
 
@@ -187,17 +200,23 @@ class Geometry(object):
         self.solid_displacements = []
 
     def add_solid(self, solid, rotation=np.identity(3), displacement=(0,0,0)):
+        """
+        Add the solid `solid` to the geometry. When building the final triangle
+        mesh, `solid` will be placed by rotating it with the rotation matrix
+        `rotation` and displacing it by the vector `displacement`.
+        """
+
         rotation = np.asarray(rotation, dtype=np.float32)
 
         if rotation.shape != (3,3):
-            raise ValueError('shape mismatch')
+            raise ValueError('rotation matrix has the wrong shape.')
 
         self.solid_rotations.append(rotation.astype(np.float32))
 
         displacement = np.asarray(displacement, dtype=np.float32)
 
         if displacement.shape != (3,):
-            raise ValueError('shape mismatch')
+            raise ValueError('displacement vector has the wrong shape.')
 
         self.solid_displacements.append(displacement)
 
@@ -207,21 +226,32 @@ class Geometry(object):
 
     @timeit
     def build(self, bits=8, shift=3):
-        offsets = [ (0,0) ]
-        for solid in self.solids:
-            offsets.append( (offsets[-1][0] + len(solid.mesh.vertices),
-                             offsets[-1][1] + len(solid.mesh.triangles)) )
-
-        vertices = np.zeros(shape=(offsets[-1][0], 3), dtype=np.float32)
-        triangles = np.zeros(shape=(offsets[-1][1],3), dtype=np.int32)
-
-        for i, (solid, (vertex_offset, triangle_offset)) in \
-                enumerate(zip(self.solids, offsets[:-1])):
-            triangles[triangle_offset:triangle_offset+len(solid.mesh.triangles),:] = \
-                solid.mesh.triangles + vertex_offset
-            vertices[vertex_offset:vertex_offset+len(solid.mesh.vertices),:] = \
+        """
+        Build the bounding volume hierarchy, material/surface code arrays, and
+        color array for this geometry. If the bounding volume hierarchy is
+        cached, load the cache instead of rebuilding, else build and cache it.
+
+        Args:
+            - bits: int, *optional*
+                The number of bits to quantize each linear dimension with when
+                morton ordering the triangle centers for building the bounding
+                volume hierarchy. Defaults to 8.
+            - shift: int, *optional*
+                The number of bits to shift the zvalue of each node when
+                building the next layer of the bounding volume hierarchy.
+                Defaults to 3.
+        """
+        nv = np.cumsum([0] + [len(solid.mesh.vertices) for solid in self.solids])
+        nt = np.cumsum([0] + [len(solid.mesh.triangles) for solid in self.solids])
+
+        vertices = np.empty((nv[-1],3), dtype=np.float32)
+        triangles = np.empty((nt[-1],3), dtype=np.uint32)
+
+        for i, solid in enumerate(self.solids):
+            vertices[nv[i]:nv[i+1]] = \
                 np.inner(solid.mesh.vertices, self.solid_rotations[i]) + self.solid_displacements[i]
-            
+            triangles[nt[i]:nt[i+1]] = solid.mesh.triangles + nv[i]
+
         self.mesh = Mesh(vertices, triangles)
 
         self.colors = np.concatenate([solid.color for solid in self.solids])
@@ -232,21 +262,18 @@ class Geometry(object):
 
         material_lookup = dict(zip(self.unique_materials, range(len(self.unique_materials))))
 
-        self.material1_index = \
-            np.fromiter(imap(material_lookup.get, chain(*[solid.material1 for solid in self.solids])), dtype=np.int32)
+        self.material1_index = np.fromiter(imap(material_lookup.get, chain(*[solid.material1 for solid in self.solids])), dtype=np.int32)
 
-        self.material2_index = \
-            np.fromiter(imap(material_lookup.get, chain(*[solid.material2 for solid in self.solids])), dtype=np.int32)
+        self.material2_index = np.fromiter(imap(material_lookup.get, chain(*[solid.material2 for solid in self.solids])), dtype=np.int32)
 
         self.unique_surfaces = list(np.unique(np.concatenate([solid.unique_surfaces for solid in self.solids])))
 
         surface_lookup = dict(zip(self.unique_surfaces, range(len(self.unique_surfaces))))
 
-        self.surface_index = \
-            np.fromiter(imap(surface_lookup.get, chain(*[solid.surface for solid in self.solids])), dtype=np.int32)
+        self.surface_index = np.fromiter(imap(surface_lookup.get, chain(*[solid.surface for solid in self.solids])), dtype=np.int32)
 
         try:
-            self.surface_index[self.surface_index == self.unique_surfaces.index(None)] = -1
+            self.surface_index[self.surface_index == surface_lookup[None]] = -1
         except ValueError:
             pass