merge documentation changes

4 files changed, 100 insertions, 101 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
 
diff --git a/pi0.py b/pi0.py
index fa624d0..91e8097 100644
--- a/pi0.py
+++ b/pi0.py
@@ -1,10 +1,7 @@
 import numpy as np
 
-#c = 2.99792458e8
-#h = 6.6260689633e-34
-#joules_per_MeV = 1.602176487e-19/1e-6
-kg_per_MeV = 1.782661758e-36/1e-6
-pi0_mass = 134.9766*kg_per_MeV
+_kg_per_MeV = 1.782661758e-36/1e-6
+_pi0_mass = 134.9766*_kg_per_MeV
 
 def rocket_to_lab(energy, momentum, v):
     """
@@ -13,23 +10,18 @@ def rocket_to_lab(energy, momentum, v):
     respect to the lab frame.
 
     Args:
-        - energy: float
-            The energy of the particle in the rocket frame in kg.
-        - momentum: array
-            The momentum of the particle in the rocket frame in kg.
-        - v: array
-            The velocity of the rocket frame as seen in the lab frame in units
-            of c.
+        - energy: float, kg
+            The energy of the particle in the rocket frame.
+        - momentum: array, kg
+            The momentum of the particle in the rocket frame.
+        - v: array, units of c
+            The velocity of the rocket frame as seen in the lab frame.
     """
-    e0 = float(energy)#/c**2
-    p0 = np.asarray(momentum, float)#/c
-    v = np.asarray(v, float)#/c
+    e0 = float(energy)
+    p0 = np.asarray(momentum, float)
+    v = np.asarray(v, float)
 
-    try:
-        assert e0**2 - p0.dot(p0) >= -1.0e-70
-    except AssertionError:
-        print e0**2 - p0.dot(p0)
-        raise
+    assert e0**2 - p0.dot(p0) >= -1.0e-70
 
     g = 1.0/np.sqrt(1.0-v.dot(v))
 
@@ -37,7 +29,6 @@ def rocket_to_lab(energy, momentum, v):
     p = p0 + ((g-1.0)*x + g*np.linalg.norm(v)*e0)*v/np.linalg.norm(v)
     e = np.sqrt(e0**2 - p0.dot(p0) + p.dot(p))
 
-    #return e*c**2, p*c
     return e, p
 
 def pi0_decay(energy, direction, theta, phi):
@@ -47,56 +38,32 @@ def pi0_decay(energy, direction, theta, phi):
     angles `theta` and `phi` in the pi0 rest frame.
 
     Args:
-        - energy: float
-            The total energy of the pi0 in MeV.
+        - energy: float, MeV
+            The total energy of the pi0.
         - direction: array
             The direction of the pi0's velocity in the lab frame.
+
+    Returns:
+        - e1: float, MeV
+            The energy of the first photon in the lab frame.
+        - v1: array,
+            The direction of the first photon in the lab frame.
+        - e2: float, MeV
+            The energy of the second photon in the lab frame.
+        - v2: array,
+            The direction of the second photon in the lab frame.
     """
     direction = np.asarray(direction)/np.linalg.norm(direction)
-    pi0_e = float(energy)*kg_per_MeV#/c**2
-    pi0_p = np.sqrt(pi0_e**2-pi0_mass**2)*direction
+    pi0_e = float(energy)*_kg_per_MeV
+    pi0_p = np.sqrt(pi0_e**2-_pi0_mass**2)*direction
     pi0_v = pi0_p/pi0_e
 
-    photon_e0 = pi0_mass/2.0
+    photon_e0 = _pi0_mass/2.0
     photon_p0 = photon_e0*np.array([np.cos(theta)*np.sin(phi), np.sin(theta)*np.sin(phi), np.cos(phi)])
 
-    #print photon_p0/np.linalg.norm(photon_p0)
-
     e1, p1 = rocket_to_lab(photon_e0, photon_p0, pi0_v)
     v1 = p1/np.linalg.norm(p1)
     e2, p2 = rocket_to_lab(photon_e0, -photon_p0, pi0_v)
     v2 = p2/np.linalg.norm(p2)
 
-    return (e1/kg_per_MeV, v1), (e2/kg_per_MeV, v2)
-
-if __name__ == '__main__':
-    import sys
-
-    npi0s = 10000
-
-    pi0_e = 300.0*kg_per_MeV
-    pi0_p = np.sqrt(pi0_e**2 - pi0_mass**2)
-
-    print 'pi0 e: %f MeV' % (pi0_e/kg_per_MeV)
-    print 'pi0 p: %f MeV' % (pi0_p/kg_per_MeV)
-
-    e, cos_theta = [], []
-    for i, (theta, phi) in enumerate(zip(np.random.rand(npi0s), np.arccos(np.random.uniform(-1.0, 1.0, size=npi0s)))):
-        print '\r%i' % (i+1),
-        sys.stdout.flush()
-
-        (e1, v1), (e2, v2) = pi0_decay(pi0_e/kg_per_MeV, [0,0,1], theta, phi)
-        e += [e1, e2]
-        cos_theta += [v1.dot(v2)]
-
-    import matplotlib.pyplot as plt
-
-    plt.figure()
-    plt.title('Energy Distribution of Photons from a %.0f MeV Pi0 Decay' % (pi0_e/kg_per_MeV))
-    plt.hist(e, 100)
-    plt.xlabel('Energy (MeV)')
-    plt.figure()
-    plt.title('Opening Angle between Photons in Lab from a %.0f MeV Pi0 Decay' % (pi0_e/kg_per_MeV))
-    plt.hist(cos_theta, 100)
-    plt.xlabel('Cos($\theta$)')
-    plt.show()
+    return (e1/_kg_per_MeV, v1), (e2/_kg_per_MeV, v2)
diff --git a/stl.py b/stl.py
index 9a63306..9bc14de 100644
--- a/stl.py
+++ b/stl.py
@@ -5,6 +5,7 @@ from geometry import Mesh
 import bz2
 
 def mesh_from_stl(filename):
+    "Return a mesh from an stl file."
     if filename.endswith('.bz2'):
         f = bz2.BZ2File(filename)
     else:
@@ -19,6 +20,7 @@ def mesh_from_stl(filename):
     return mesh_from_ascii_stl(filename)
 
 def mesh_from_ascii_stl(filename):
+    "Return a mesh from an ascii stl file."
     if filename.endswith('.bz2'):
         f = bz2.BZ2File(filename)
     else:
@@ -57,6 +59,7 @@ def mesh_from_ascii_stl(filename):
     return Mesh(np.array(vertices), np.array(triangles, dtype=np.uint32))
 
 def mesh_from_binary_stl(filename):
+    "Return a mesh from a binary stl file."
     f = open(filename)
 
     vertices = []
diff --git a/tools.py b/tools.py
index ba36f6d..1338526 100644
--- a/tools.py
+++ b/tools.py
@@ -17,16 +17,18 @@ def debugger_hook(type, value, tb):
        pdb.pm()
 
 def enable_debug_on_crash():
-    '''Start the PDB console when an uncaught exception propagates to the top.'''
+    "Start the PDB console when an uncaught exception propagates to the top."
     sys.excepthook = debugger_hook
 
-# Allow profile decorator to exist, but do nothing if not running under kernprof
+# allow profile decorator to exist, but do nothing if not running under
+# kernprof
 try:
     profile_if_possible = profile
 except NameError:
     profile_if_possible = lambda x: x
 
 def timeit(func):
+    "A decorator to print the time elapsed in a function call."
     def f(*args, **kwargs):
         t0 = time.time()
         retval = func(*args, **kwargs)