added stl mesh viewer

4 files changed, 186 insertions, 106 deletions

diff --git a/gpu.py b/gpu.py
new file mode 100644
index 0000000..dbcb060
--- /dev/null
+++ b/gpu.py
@@ -0,0 +1,76 @@
+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
+
+def make_vector(arr, dtype=gpuarray.vec.float3):
+    if len(arr.shape) != 2 or arr.shape[-1] != 3:
+        raise Exception('shape mismatch')
+
+    v = np.empty(arr.shape[0], dtype)
+    v['x'] = arr[:,0]
+    v['y'] = arr[:,1]
+    v['z'] = arr[:,2]
+
+    return v
+
+print 'device %s' % autoinit.device.name()
+
+source_directory = os.path.split(os.path.realpath(__file__))[0] + '/src'
+
+source = open(source_directory + '/kernel.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):
+    def __init__(self):
+        pass
+
+    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.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.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')
+
+        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.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.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)
diff --git a/src/kernel.cu b/src/kernel.cu
index d8f2300..c2b3fb2 100644
--- a/src/kernel.cu
+++ b/src/kernel.cu
@@ -163,7 +163,7 @@ __global__ void rotate(int max_idx, float3 *pt, float phi, float3 axis)
 	pt[idx] = rotate(pt[idx], phi, axis);
 }
 
-__global__ void intersect_mesh(int max_idx, float3 *origin_arr, float3 *direction_arr, int *pixel_arr, int first_leaf, int *state_arr)
+__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;
 
@@ -172,6 +172,7 @@ __global__ void intersect_mesh(int max_idx, float3 *origin_arr, float3 *directio
 
 	float3 origin = origin_arr[idx];
 	float3 direction = direction_arr[idx];
+	direction /= norm(direction);
 
 	int *pixel = pixel_arr+idx;
 	int *state = state_arr+idx;
diff --git a/test.py b/test.py
deleted file mode 100644
index 236d1ec..0000000
--- a/test.py
+++ /dev/null
@@ -1,105 +0,0 @@
-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
-
-from stl import *
-from geometry import *
-from materials import *
-from camera import *
-from vector import *
-import detectors
-
-import pygame
-
-print 'device %s' % autoinit.device.name()
-
-source_directory = os.path.split(os.path.realpath(__file__))[0] + '/src'
-
-source = open(source_directory + '/kernel.cu').read()
-mod = SourceModule(source, options=['-I ' + source_directory], no_extern_c=True, cache_dir=False)
-intersect_mesh = mod.get_function('intersect_mesh')
-rotate = mod.get_function('rotate')
-translate = mod.get_function('translate')
-
-size = width, height = 800, 600
-screen = pygame.display.set_mode(size)
-camera = Camera(size)
-camera.position((0,100,10))
-
-origin, direction = camera.get_rays()
-
-for i in range(direction.shape[0]):
-    direction[i] /= np.linalg.norm(direction[i])
-
-origin, direction = make_vector(origin), make_vector(direction)
-
-origin_gpu = cuda.to_device(origin)
-direction_gpu = cuda.to_device(direction)
-
-geometry = detectors.build_lbne(bits=6)
-
-mesh = geometry.mesh
-mesh = mesh.reshape(mesh.shape[0]*3,3)
-mesh = make_vector(mesh, dtype=gpuarray.vec.float4)
-lower_bound = make_vector(geometry.lower_bound, dtype=gpuarray.vec.float4)
-upper_bound = make_vector(geometry.upper_bound, dtype=gpuarray.vec.float4)
-child_map = geometry.child_map.astype(np.uint32)
-child_len = geometry.child_len.astype(np.uint32)
-first_leaf = np.int32(geometry.first_leaf)
-
-mesh_gpu = cuda.to_device(mesh)
-lower_bound_gpu = cuda.to_device(lower_bound)
-upper_bound_gpu = cuda.to_device(upper_bound)
-child_map_gpu = cuda.to_device(child_map)
-child_len_gpu = cuda.to_device(child_len)
-
-mesh_tex = mod.get_texref('mesh')
-lower_bound_tex = mod.get_texref('lower_bound_arr')
-upper_bound_tex = mod.get_texref('upper_bound_arr')
-child_map_tex = mod.get_texref('child_map_arr')
-child_len_tex = mod.get_texref('child_len_arr')
-
-mesh_tex.set_address(mesh_gpu, mesh.nbytes)
-lower_bound_tex.set_address(lower_bound_gpu, lower_bound.nbytes)
-upper_bound_tex.set_address(upper_bound_gpu, upper_bound.nbytes)
-child_map_tex.set_address(child_map_gpu, child_map.nbytes)
-child_len_tex.set_address(child_len_gpu, child_len.nbytes)
-
-mesh_tex.set_format(cuda.array_format.FLOAT, 4)
-lower_bound_tex.set_format(cuda.array_format.FLOAT, 4)
-upper_bound_tex.set_format(cuda.array_format.FLOAT, 4)
-child_map_tex.set_format(cuda.array_format.UNSIGNED_INT32, 1)
-child_len_tex.set_format(cuda.array_format.UNSIGNED_INT32, 1)
-
-pixels = np.empty(width*height, dtype=np.int32)
-pixels_gpu = cuda.to_device(pixels)
-
-states = np.empty(width*height, dtype=np.int32)
-states_gpu = cuda.to_device(states)
-
-block_size = 64
-for i in range(100000):
-    rotate(np.int32(origin.size), origin_gpu, np.float32(np.pi/100), gpuarray.vec.make_float3(0,0,1), block=(block_size,1,1), grid=(width*height//block_size+1,1))
-
-    rotate(np.int32(direction.size), direction_gpu, np.float32(np.pi/100), gpuarray.vec.make_float3(0,0,1), block=(block_size,1,1), grid=(width*height//block_size+1,1))
-
-    #translate(np.int32(origin.size), origin_gpu, gpuarray.vec.make_float3(-np.sin(i*np.pi/100)/2, -np.cos(i*np.pi/100)/2, 0), block=(block_size,1,1), grid=(width*height//block_size+1,1))
-
-    t0 = time.time()
-    intersect_mesh(np.int32(origin.size), origin_gpu, direction_gpu, pixels_gpu, first_leaf, states_gpu, block=(block_size,1,1), grid=(width*height//block_size+1,1), texrefs=[mesh_tex, lower_bound_tex, upper_bound_tex, child_map_tex, child_len_tex])
-    cuda.Context.synchronize()
-    elapsed = time.time() - t0
-
-    print '%i triangles, %i photons, %f sec; (%f photons/sec)' % \
-        (mesh.size//3, pixels.size, elapsed, pixels.size/elapsed)
-
-    cuda.memcpy_dtoh(pixels, pixels_gpu)
-    pygame.surfarray.blit_array(screen, pixels.reshape(size))
-    pygame.display.flip()
-
-raw_input('press enter to exit')
diff --git a/view.py b/view.py
new file mode 100644
index 0000000..6ebf397
--- /dev/null
+++ b/view.py
@@ -0,0 +1,108 @@
+import numpy as np
+from stl import *
+from geometry import *
+from materials import *
+from camera import *
+from gpu import *
+
+import sys
+import optparse
+
+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()
+
+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)
+
+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])])
+
+scale = np.linalg.norm(upper_bound-lower_bound)/10.0
+
+gpu = GPU()
+gpu.load_geometry(geometry)
+
+pygame.init()
+size = width, height = 800, 600
+screen = pygame.display.set_mode(size)
+pygame.display.set_caption(os.path.split(args[0])[-1])
+
+camera = Camera(size)
+camera.position((0, upper_bound[1]*5, np.mean([lower_bound[2], upper_bound[2]])))
+
+origin, direction = camera.get_rays()
+
+pixels = np.empty(width*height, dtype=np.int32)
+states = np.empty(width*height, dtype=np.int32)
+
+origin_gpu = cuda.to_device(make_vector(origin))
+direction_gpu = cuda.to_device(make_vector(direction))
+
+pixels_gpu = cuda.to_device(pixels)
+states_gpu = cuda.to_device(states)
+
+block_size = 64
+gpu_kwargs = {'block': (block_size,1,1), 'grid': (pixels.size//block_size+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)
+
+    cuda.memcpy_dtoh(pixels, pixels_gpu)
+    
+    pygame.surfarray.blit_array(screen, pixels.reshape(size))
+    pygame.display.flip()
+
+render()
+
+from transform import *
+
+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.button == 5:
+                cuda_translate(np.int32(pixels.size), origin_gpu, gpuarray.vec.make_float3(*(-scale*to_origin)), **gpu_kwargs)
+                render()
+
+            if event.button == 1:
+                clicked = True
+                mouse_position = pygame.mouse.get_rel()
+
+        if event.type == MOUSEBUTTONUP:
+            if event.button == 1:
+                clicked = False
+
+        if event.type == MOUSEMOTION and clicked:
+            movement = pygame.mouse.get_rel()[0]/float(width)
+
+            if movement == 0:
+                continue
+
+            cuda_rotate(np.int32(pixels.size), origin_gpu, np.float32(movement*2*np.pi), gpuarray.vec.make_float3(0,0,1), **gpu_kwargs)
+
+            cuda_rotate(np.int32(pixels.size), direction_gpu, np.float32(movement*2*np.pi), gpuarray.vec.make_float3(0,0,1), **gpu_kwargs)
+
+            to_origin = rotate(to_origin, movement*2*np.pi, (0,0,1))
+
+            render()
+
+        if event.type == KEYUP or event.type == KEYDOWN:
+            if event.key == K_ESCAPE:
+                done = True
+                break
+