added some more documentation and a more accurate miniature version of lbne

11 files changed, 428 insertions, 367 deletions

diff --git a/__init__.py b/__init__.py
index 409c3ea..d3a7f3d 100644
--- a/__init__.py
+++ b/__init__.py
@@ -5,3 +5,4 @@ import stl
 import view
 import photon
 import gpu
+import layout
diff --git a/detectors/__init__.py b/detectors/__init__.py
index a4fef35..8bcf69b 100644
--- a/detectors/__init__.py
+++ b/detectors/__init__.py
@@ -1 +1 @@
-from lbne import *
+from lbne import LBNE
diff --git a/detectors/lbne.py b/detectors/lbne.py
index cac737a..54e11a9 100644
--- a/detectors/lbne.py
+++ b/detectors/lbne.py
@@ -1,130 +1,80 @@
-import os
 import numpy as np
-import pickle
-from chroma import *
 from copy import deepcopy
+from chroma import layout
+from chroma.stl import read_stl
+from chroma.transform import rotate
+from chroma.geometry import Geometry, Solid
+from chroma.materials import glass, h2o
+from chroma.photon import uniform_sphere
+from chroma.gpu import GPU
+from itertools import product
 from histogram import *
 
-models_directory = os.path.split(os.path.realpath(__file__))[0] + '/../models'
+endcap_spacing = .485
 
-strings = 20
-pmts_per_string = 10
-radius = 10.0
-height = 20.0
+radius = 25.0/10.0
+height = 50.0/10.0
 
-grid_spacing = height/pmts_per_string
+nstrings = 324//10
+pmts_per_string = 102//10
 
-block_size = 64
-
-
-class LBNE(geometry.Geometry):
+class LBNE(Geometry):
+    """Miniature version of the LBNE water Cerenkov detector geometry."""
     def __init__(self):
         super(LBNE, self).__init__()
 
-        pmt_mesh = stl.read_stl(models_directory + '/hamamatsu_12inch.stl')
-        pmt_mesh /= 1000.0
-
-        apmt = geometry.Solid(pmt_mesh, materials.glass, materials.h2o)
-
-        self.pmt_index = []
-        self.pmt_local_axes = []
-        self.pmt_positions = []
-
-        self.pmt_hits = []
+        pmt_mesh = read_stl(layout.models + '/hamamatsu_12inch.stl')/1000.0
+        pmt_solid = Solid(pmt_mesh, glass, h2o)
 
+        # construct the barrel
         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))
+            for j in range(nstrings):
+                pmt = deepcopy(pmt_solid)
+                pmt.mesh += (-radius,0,-height/2+i*height/(pmts_per_string-1))
+                pmt.mesh = rotate(pmt.mesh, j*2*np.pi/nstrings, (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)
+        # construct the top endcap
+        for x, y in np.array(tuple(product(\
+                    np.arange(-radius+0.075, radius, endcap_spacing),
+                    np.arange(-radius+0.075, radius, endcap_spacing)))):
+            if np.sqrt(x**2 + y**2) <= radius:
+                pmt = deepcopy(pmt_solid)
+                pmt.mesh = rotate(pmt.mesh, -np.pi/2, (0,1,0))
+                pmt.mesh += (x,y,+height/2+height/(pmts_per_string-1)/2)
+                self.add_solid(pmt)
 
-        for pmt_hit in self.pmt_hits:
-            pmt_hit.normalize()
+        # construct the bottom endcap
+        for x, y in np.array(tuple(product(\
+                    np.arange(-radius+0.075, radius, endcap_spacing),
+                    np.arange(-radius+0.075, radius, endcap_spacing)))):
+            if np.sqrt(x**2 + y**2) <= radius:
+                pmt = deepcopy(pmt_solid)
+                pmt.mesh = rotate(pmt.mesh, +np.pi/2, (0,1,0))
+                pmt.mesh += (x,y,-height/2-height/(pmts_per_string-1)/2)
+                self.add_solid(pmt)
 
-        likelihood = 0.0
-        for i in range(self.npmts):
-            probability = self.pmt_hits[i].eval(self.bin_count[i])
+    def load_geometry(self):
+        self.gpu = GPU()
+        self.texrefs = self.gpu.load_geometry(self)
+        self.propagate = self.gpu.get_function('propagate')
 
-            if probability == 0.0:
-                print 'calculating likelihood from pmt %i' % i
-                print 'bin count =', self.bin_count[i]
-                print self.pmt_hits[i].hist
+if __name__ == '__main__':
+    import sys
+    import optparse
+    import pickle
 
-            likelihood -= np.log(self.pmt_hits[i].eval(self.bin_count[i]))
+    parser = optparse.OptionParser('prog filename')
+    parser.add_option('-b', '--bits', type='int', dest='bits',
+                      help='bits for z-ordering space axes', default=6)
+    options, args = parser.parse_args()
 
-        return likelihood
+    if len(args) < 1:
+        sys.exit(parser.format_help())
 
-if __name__ == '__main__':
     lbne = LBNE()
-    view.view(lbne)
+    lbne.build(bits=options.bits)
+
+    f = open(args[0], 'wb')
+    pickle.dump(lbne, f)
+    f.close()
diff --git a/geometry.py b/geometry.py
index 2c136a3..3a11daf 100644
--- a/geometry.py
+++ b/geometry.py
@@ -12,6 +12,14 @@ def compress(data, selectors):
     """
     return (d for d, s in zip(data, selectors) if s)
 
+def get_bounds(mesh):
+    """Returns the lower/upper bounds for `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])])
+    return lower_bound, upper_bound
+
 class Leaf(object):
     """
     Leaf object represents the last layer in the bounding volume hierarchy
@@ -33,8 +41,7 @@ class Leaf(object):
         self.zvalue = zvalue
         self.mesh_idx = mesh_idx
 
-        self.lower_bound = np.array([np.min(mesh[:,:,0]), np.min(mesh[:,:,1]), np.min(mesh[:,:,2])])
-        self.upper_bound = np.array([np.max(mesh[:,:,0]), np.max(mesh[:,:,1]), np.max(mesh[:,:,2])])
+        self.lower_bound, self.upper_bound = get_bounds(mesh)
 
         self.size = mesh.shape[0]
 
@@ -86,8 +93,7 @@ def morton_order(mesh, bits):
     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])])
+    lower_bound, upper_bound = get_bounds(mesh)
 
     max_value = 2**bits - 1
     
@@ -227,15 +233,15 @@ class Geometry(object):
                 nodes.append(node)
 
 
-        self.lower_bound = np.empty((len(nodes),3))
-        self.upper_bound = np.empty((len(nodes),3))
+        self.lower_bounds = np.empty((len(nodes),3))
+        self.upper_bounds = np.empty((len(nodes),3))
         self.child_map = np.empty(len(nodes))
         self.child_len = np.empty(len(nodes))
         self.first_leaf = None
 
         for i, node in enumerate(nodes):
-            self.lower_bound[i] = node.lower_bound
-            self.upper_bound[i] = node.upper_bound
+            self.lower_bounds[i] = node.lower_bound
+            self.upper_bounds[i] = node.upper_bound
             
             self.child_len[i] = node.size
 
diff --git a/gpu.py b/gpu.py
index 3d30387..6ab525d 100644
--- a/gpu.py
+++ b/gpu.py
@@ -1,13 +1,17 @@
 import os
-import time
 import numpy as np
-
 from pycuda import autoinit
 from pycuda.compiler import SourceModule
 import pycuda.driver as cuda
 from pycuda import gpuarray
+import layout
+
+float3 = gpuarray.vec.float3
+float4 = gpuarray.vec.float4
 
-def make_vector(arr, dtype=gpuarray.vec.float3):
+source = open(layout.source + '/kernel.cu').read()
+
+def make_vector(arr, dtype=float3):
     if len(arr.shape) != 2 or arr.shape[-1] != 3:
         raise Exception('shape mismatch')
 
@@ -18,59 +22,52 @@ def make_vector(arr, dtype=gpuarray.vec.float3):
 
     return v
 
-print 'device %s' % autoinit.device.name()
-
-source_directory = os.path.split(os.path.realpath(__file__))[0] + '/src'
-
-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')
-cuda_rotate = module.get_function('rotate')
-cuda_translate = module.get_function('translate')
-
 class GPU(object):
+    """
+    Object to handle all of the texture allocation/referencing when loading
+    a geometry onto the GPU.
+    """
     def __init__(self):
-        pass
+        print 'device %s' % autoinit.device.name()
 
-    def load_geometry(self, geometry):
-        self.mesh = geometry.mesh
-        self.lower_bound = geometry.lower_bound
-        self.upper_bound = geometry.upper_bound
-        self.child_map = geometry.child_map.astype(np.uint32)
-        self.child_len = geometry.child_len.astype(np.uint32)
-        self.first_leaf = np.int32(geometry.first_leaf)
-
-        self.mesh_vec = make_vector(self.mesh.reshape(self.mesh.shape[0]*3,3), dtype=gpuarray.vec.float4)
-        self.lower_bound_vec = make_vector(geometry.lower_bound, dtype=gpuarray.vec.float4)
-        self.upper_bound_vec = make_vector(geometry.upper_bound, dtype=gpuarray.vec.float4)
+        self.module = SourceModule(source, options=['-I' + layout.source],
+                                   no_extern_c=True, cache_dir=False)
 
-        self.mesh_gpu = cuda.to_device(self.mesh_vec)
-        self.lower_bound_gpu = cuda.to_device(self.lower_bound_vec)
-        self.upper_bound_gpu = cuda.to_device(self.upper_bound_vec)
-        self.child_map_gpu = cuda.to_device(self.child_map)
-        self.child_len_gpu = cuda.to_device(self.child_len)
+        self.get_function = self.module.get_function
 
-        self.mesh_tex = module.get_texref('mesh')
-        self.lower_bound_tex = module.get_texref('lower_bound_arr')
-        self.upper_bound_tex = module.get_texref('upper_bound_arr')
-        self.child_map_tex = module.get_texref('child_map_arr')
-        self.child_len_tex = module.get_texref('child_len_arr')
+    def load_geometry(self, geometry):
+        """
+        Load all the textures from `geometry` onto the GPU and return a list
+        of texture references.
+        """
+        self.mesh_vec = make_vector(geometry.mesh.reshape(geometry.mesh.shape[0]*3,3), float4)
+        self.lower_bounds_vec = make_vector(geometry.lower_bounds, float4)
+        self.upper_bounds_vec = make_vector(geometry.upper_bounds, float4)
+        self.uchild_map = geometry.child_map.astype(np.uint32)
+        self.uchild_len = geometry.child_len.astype(np.uint32)
 
+        self.mesh_gpu = cuda.to_device(self.mesh_vec)
+        self.lower_bounds_gpu = cuda.to_device(self.lower_bounds_vec)
+        self.upper_bounds_gpu = cuda.to_device(self.upper_bounds_vec)
+        self.child_map_gpu = cuda.to_device(self.uchild_map)
+        self.child_len_gpu = cuda.to_device(self.uchild_len)
+
+        self.mesh_tex = self.module.get_texref('mesh')
+        self.lower_bounds_tex = self.module.get_texref('lower_bounds')
+        self.upper_bounds_tex = self.module.get_texref('upper_bounds')
+        self.child_map_tex = self.module.get_texref('child_map_arr')
+        self.child_len_tex = self.module.get_texref('child_len_arr')
+        
         self.mesh_tex.set_address(self.mesh_gpu, self.mesh_vec.nbytes)
-        self.lower_bound_tex.set_address(self.lower_bound_gpu, self.lower_bound_vec.nbytes)
-        self.upper_bound_tex.set_address(self.upper_bound_gpu, self.upper_bound_vec.nbytes)
-        self.child_map_tex.set_address(self.child_map_gpu, self.child_map.nbytes)
-        self.child_len_tex.set_address(self.child_len_gpu, self.child_len.nbytes)
+        self.lower_bounds_tex.set_address(self.lower_bounds_gpu, self.lower_bounds_vec.nbytes)
+        self.upper_bounds_tex.set_address(self.upper_bounds_gpu, self.upper_bounds_vec.nbytes)
+        self.child_map_tex.set_address(self.child_map_gpu, self.uchild_map.nbytes)
+        self.child_len_tex.set_address(self.child_len_gpu, self.uchild_len.nbytes)
 
         self.mesh_tex.set_format(cuda.array_format.FLOAT, 4)
-        self.lower_bound_tex.set_format(cuda.array_format.FLOAT, 4)
-        self.upper_bound_tex.set_format(cuda.array_format.FLOAT, 4)
+        self.lower_bounds_tex.set_format(cuda.array_format.FLOAT, 4)
+        self.upper_bounds_tex.set_format(cuda.array_format.FLOAT, 4)
         self.child_map_tex.set_format(cuda.array_format.UNSIGNED_INT32, 1)
         self.child_len_tex.set_format(cuda.array_format.UNSIGNED_INT32, 1)
 
-        self.geometry = geometry
-
-    def call(self, *args, **kwargs):
-        kwargs['texrefs'] = [self.mesh_tex, self.lower_bound_tex, self.upper_bound_tex, self.child_map_tex, self.child_len_tex]
-        cuda_intersect(*args, **kwargs)
+        return [self.mesh_tex, self.lower_bounds_tex, self.upper_bounds_tex, self.child_map_tex, self.child_len_tex]
diff --git a/layout.py b/layout.py
new file mode 100644
index 0000000..e192bd8
--- /dev/null
+++ b/layout.py
@@ -0,0 +1,6 @@
+import os
+
+dir = os.path.split(os.path.realpath(__file__))[0]
+
+models = dir + '/models'
+source = dir + '/src'
diff --git a/photon.py b/photon.py
index b58a97b..d1574ea 100644
--- a/photon.py
+++ b/photon.py
@@ -17,14 +17,11 @@ def uniform_sphere(size=None):
 
     c = np.sqrt(1-u**2)
 
-    x = c*np.cos(theta)
-    y = c*np.sin(theta)
-    z = u
-
     points = np.empty((x.size, 3))
-    points[:,0] = x
-    points[:,1] = y
-    points[:,2] = z
+
+    points[:,0] = c*np.cos(theta)
+    points[:,1] = c*np.sin(theta)
+    points[:,2] = u
 
     if size is None:
         return points[0]
diff --git a/src/intersect.cu b/src/intersect.h
index 1402aa1..b984612 100644
--- a/src/intersect.cu
+++ b/src/intersect.h
@@ -5,22 +5,6 @@
 #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;
-
-__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
@@ -155,146 +139,3 @@ __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));
-	float3 upper_bound = make_float3(tex1Dfetch(upper_bound_arr, i));
-
-	return intersect_box(origin, direction, lower_bound, upper_bound);
-}
-
-extern "C"
-{
-
-__global__ void translate(int max_idx, float3 *pt, float3 v)
-{
-	int idx = blockIdx.x*blockDim.x + threadIdx.x;
-
-	if (idx > max_idx)
-		return;
-
-	pt[idx] += v;
-}
-
-__global__ void rotate(int max_idx, float3 *pt, float phi, float3 axis)
-{
-	int idx = blockIdx.x*blockDim.x + threadIdx.x;
-
-	if (idx > max_idx)
-		return;
-
-	pt[idx] = rotate(pt[idx], phi, axis);
-}
-
-__global__ void intersect_mesh(int max_idx, float3 *origin_arr, float3 *direction_arr, int first_leaf, int *state_arr, int *pixel_arr)
-{
-	int idx = blockIdx.x*blockDim.x + threadIdx.x;
-
-	if (idx > max_idx)
-		return;
-
-	float3 origin = origin_arr[idx];
-	float3 direction = direction_arr[idx];
-	direction /= norm(direction);
-
-	int *pixel = pixel_arr+idx;
-	int *state = state_arr+idx;
-
-	bool hit = false;
-
-	float distance;
-	float min_distance;
-	
-	if (!intersect_node(origin, direction, 0))
-	{
-		*pixel = 0;
-		return;
-	}
-
-	int stack[100];
-
-	int *head = &stack[0];
-	int *node = &stack[1];
-	*node = 0;
-
-	int i;
-
-	bool show_leafs = false;
-
-	do
-	{
-		int child_map = tex1Dfetch(child_map_arr, *node);
-		int child_len = tex1Dfetch(child_len_arr, *node);
-
-		if (*node < first_leaf)
-		{
-			for (i=0; i < child_len; i++)
-				if (intersect_node(origin, direction, child_map+i))
-					*node++ = child_map+i;
-
-			if (node == head)
-				break;
-			
-			node--;
-
-		}
-		else if (show_leafs)
-		{
-			hit = true;
-			*pixel = 255;
-			node--;
-		}
-		else // node is a leaf
-		{
-			for (i=0; i < child_len; i++)
-			{
-				int mesh_idx = 3*(child_map + i);
-
-				float3 v0 = make_float3(tex1Dfetch(mesh, mesh_idx));
-				float3 v1 = make_float3(tex1Dfetch(mesh, mesh_idx+1));
-				float3 v2 = make_float3(tex1Dfetch(mesh, mesh_idx+2));
-				
-				if (intersect_triangle(origin, direction, v0, v1, v2, distance))
-				{
-					if (!hit)
-					{
-						*pixel = get_color(direction, v0, v1, v2);
-						*state = child_map + i;
-
-						min_distance = distance;
-
-						hit = true;
-						continue;
-					}
-					
-					if (distance < min_distance)
-					{
-						*pixel = get_color(direction, v0, v1, v2);
-						*state = child_map + i;
-						
-						min_distance = distance;
-					}
-				}
-
-			} // triangle loop
-
-			node--;
-
-		} // node is a leaf
-
-	} // while loop
-	while (node != head);
-
-	if (!hit)
-	{
-		*state = -1;
-		*pixel = 0;
-	}
-
-} // intersect mesh
-
-} // extern "c"
diff --git a/src/kernel.cu b/src/kernel.cu
new file mode 100644
index 0000000..ed53032
--- /dev/null
+++ b/src/kernel.cu
@@ -0,0 +1,193 @@
+//-*-c-*-
+#include <math_constants.h>
+
+#include "linalg.h"
+#include "matrix.h"
+#include "rotate.h"
+#include "intersect.h"
+
+#define STACK_SIZE 100
+
+/* 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_bounds;
+texture<float4, 1, cudaReadModeElementType> lower_bounds;
+
+/* 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;
+
+__device__ float3 make_float3(const float4 &a)
+{
+	return make_float3(a.x, a.y, a.z);
+}
+
+/* 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_bounds, i));
+	float3 upper_bound = make_float3(tex1Dfetch(upper_bounds, i));
+
+	return intersect_box(origin, direction, lower_bound, upper_bound);
+}
+
+/* Find the intersection between a ray starting at `origin` traveling in the
+   direction `direction` and the global mesh texture. If the ray intersects
+   the texture return the index of the triangle which the ray intersected,
+   else return -1. */
+__device__ int intersect_mesh(const float3 &origin, const float3& direction, const int first_leaf)
+{
+	int triangle_idx = -1;
+
+	float distance;
+	float min_distance;
+
+	if (!intersect_node(origin, direction, 0))
+		return -1;
+
+	int stack[STACK_SIZE];
+
+	int *head = &stack[0];
+	int *node = &stack[1];
+	int *tail = &stack[STACK_SIZE-1];
+	*node = 0;
+
+	int i;
+
+	do
+	{
+		int child_map = tex1Dfetch(child_map_arr, *node);
+		int child_len = tex1Dfetch(child_len_arr, *node);
+
+		if (*node < first_leaf)
+		{
+			for (i=0; i < child_len; i++)
+				if (intersect_node(origin, direction, child_map+i))
+					*node++ = child_map+i;
+
+			if (node == head)
+				break;
+
+			node--;
+		}
+		else // node is a leaf
+		{
+			for (i=0; i < child_len; i++)
+			{
+				int mesh_idx = 3*(child_map + i);
+
+				float3 v0 = make_float3(tex1Dfetch(mesh, mesh_idx));
+				float3 v1 = make_float3(tex1Dfetch(mesh, mesh_idx+1));
+				float3 v2 = make_float3(tex1Dfetch(mesh, mesh_idx+2));
+
+				if (intersect_triangle(origin, direction, v0, v1, v2, distance))
+				{
+					if (triangle_idx == -1)
+					{
+						triangle_idx = child_map + i;
+						min_distance = distance;
+						continue;
+					}
+
+					if (distance < min_distance)
+					{
+						triangle_idx = child_map + i;
+						min_distance = distance;
+					}
+				}
+			} // triangle loop
+
+			node--;
+
+		} // node is a leaf
+
+	} // while loop
+	while (node != head);
+
+	return triangle_idx;
+}
+
+extern "C"
+{
+
+/* Translate `points` by the vector `v` */
+__global__ void translate(int max_idx, float3 *points, float3 v)
+{
+	int idx = blockIdx.x*blockDim.x + threadIdx.x;
+
+	if (idx > max_idx)
+		return;
+
+	*(points+idx) += v;
+}
+
+/* Rotate `points` through an angle `phi` counter-clockwise about the
+   axis `axis` (when looking towards +infinity). */
+__global__ void rotate(int max_idx, float3 *points, float phi, float3 axis)
+{
+	int idx = blockIdx.x*blockDim.x + threadIdx.x;
+
+	if (idx > max_idx)
+		return;
+
+	*(points+idx) = rotate(*(points+idx), phi, axis);
+}
+
+/* Trace the rays starting at `origins` traveling in the direction `directions`
+   to their intersection with the global mesh. If the ray intersects the mesh
+   set the pixel associated with the ray to a 32 bit color whose brightness is
+   determined by the cosine of the angle between the ray and the normal of the
+   triangle it intersected, else set the pixel to 0. */
+__global__ void ray_trace(int max_idx, float3 *origins, float3 *directions, int first_leaf, int *pixels)
+{
+	int idx = blockIdx.x*blockDim.x + threadIdx.x;
+
+	if (idx > max_idx)
+		return;
+
+	float3 origin = *(origins+idx);
+	float3 direction = *(directions+idx);
+	direction /= norm(direction);
+
+	int intersection_idx = intersect_mesh(origin, direction, first_leaf);
+
+	if (intersection_idx == -1)
+	{
+		*(pixels+idx) = 0;
+	}
+	else
+	{
+		int mesh_idx = 3*intersection_idx;
+
+		float3 v0 = make_float3(tex1Dfetch(mesh, mesh_idx));
+		float3 v1 = make_float3(tex1Dfetch(mesh, mesh_idx+1));
+		float3 v2 = make_float3(tex1Dfetch(mesh, mesh_idx+2));
+
+		*(pixels+idx) = get_color(direction, v0, v1, v2);
+	}
+} // ray_trace
+
+/* Propagate the photons starting at `origins` traveling in the direction
+   `directions` to their intersection with the global mesh. If the ray
+   intersects the mesh set the hit_solid array value associated with the
+   photon to the triangle index of the triangle the photon intersected, else
+   set the hit_solid array value to -1. */
+__global__ void propagate(int max_idx, float3 *origins, float3 *directions, int first_leaf, int *hit_solids)
+{
+	int idx = blockIdx.x*blockDim.x + threadIdx.x;
+
+	if (idx > max_idx)
+		return;
+
+	float3 origin = *(origins+idx);
+	float3 direction = *(directions+idx);
+	direction /= norm(direction);
+
+	*(hit_solids+idx) = intersect_mesh(origin, direction, first_leaf);
+} // propagate
+
+} // extern "c"
diff --git a/transform.py b/transform.py
index cab23fb..1321c8c 100644
--- a/transform.py
+++ b/transform.py
@@ -8,7 +8,7 @@ def rotate(x, phi, n):
     Source: Weisstein, Eric W. "Rotation Formula." Mathworld.
     """
     x = np.asarray(x)
-    n = np.asarray(n)
+    n = np.asarray(n)/np.linalg.norm(n)
 
     r = np.cos(phi)*np.identity(3) + (1-np.cos(phi))*np.outer(n,n) + \
         np.sin(phi)*np.array([[0,n[2],-n[1]],[-n[2],0,n[0]],[n[1],-n[0],0]])
diff --git a/view.py b/view.py
index b3632de..2c37d2e 100755
--- a/view.py
+++ b/view.py
@@ -1,24 +1,33 @@
 #!/usr/bin/env python
 import numpy as np
-from stl import *
-from geometry import *
-from materials import *
-from camera import *
-from gpu import *
-from transform import *
 
 import pygame
 from pygame.locals import *
 
+from camera import *
+from geometry import *
+from transform import *
+from gpu import *
+
+
 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])])
+    """
+    Render `geometry` in a pygame window.
 
-    scale = np.linalg.norm(upper_bound-lower_bound)/10.0
+    Movement:
+        - zoom: scroll the mouse wheel
+        - rotate: click and drag the mouse
+        - move: shift+click and drag the mouse
+    """
+    lower_bound, upper_bound = get_bounds(geometry.mesh)
+
+    scale = np.linalg.norm(upper_bound-lower_bound)
 
     gpu = GPU()
     gpu.load_geometry(geometry)
+    cuda_raytrace = gpu.get_function('ray_trace')
+    cuda_rotate = gpu.get_function('rotate')
+    cuda_translate = gpu.get_function('translate')
 
     pygame.init()
     size = width, height = 800, 600
@@ -26,27 +35,33 @@ def view(geometry, name=''):
     pygame.display.set_caption(name)
 
     camera = Camera(size)
-    camera.position((0, upper_bound[1]*5, np.mean([lower_bound[2], upper_bound[2]])))
+
+    diagonal = np.linalg.norm(upper_bound-lower_bound)
+
+    point = np.array([0, diagonal*2, (lower_bound[2]+upper_bound[2])/2])
+    axis1 = np.array([1,0,0], dtype=np.double)
+    axis2 = np.array([0,0,1], dtype=np.double)
+
+    camera.position(point)
 
     origin, direction = camera.get_rays()
 
     pixels = np.empty(width*height, dtype=np.int32)
-    states = np.empty(width*height, dtype=np.int32)
+    pixels_gpu = cuda.to_device(pixels)
 
-    origin_gpu = cuda.to_device(make_vector(origin))
-    direction_gpu = cuda.to_device(make_vector(direction))
+    origins_gpu = cuda.to_device(make_vector(origin))
+    directions_gpu = cuda.to_device(make_vector(direction))
 
-    pixels_gpu = cuda.to_device(pixels)
-    states_gpu = cuda.to_device(states)
+    nblocks = 64
 
-    block_size = 64
-    gpu_kwargs = {'block': (block_size,1,1), 'grid': (pixels.size//block_size+1,1)}
+    gpu_kwargs = {'block': (nblocks,1,1), 'grid':(pixels.size/nblocks+1,1)}
 
     def render():
-        gpu.call(np.int32(pixels.size), origin_gpu, direction_gpu, np.int32(geometry.first_leaf), states_gpu, pixels_gpu, **gpu_kwargs)
+        """Render the mesh and display to screen."""
+        cuda_raytrace(np.int32(pixels.size), origins_gpu, directions_gpu,
+                      np.int32(geometry.first_leaf), pixels_gpu, **gpu_kwargs)
 
         cuda.memcpy_dtoh(pixels, pixels_gpu)
-
         pygame.surfarray.blit_array(screen, pixels.reshape(size))
         pygame.display.flip()
 
@@ -54,16 +69,27 @@ def view(geometry, name=''):
 
     done = False
     clicked = False
-    to_origin = np.array([0,-1,0])
+    shift = False
+
     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)
+                    v = scale*np.cross(axis1,axis2)/10.0
+
+                    cuda_translate(np.int32(pixels.size), origins_gpu, gpuarray.vec.make_float3(*v), **gpu_kwargs)
+
+                    point += v
+
                     render()
 
                 if event.button == 5:
-                    cuda_translate(np.int32(pixels.size), origin_gpu, gpuarray.vec.make_float3(*(-scale*to_origin)), **gpu_kwargs)
+                    v = -scale*np.cross(axis1,axis2)/10.0
+
+                    cuda_translate(np.int32(pixels.size), origins_gpu, gpuarray.vec.make_float3(*v), **gpu_kwargs)
+
+                    point += v
+
                     render()
 
                 if event.button == 1:
@@ -75,37 +101,81 @@ def view(geometry, name=''):
                     clicked = False
 
             if event.type == MOUSEMOTION and clicked:
-                movement = pygame.mouse.get_rel()[0]/float(width)
+                movement = np.array(pygame.mouse.get_rel())
 
-                if movement == 0:
+                if (movement == 0).all():
                     continue
 
-                cuda_rotate(np.int32(pixels.size), origin_gpu, np.float32(movement*2*np.pi), gpuarray.vec.make_float3(0,0,1), **gpu_kwargs)
+                length = np.linalg.norm(movement)
 
-                cuda_rotate(np.int32(pixels.size), direction_gpu, np.float32(movement*2*np.pi), gpuarray.vec.make_float3(0,0,1), **gpu_kwargs)
+                mouse_direction = movement[0]*axis1 + movement[1]*axis2
+                mouse_direction /= np.linalg.norm(mouse_direction)
+
+                if shift:
+                    v = mouse_direction*scale*length/float(width)
+
+                    cuda_translate(np.int32(pixels.size), origins_gpu, gpuarray.vec.make_float3(*v), **gpu_kwargs)
+
+                    point += v
+
+                    render()
+                else:
+                    phi = np.float32(2*np.pi*length/float(width))
+                    n = rotate(mouse_direction, np.pi/2, \
+                                   -np.cross(axis1,axis2))
 
-                to_origin = rotate(to_origin, movement*2*np.pi, (0,0,1))
+                    cuda_rotate(np.int32(pixels.size), origins_gpu, phi, gpuarray.vec.make_float3(*n), **gpu_kwargs)
 
-                render()
+                    cuda_rotate(np.int32(pixels.size), directions_gpu, phi, gpuarray.vec.make_float3(*n), **gpu_kwargs)
+
+                    point = rotate(point, phi, n)
+                    axis1 = rotate(axis1, phi, n)
+                    axis2 = rotate(axis2, phi, n)
+
+                    render()
+
+            if event.type == KEYDOWN:
+                if event.key == K_LSHIFT or event.key == K_RSHIFT:
+                    shift = True
 
-            if event.type == KEYUP or event.type == KEYDOWN:
                 if event.key == K_ESCAPE:
                     done = True
                     break
 
+            if event.type == KEYUP:
+                if event.key == K_LSHIFT or event.key == K_RSHIFT:
+                    shift = False
+
 if __name__ == '__main__':
+    import os
     import sys
     import optparse
 
+    from stl import *
+    from materials 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)
+    parser.add_option('-b', '--bits', type='int', dest='bits',
+                      help='bits for z-ordering space axes', default=6)
     options, args = parser.parse_args()
 
     if len(args) < 1:
         sys.exit(parser.format_help())
 
-    geometry = Geometry()
-    geometry.add_solid(Solid(read_stl(args[0]), vacuum, vacuum))
-    geometry.build(options.bits)
+    head, tail = os.path.split(args[0])
+    root, ext = os.path.splitext(tail)
+
+    if ext.lower() == '.stl':
+        geometry = Geometry()
+        geometry.add_solid(Solid(read_stl(args[0]), vacuum, vacuum))
+        geometry.build(options.bits)
+
+    elif ext.lower() == '.pkl':
+        import pickle
+        from detectors import *
+
+        f = open(args[0], 'rb')
+        geometry = pickle.load(f)
+        f.close()
 
-    view(geometry, os.path.split(args[0])[-1])
+    view(geometry, tail)