added documentation

12 files changed, 390 insertions, 158 deletions

diff --git a/__init__.py b/__init__.py
index 2d61de9..409c3ea 100644
--- a/__init__.py
+++ b/__init__.py
@@ -2,3 +2,6 @@ import geometry
 import materials
 import transform
 import stl
+import view
+import photon
+import gpu
diff --git a/camera.py b/camera.py
index f00cc56..5a05535 100644
--- a/camera.py
+++ b/camera.py
@@ -1,22 +1,33 @@
 import numpy as np
+from itertools import product
 
 class Camera(object):
-    def __init__(self, size = (800, 600), film_size = (0.035, 0.024), focal_length=0.05):
-        width, height = size
+    """
+    Pinhole camera object.
 
-        grid = []
-        for i, x in enumerate(np.linspace(-film_size[0]/2, film_size[0]/2, width)):
-            for j, z in enumerate(np.linspace(-film_size[1]/2, film_size[1]/2, height)):
-                grid.append((x,0,z))
+    Args:
+        - size: tuple, *optional*
+            Pixel array shape.
+        - film_size: tuple, *optional*
+            Physical size of photographic film. Defaults to 35mm film size.
+        - focal_length: float, *optional*
+            Focal length of camera.
+    """
+    def __init__(self, size = (800, 600), film_size = (0.035, 0.024), \
+                     focal_length=0.05):
+        x = np.linspace(-film_size[0]/2, film_size[0]/2, size[0])
+        z = np.linspace(-film_size[1]/2, film_size[1]/2, size[1])
 
-        self.grid = np.array(grid)
-        self.grid += (0,focal_length,0)
+        self.grid = np.array(tuple(product(x,[0],z)))
 
+        self.grid += (0,focal_length,0)
         self.focal_point = np.zeros(3)
 
     def position(self, position):
+        """Translate the camera to `position`."""
         self.grid += position
         self.focal_point += position
 
     def get_rays(self):
+        """Return the position and direction for each pixel ray."""
         return self.grid, self.focal_point-self.grid
diff --git a/detectors/lbne.py b/detectors/lbne.py
index 4473084..cac737a 100644
--- a/detectors/lbne.py
+++ b/detectors/lbne.py
@@ -2,43 +2,129 @@ import os
 import numpy as np
 import pickle
 from chroma import *
+from copy import deepcopy
+from histogram import *
 
 models_directory = os.path.split(os.path.realpath(__file__))[0] + '/../models'
 
-strings = 10
+strings = 20
 pmts_per_string = 10
 radius = 10.0
 height = 20.0
 
-def build_lbne(bits=4):
-    lbne = geometry.Geometry()
+grid_spacing = height/pmts_per_string
 
-    sphere_mesh = stl.read_stl(models_directory + '/hamamatsu_12inch.stl')
-    sphere_mesh /= 1000.0
+block_size = 64
 
-    solids = []
-    for i in range(pmts_per_string):
-        for j in range(strings):
-            sphere = np.copy(sphere_mesh)
-            sphere += (-radius,0,i*(height/pmts_per_string))
-            sphere = transform.rotate(sphere, j*2*np.pi/strings, (0,0,1))
-            lbne.add_solid(geometry.Solid(sphere, materials.vacuum, materials.h2o))
 
-    lbne.build(bits)
+class LBNE(geometry.Geometry):
+    def __init__(self):
+        super(LBNE, self).__init__()
 
-    return lbne
+        pmt_mesh = stl.read_stl(models_directory + '/hamamatsu_12inch.stl')
+        pmt_mesh /= 1000.0
 
-def cache_lbne(filename):
-    lbne = build_lbne(bits=8)
-    f = open(filename, 'wb')
-    pickle.dump(lbne, f)
-    f.close()
+        apmt = geometry.Solid(pmt_mesh, materials.glass, materials.h2o)
 
-def load_lbne(filename):
-    f =  open(filename, 'rb')
-    lbne = pickle.load(f)
-    f.close()
-    return lbne
+        self.pmt_index = []
+        self.pmt_local_axes = []
+        self.pmt_positions = []
+
+        self.pmt_hits = []
+
+        for i in range(pmts_per_string):
+            for j in range(strings):
+                pmt = deepcopy(apmt)
+                pmt.mesh += (-radius,0,i*(height/pmts_per_string))
+                pmt.mesh = transform.rotate(pmt.mesh, j*2*np.pi/strings, (0,0,1))
+                self.add_solid(pmt)
+                self.pmt_hits.append(Histogram(10000, (-0.5, 9999.5)))
+
+        for x in np.arange(-radius, radius, grid_spacing):
+            for y in np.arange(-radius, radius, grid_spacing):
+                if np.sqrt(x**2+y**2) <= radius:
+                    pmt = deepcopy(apmt)
+                    pmt.mesh = transform.rotate(pmt.mesh, np.pi/2, (0,1,0))
+                    pmt.mesh += (x,y,0)
+                    self.add_solid(pmt)
+                    self.pmt_hits.append(Histogram(10000, (-0.5, 9999.5)))
+
+        for x in np.arange(-radius, radius, grid_spacing):
+            for y in np.arange(-radius, radius, grid_spacing):
+                if np.sqrt(x**2+y**2) <= radius:
+                    pmt = deepcopy(apmt)
+                    pmt.mesh = transform.rotate(pmt.mesh, -np.pi/2, (0,1,0))
+                    pmt.mesh += (x,y,height)
+                    self.add_solid(pmt)
+                    self.pmt_hits.append(Histogram(10000, (-0.5, 9999.5)))
+
+        self.build(bits=4)
+
+        self.npmts = len(self.pmt_hits)
+
+        self.gpu = gpu.GPU()
+        self.gpu.load_geometry(self)
+
+    def throw_photon_bomb(self, z_position, nphotons=100000):
+        origin = np.zeros((nphotons,3)) + (0,0,z_position)
+
+        direction = photon.uniform_sphere(nphotons)
+
+        origin_gpu = gpu.cuda.to_device(gpu.make_vector(origin))
+        direction_gpu = gpu.cuda.to_device(gpu.make_vector(direction))
+
+        pixels = np.empty(nphotons, dtype=np.int32)
+        states = np.empty(nphotons, dtype=np.int32)
+
+        pixels_gpu = gpu.cuda.to_device(pixels)
+        states_gpu = gpu.cuda.to_device(states)
+
+        gpu_kwargs = {'block': (block_size,1,1), 'grid': (nphotons//block_size+1,1)}
+        self.gpu.call(np.int32(nphotons), origin_gpu, direction_gpu, np.int32(self.first_leaf), states_gpu, pixels_gpu, **gpu_kwargs)
+
+        gpu.cuda.memcpy_dtoh(states, states_gpu)
+
+        pmt_indices = self.solid_index[states[(states != -1)]]
+
+        bin_count = np.bincount(pmt_indices)
+
+        bin_count = np.append(bin_count, np.zeros(self.npmts-bin_count.size))
+
+        return bin_count
+
+    def generate_event(self, z_position):
+        self.bin_count = self.throw_photon_bomb(z_position)
+
+    def get_likelihood(self, z_position, calls=1000):
+        if not hasattr(self, 'bin_count'):
+            raise Exception('must call generate_event() first')
+
+        for pmt_hit in self.pmt_hits:
+            pmt_hit.reset()
+
+        for i in range(calls):
+            print 'throwing bomb %i' % i
+            bin_count = self.throw_photon_bomb(z_position)
+
+            for i, count in enumerate(bin_count):
+                self.pmt_hits[i].fill(count)
+
+        for pmt_hit in self.pmt_hits:
+            pmt_hit.normalize()
+
+        likelihood = 0.0
+        for i in range(self.npmts):
+            probability = self.pmt_hits[i].eval(self.bin_count[i])
+
+            if probability == 0.0:
+                print 'calculating likelihood from pmt %i' % i
+                print 'bin count =', self.bin_count[i]
+                print self.pmt_hits[i].hist
+
+            likelihood -= np.log(self.pmt_hits[i].eval(self.bin_count[i]))
+
+        return likelihood
 
 if __name__ == '__main__':
-    build_lbne()
+    lbne = LBNE()
+    view.view(lbne)
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
diff --git a/gpu.py b/gpu.py
index dbcb060..3d30387 100644
--- a/gpu.py
+++ b/gpu.py
@@ -22,7 +22,7 @@ print 'device %s' % autoinit.device.name()
 
 source_directory = os.path.split(os.path.realpath(__file__))[0] + '/src'
 
-source = open(source_directory + '/kernel.cu').read()
+source = open(source_directory + '/intersect.cu').read()
 module = SourceModule(source, options=['-I' + source_directory], no_extern_c=True, cache_dir=False)
 
 cuda_intersect = module.get_function('intersect_mesh')
diff --git a/materials.py b/materials.py
index 4245212..4a8058f 100644
--- a/materials.py
+++ b/materials.py
@@ -1,4 +1,5 @@
 class Material(object):
+    """Material optical properties."""
     def __init__(self, name='none'):
         self.name = name
 
@@ -7,6 +8,13 @@ class Material(object):
         self.absorption_length = None
         self.scattering_length = None
 
+class Surface(object):
+    """Surface optical properties."""
+    def __init__(self, name='none'):
+        self.name = name
+
+        self.wavelength = None
+
 air = Material('air')
 h2o = Material('h2o')
 glass = Material('glass')
diff --git a/photon.py b/photon.py
index 3472c1e..b58a97b 100644
--- a/photon.py
+++ b/photon.py
@@ -1,4 +1,3 @@
-import time
 import numpy as np
 
 def uniform_sphere(size=None):
@@ -13,8 +12,6 @@ def uniform_sphere(size=None):
     source: Weisstein, Eric W. "Sphere Point Picking." Mathworld. 
     """
 
-    t0 = time.time()
-
     theta, u = np.random.uniform(0.0, 2*np.pi, size), \
         np.random.uniform(-1.0, 1.0, size)
 
diff --git a/src/kernel.cu b/src/intersect.cu
index c2b3fb2..1402aa1 100644
--- a/src/kernel.cu
+++ b/src/intersect.cu
@@ -5,11 +5,14 @@
 #include "matrix.h"
 #include "rotate.h"
 
+/* flattened triangle mesh */
 texture<float4, 1, cudaReadModeElementType> mesh;
 
+/* lower/upper bounds for the bounding box associated with each node/leaf */
 texture<float4, 1, cudaReadModeElementType> upper_bound_arr;
 texture<float4, 1, cudaReadModeElementType> lower_bound_arr;
 
+/* map to child nodes/triangles and the number of child nodes/triangles */
 texture<uint, 1, cudaReadModeElementType> child_map_arr;
 texture<uint, 1, cudaReadModeElementType> child_len_arr;
 
@@ -18,6 +21,12 @@ __device__ float3 make_float3(const float4 &a)
 	return make_float3(a.x, a.y, a.z);
 }
 
+/* Test the intersection between a ray starting from `origin` traveling in the
+   direction `direction` and a triangle defined by the vertices `v0`, `v1`, and
+   `v2`. If the ray intersects the triangle, set `distance` to the distance
+   between `origin` and the intersection and return true, else return false.
+
+   `direction` must be normalized. */
 __device__ bool intersect_triangle(const float3 &origin, const float3 &direction, const float3 &v0, const float3 &v1, const float3 &v2, float &distance)
 {
 	Matrix m = make_matrix(v1-v0, v2-v0, -direction);
@@ -29,26 +38,38 @@ __device__ bool intersect_triangle(const float3 &origin, const float3 &direction
 
 	float3 b = origin-v0;
 
-	float u1 = ((m.a11*m.a22 - m.a12*m.a21)*b.x + (m.a02*m.a21 - m.a01*m.a22)*b.y + (m.a01*m.a12 - m.a02*m.a11)*b.z)/determinant;
+	float u1 = ((m.a11*m.a22 - m.a12*m.a21)*b.x +
+		    (m.a02*m.a21 - m.a01*m.a22)*b.y +
+		    (m.a01*m.a12 - m.a02*m.a11)*b.z)/determinant;
 
 	if (u1 < 0.0f)
 		return false;
 
-	float u2 = ((m.a12*m.a20 - m.a10*m.a22)*b.x + (m.a00*m.a22 - m.a02*m.a20)*b.y + (m.a02*m.a10 - m.a00*m.a12)*b.z)/determinant;
+	float u2 = ((m.a12*m.a20 - m.a10*m.a22)*b.x +
+		    (m.a00*m.a22 - m.a02*m.a20)*b.y +
+		    (m.a02*m.a10 - m.a00*m.a12)*b.z)/determinant;
 
 	if (u2 < 0.0f)
 		return false;
 
-	float u3 = ((m.a10*m.a21 - m.a11*m.a20)*b.x + (m.a01*m.a20 - m.a00*m.a21)*b.y + (m.a00*m.a11 - m.a01*m.a10)*b.z)/determinant;
+	float u3 = ((m.a10*m.a21 - m.a11*m.a20)*b.x +
+		    (m.a01*m.a20 - m.a00*m.a21)*b.y +
+		    (m.a00*m.a11 - m.a01*m.a10)*b.z)/determinant;
 
 	if (u3 < 0.0f || (1-u1-u2) < 0.0f)
 		return false;
 
-	distance = norm(direction*u3);
+	distance = u3;
 
 	return true;
 }
 
+/* Return the 32 bit color associated with the intersection between a ray
+   starting from `origin` traveling in the direction `direction` and the
+   plane defined by the points `v0`, `v1`, and `v2` using the cosine of the
+   angle between the ray and the plane normal to determine the brightness.
+
+   `direction` must be normalized. */
 __device__ int get_color(const float3 &direction, const float3 &v0, const float3& v1, const float3 &v2)
 {
 	float scale = 255*dot(normalize(cross(v1-v0,v2-v0)),-direction);
@@ -61,7 +82,10 @@ __device__ int get_color(const float3 &direction, const float3 &v0, const float3
 	return rgb << 16 | rgb << 8 | rgb;
 }
 
-
+/* Test the intersection between a ray starting from `origin` traveling in the
+   direction `direction` and the axis-aligned box defined by the opposite
+   vertices `lower_bound` and `upper_bound`. If the ray intersects the box
+   return True, else return False. */
 __device__ bool intersect_box(const float3 &origin, const float3 &direction, const float3 &lower_bound, const float3 &upper_bound)
 {
 	float kmin, kmax, kymin, kymax, kzmin, kzmax;
@@ -132,6 +156,9 @@ __device__ bool intersect_box(const float3 &origin, const float3 &direction, con
 	return true;
 }
 
+/* Test the intersection between a ray starting at `origin` traveling in the
+   direction `direction` and the bounding box around node `i`. If the ray
+   intersects the bounding box return true, else return false. */
 __device__ bool intersect_node(const float3 &origin, const float3 &direction, const int &i)
 {
 	float3 lower_bound = make_float3(tex1Dfetch(lower_bound_arr, i));
@@ -236,7 +263,7 @@ __global__ void intersect_mesh(int max_idx, float3 *origin_arr, float3 *directio
 					if (!hit)
 					{
 						*pixel = get_color(direction, v0, v1, v2);
-						*state = mesh_idx;
+						*state = child_map + i;
 
 						min_distance = distance;
 
@@ -247,7 +274,7 @@ __global__ void intersect_mesh(int max_idx, float3 *origin_arr, float3 *directio
 					if (distance < min_distance)
 					{
 						*pixel = get_color(direction, v0, v1, v2);
-						*state = mesh_idx;
+						*state = child_map + i;
 						
 						min_distance = distance;
 					}
@@ -263,7 +290,10 @@ __global__ void intersect_mesh(int max_idx, float3 *origin_arr, float3 *directio
 	while (node != head);
 
 	if (!hit)
+	{
+		*state = -1;
 		*pixel = 0;
+	}
 
 } // intersect mesh
 
diff --git a/stl.py b/stl.py
index 276a645..a458708 100644
--- a/stl.py
+++ b/stl.py
@@ -3,6 +3,7 @@ import string
 import struct
 
 def read_stl(filename):
+    """Return a triangle mesh from `filename`."""
     f = open(filename)
     buf = f.read(200)
     f.close()
diff --git a/transform.py b/transform.py
index 4e5eb9c..cab23fb 100644
--- a/transform.py
+++ b/transform.py
@@ -1,6 +1,12 @@
 import numpy as np
 
 def rotate(x, phi, n):
+    """
+    Rotate an array of points `x` through an angle phi counter-clockwise
+    around the axis `n` (when looking towards +infinity).
+
+    Source: Weisstein, Eric W. "Rotation Formula." Mathworld.
+    """
     x = np.asarray(x)
     n = np.asarray(n)
 
diff --git a/vector.py b/vector.py
deleted file mode 100644
index 7b36e07..0000000
--- a/vector.py
+++ /dev/null
@@ -1,13 +0,0 @@
-import numpy as np
-from pycuda import gpuarray
-
-def make_vector(arr, dtype=gpuarray.vec.float3):
-    if len(arr.shape) != 2 or arr.shape[-1] != 3:
-        raise Exception('shape mismatch')
-
-    x = np.empty(arr.shape[0], dtype)
-    x['x'] = arr[:,0]
-    x['y'] = arr[:,1]
-    x['z'] = arr[:,2]
-
-    return x
diff --git a/view.py b/view.py
index 6ebf397..b3632de 100644..100755
--- a/view.py
+++ b/view.py
@@ -1,108 +1,111 @@
+#!/usr/bin/env python
 import numpy as np
 from stl import *
 from geometry import *
 from materials import *
 from camera import *
 from gpu import *
-
-import sys
-import optparse
+from transform import *
 
 import pygame
 from pygame.locals import *
 
-parser = optparse.OptionParser('%prog filename.stl')
-parser.add_option('-b', '--bits', type='int', dest='bits', help='number of bits for z ordering space axes', default=4)
-options, args = parser.parse_args()
+def view(geometry, name=''):
+    mesh = geometry.mesh
+    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])])
 
-if len(args) < 1:
-    sys.exit(parser.format_help())
+    scale = np.linalg.norm(upper_bound-lower_bound)/10.0
 
-geometry = Geometry()
-geometry.add_solid(Solid(read_stl(args[0]), vacuum, vacuum))
-geometry.build(options.bits)
+    gpu = GPU()
+    gpu.load_geometry(geometry)
 
-mesh = geometry.mesh
-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])])
+    pygame.init()
+    size = width, height = 800, 600
+    screen = pygame.display.set_mode(size)
+    pygame.display.set_caption(name)
 
-scale = np.linalg.norm(upper_bound-lower_bound)/10.0
+    camera = Camera(size)
+    camera.position((0, upper_bound[1]*5, np.mean([lower_bound[2], upper_bound[2]])))
 
-gpu = GPU()
-gpu.load_geometry(geometry)
+    origin, direction = camera.get_rays()
 
-pygame.init()
-size = width, height = 800, 600
-screen = pygame.display.set_mode(size)
-pygame.display.set_caption(os.path.split(args[0])[-1])
+    pixels = np.empty(width*height, dtype=np.int32)
+    states = np.empty(width*height, dtype=np.int32)
 
-camera = Camera(size)
-camera.position((0, upper_bound[1]*5, np.mean([lower_bound[2], upper_bound[2]])))
+    origin_gpu = cuda.to_device(make_vector(origin))
+    direction_gpu = cuda.to_device(make_vector(direction))
 
-origin, direction = camera.get_rays()
+    pixels_gpu = cuda.to_device(pixels)
+    states_gpu = cuda.to_device(states)
 
-pixels = np.empty(width*height, dtype=np.int32)
-states = np.empty(width*height, dtype=np.int32)
+    block_size = 64
+    gpu_kwargs = {'block': (block_size,1,1), 'grid': (pixels.size//block_size+1,1)}
 
-origin_gpu = cuda.to_device(make_vector(origin))
-direction_gpu = cuda.to_device(make_vector(direction))
+    def render():
+        gpu.call(np.int32(pixels.size), origin_gpu, direction_gpu, np.int32(geometry.first_leaf), states_gpu, pixels_gpu, **gpu_kwargs)
 
-pixels_gpu = cuda.to_device(pixels)
-states_gpu = cuda.to_device(states)
+        cuda.memcpy_dtoh(pixels, pixels_gpu)
 
-block_size = 64
-gpu_kwargs = {'block': (block_size,1,1), 'grid': (pixels.size//block_size+1,1)}
+        pygame.surfarray.blit_array(screen, pixels.reshape(size))
+        pygame.display.flip()
 
-def render():
-    gpu.call(np.int32(pixels.size), origin_gpu, direction_gpu, np.int32(geometry.first_leaf), states_gpu, pixels_gpu, **gpu_kwargs)
+    render()
 
-    cuda.memcpy_dtoh(pixels, pixels_gpu)
-    
-    pygame.surfarray.blit_array(screen, pixels.reshape(size))
-    pygame.display.flip()
+    done = False
+    clicked = False
+    to_origin = np.array([0,-1,0])
+    while not done:
+        for event in pygame.event.get():
+            if event.type == MOUSEBUTTONDOWN:
+                if event.button == 4:
+                    cuda_translate(np.int32(pixels.size), origin_gpu, gpuarray.vec.make_float3(*(scale*to_origin)), **gpu_kwargs)
+                    render()
 
-render()
+                if event.button == 5:
+                    cuda_translate(np.int32(pixels.size), origin_gpu, gpuarray.vec.make_float3(*(-scale*to_origin)), **gpu_kwargs)
+                    render()
 
-from transform import *
+                if event.button == 1:
+                    clicked = True
+                    mouse_position = pygame.mouse.get_rel()
 
-done = False
-clicked = False
-to_origin = np.array([0,-1,0])
-while not done:
-    for event in pygame.event.get():
-        if event.type == MOUSEBUTTONDOWN:
-            if event.button == 4:
-                cuda_translate(np.int32(pixels.size), origin_gpu, gpuarray.vec.make_float3(*(scale*to_origin)), **gpu_kwargs)
-                render()
+            if event.type == MOUSEBUTTONUP:
+                if event.button == 1:
+                    clicked = False
 
-            if event.button == 5:
-                cuda_translate(np.int32(pixels.size), origin_gpu, gpuarray.vec.make_float3(*(-scale*to_origin)), **gpu_kwargs)
-                render()
+            if event.type == MOUSEMOTION and clicked:
+                movement = pygame.mouse.get_rel()[0]/float(width)
+
+                if movement == 0:
+                    continue
 
-            if event.button == 1:
-                clicked = True
-                mouse_position = pygame.mouse.get_rel()
+                cuda_rotate(np.int32(pixels.size), origin_gpu, np.float32(movement*2*np.pi), gpuarray.vec.make_float3(0,0,1), **gpu_kwargs)
 
-        if event.type == MOUSEBUTTONUP:
-            if event.button == 1:
-                clicked = False
+                cuda_rotate(np.int32(pixels.size), direction_gpu, np.float32(movement*2*np.pi), gpuarray.vec.make_float3(0,0,1), **gpu_kwargs)
 
-        if event.type == MOUSEMOTION and clicked:
-            movement = pygame.mouse.get_rel()[0]/float(width)
+                to_origin = rotate(to_origin, movement*2*np.pi, (0,0,1))
 
-            if movement == 0:
-                continue
+                render()
 
-            cuda_rotate(np.int32(pixels.size), origin_gpu, np.float32(movement*2*np.pi), gpuarray.vec.make_float3(0,0,1), **gpu_kwargs)
+            if event.type == KEYUP or event.type == KEYDOWN:
+                if event.key == K_ESCAPE:
+                    done = True
+                    break
 
-            cuda_rotate(np.int32(pixels.size), direction_gpu, np.float32(movement*2*np.pi), gpuarray.vec.make_float3(0,0,1), **gpu_kwargs)
+if __name__ == '__main__':
+    import sys
+    import optparse
 
-            to_origin = rotate(to_origin, movement*2*np.pi, (0,0,1))
+    parser = optparse.OptionParser('%prog filename.stl')
+    parser.add_option('-b', '--bits', type='int', dest='bits', help='number of bits for z ordering space axes', default=4)
+    options, args = parser.parse_args()
 
-            render()
+    if len(args) < 1:
+        sys.exit(parser.format_help())
 
-        if event.type == KEYUP or event.type == KEYDOWN:
-            if event.key == K_ESCAPE:
-                done = True
-                break
+    geometry = Geometry()
+    geometry.add_solid(Solid(read_stl(args[0]), vacuum, vacuum))
+    geometry.build(options.bits)
 
+    view(geometry, os.path.split(args[0])[-1])