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

1 files changed, 105 insertions, 35 deletions

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)