merge

1 files changed, 202 insertions, 110 deletions

diff --git a/camera.py b/camera.py
index 03b07f5..fc2f4aa 100755
--- a/camera.py
+++ b/camera.py
@@ -9,15 +9,17 @@ from chroma.geometry import Mesh, Solid, Geometry
 import chroma.make as make
 import matplotlib.pyplot as plt
 
-from chroma.gpu import GPU, CUDAFuncs
+from chroma import gpu
 from chroma.tools import timeit
 from chroma.transform import rotate
 from chroma.optics import vacuum, lambertian_surface
+import chroma.project as project
+from chroma.fileio.root import RootReader
 
 import pygame
 from pygame.locals import *
 
-from pycuda import gpuarray
+from pycuda import gpuarray as ga
 from pycuda.characterize import sizeof
 import pycuda.driver as cuda
 
@@ -98,28 +100,14 @@ def encode_movie(dir):
 
     print 'movie saved to %s.' % filename
 
-def get_rays(position, size = (800, 600), width = 0.035, focal_length=0.018):
-    height = width*(size[1]/float(size[0]))
-
-    x = np.linspace(-width/2, width/2, size[0])
-    z = np.linspace(-height/2, height/2, size[1])
-
-    grid = np.array(tuple(product(x,[0],z)))
-
-    grid += (0,focal_length,0)
-    focal_point = np.zeros(3)
-
-    grid += position
-    focal_point += position
-
-    return grid, focal_point-grid
-
 class Camera(Thread):
-    def __init__(self, geometry, size=(800,600), device_id=None):
+    def __init__(self, geometry, size=(800,600), device_id=None, enable3d=False, green_magenta=False):
         Thread.__init__(self)
         self.geometry = geometry
         self.device_id = device_id
         self.size = size
+        self.enable3d = enable3d
+        self.green_magenta = green_magenta
 
         self.unique_bvh_layers = np.unique(self.geometry.layers)
         self.currentlayer = None
@@ -133,11 +121,14 @@ class Camera(Thread):
             self.spnav = False
 
     def init_gpu(self):
-        self.gpu = GPU(self.device_id)
-        self.gpu.load_geometry(geometry)
+        self.gpu_instance = gpu.GPU(self.device_id)
+        self.gpu_geometry = gpu.GPUGeometry(self.gpu_instance, geometry)
+        self.gpu_funcs = gpu.GPUFuncs(self.gpu_instance.module)
 
         self.width, self.height = size
 
+        self.npixels = self.width*self.height
+
         pygame.init()
         self.screen = pygame.display.set_mode(size)
         pygame.display.set_caption('')
@@ -159,15 +150,38 @@ class Camera(Thread):
 
         self.nblocks = 64
 
-        self.point = np.array([0, self.scale*1.0, (lower_bound[2]+upper_bound[2])/2])
-        self.axis1 = np.array([1,0,0], float)
-        self.axis2 = np.array([0,0,1], float)
+        self.point = np.array([0, -self.scale*1.0, (lower_bound[2]+upper_bound[2])/2])
+        self.axis1 = np.array([0,0,1], float)
+        self.axis2 = np.array([1,0,0], float)
+
+        self.film_width = 0.035
+
+        if self.enable3d:
+            self.point1 = self.point-(self.scale/60,0,0)
+            self.point2 = self.point+(self.scale/60,0,0)
+
+            self.viewing_angle = 0.0
 
-        origins, directions = get_rays(self.point, self.size)
+            pos1, dir1 = project.from_film(self.point1, size=self.size, width=self.film_width)
+            pos2, dir2 = project.from_film(self.point2, size=self.size, width=self.film_width)
 
-        self.origins_gpu = gpuarray.to_gpu(origins.astype(np.float32).view(gpuarray.vec.float3))
-        self.directions_gpu = gpuarray.to_gpu(directions.astype(np.float32).view(gpuarray.vec.float3))
-        self.pixels_gpu = gpuarray.zeros(self.width*self.height, dtype=np.int32)
+            self.rays1 = gpu.GPURays(pos1, dir1)
+            self.rays2 = gpu.GPURays(pos2, dir2)
+
+            scope_pos, scope_dir = project.from_film(self.point, size=np.array(self.size)/4.0, width=self.film_width/4.0)
+
+            self.scope_rays = gpu.GPURays(scope_pos, scope_dir)
+
+            self.pixels1_gpu = ga.empty(self.width*self.height, dtype=np.int32)
+            self.pixels2_gpu = ga.empty(self.width*self.height, dtype=np.int32)
+
+            self.distances_gpu = ga.empty(self.scope_rays.pos.size, dtype=np.float32)
+        else:
+            pos, dir = project.from_film(self.point, size=self.size, width=self.film_width)
+
+            self.rays = gpu.GPURays(pos, dir)
+
+            self.pixels_gpu = ga.empty(self.npixels, dtype=np.int32)
 
         self.alpha = True
         self.movie = False
@@ -177,12 +191,15 @@ class Camera(Thread):
 
     @timeit
     def initialize_render(self):
-        self.rng_states_gpu = cuda.mem_alloc(self.width*self.height*sizeof('curandStateXORWOW', '#include <curand_kernel.h>'))
-        self.gpu.kernels.init_rng(np.int32(self.width*self.height), self.rng_states_gpu, np.int32(0), np.int32(0), block=(self.nblocks,1,1), grid=(self.width*self.height//self.nblocks+1,1))
-        self.xyz_lookup1_gpu = gpuarray.zeros(len(self.geometry.mesh.triangles), dtype=gpuarray.vec.float3)
-        self.xyz_lookup2_gpu = gpuarray.zeros(len(self.geometry.mesh.triangles), dtype=gpuarray.vec.float3)
-        self.image_gpu = gpuarray.zeros(self.width*self.height, dtype=gpuarray.vec.float3)
-        self.gpu.context.synchronize()
+        self.rng_states_gpu = gpu.get_rng_states(self.npixels)
+        self.xyz_lookup1_gpu = ga.zeros(len(self.geometry.mesh.triangles), dtype=ga.vec.float3)
+        self.xyz_lookup2_gpu = ga.zeros(len(self.geometry.mesh.triangles), dtype=ga.vec.float3)
+
+        if self.enable3d:
+            self.image1_gpu = ga.zeros(self.npixels, dtype=ga.vec.float3)
+            self.image2_gpu = ga.zeros(self.npixels, dtype=ga.vec.float3)
+        else:
+            self.image_gpu = ga.zeros(self.npixels, dtype=ga.vec.float3)
 
         self.source_position = self.point
 
@@ -191,39 +208,48 @@ class Camera(Thread):
         self.max_steps = 10
 
     def clear_xyz_lookup(self):
-        self.xyz_lookup1_gpu.fill(gpuarray.vec.make_float3(0.0,0.0,0.0))
-        self.xyz_lookup2_gpu.fill(gpuarray.vec.make_float3(0.0,0.0,0.0))
+        self.xyz_lookup1_gpu.fill(ga.vec.make_float3(0.0,0.0,0.0))
+        self.xyz_lookup2_gpu.fill(ga.vec.make_float3(0.0,0.0,0.0))
 
         self.nlookup_calls = 0
 
     def update_xyz_lookup(self, source_position):
-        for i in range(self.xyz_lookup1_gpu.size//(self.width*self.height)+1):
-            self.gpu.kernels.update_xyz_lookup(np.int32(self.width*self.height), np.int32(self.xyz_lookup1_gpu.size), np.int32(i*self.width*self.height), gpuarray.vec.make_float3(*source_position), self.rng_states_gpu, np.float32(685.0), gpuarray.vec.make_float3(1.0,0.0,0.0), self.xyz_lookup1_gpu, self.xyz_lookup2_gpu, np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.width*self.height//self.nblocks+1,1))
-
-        for i in range(self.xyz_lookup1_gpu.size//(self.width*self.height)+1):
-            self.gpu.kernels.update_xyz_lookup(np.int32(self.width*self.height), np.int32(self.xyz_lookup1_gpu.size), np.int32(i*self.width*self.height), gpuarray.vec.make_float3(*source_position), self.rng_states_gpu, np.float32(545.0), gpuarray.vec.make_float3(0.0,1.0,0.0), self.xyz_lookup1_gpu, self.xyz_lookup2_gpu, np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.width*self.height//self.nblocks+1,1))
-
-        for i in range(self.xyz_lookup1_gpu.size//(self.width*self.height)+1):
-            self.gpu.kernels.update_xyz_lookup(np.int32(self.width*self.height), np.int32(self.xyz_lookup1_gpu.size), np.int32(i*self.width*self.height), gpuarray.vec.make_float3(*source_position), self.rng_states_gpu, np.float32(445.0), gpuarray.vec.make_float3(0.0,0.0,1.0), self.xyz_lookup1_gpu, self.xyz_lookup2_gpu, np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.width*self.height//self.nblocks+1,1))
+        for wavelength, rgb_tuple in \
+                zip([685.0, 545.0, 445.0],[(1,0,0),(0,1,0),(0,0,1)]):
+            for i in range(self.xyz_lookup1_gpu.size//(self.npixels)+1):
+                self.gpu_funcs.update_xyz_lookup(np.int32(self.npixels), np.int32(self.xyz_lookup1_gpu.size), np.int32(i*self.npixels), ga.vec.make_float3(*source_position), self.rng_states_gpu, np.float32(wavelength), ga.vec.make_float3(*rgb_tuple), self.xyz_lookup1_gpu, self.xyz_lookup2_gpu, np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.npixels//self.nblocks+1,1))
 
         self.nlookup_calls += 1
 
     def clear_image(self):
-        self.image_gpu.fill(gpuarray.vec.make_float3(0.0,0.0,0.0))
+        if self.enable3d:
+            self.image1_gpu.fill(ga.vec.make_float3(0.0,0.0,0.0))
+            self.image2_gpu.fill(ga.vec.make_float3(0.0,0.0,0.0))
+        else:
+            self.image_gpu.fill(ga.vec.make_float3(0.0,0.0,0.0))
 
         self.nimages = 0
 
-    def update_image(self):
-        self.gpu.kernels.update_xyz_image(np.int32(self.width*self.height), self.rng_states_gpu, self.origins_gpu, self.directions_gpu, np.float32(685.0), gpuarray.vec.make_float3(1.0,0.0,0.0), self.xyz_lookup1_gpu, self.xyz_lookup2_gpu, self.image_gpu, np.int32(self.nlookup_calls), np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.width*self.height//self.nblocks+1,1))
-
-        self.gpu.kernels.update_xyz_image(np.int32(self.width*self.height), self.rng_states_gpu, self.origins_gpu, self.directions_gpu, np.float32(545.0), gpuarray.vec.make_float3(0.0,1.0,0.0), self.xyz_lookup1_gpu, self.xyz_lookup2_gpu, self.image_gpu, np.int32(self.nlookup_calls), np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.width*self.height//self.nblocks+1,1))
+    def update_image_from_rays(self, image, rays):
+        for wavelength, rgb_tuple in \
+                zip([685.0, 545.0, 445.0],[(1,0,0),(0,1,0),(0,0,1)]):
+            self.gpu_funcs.update_xyz_image(np.int32(rays.pos.size), self.rng_states_gpu, rays.pos, rays.dir, np.float32(wavelength), ga.vec.make_float3(*rgb_tuple), self.xyz_lookup1_gpu, self.xyz_lookup2_gpu, image, np.int32(self.nlookup_calls), np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(rays.pos.size//self.nblocks+1,1))
 
-        self.gpu.kernels.update_xyz_image(np.int32(self.width*self.height), self.rng_states_gpu, self.origins_gpu, self.directions_gpu, np.float32(445.0), gpuarray.vec.make_float3(0.0,0.0,1.0), self.xyz_lookup1_gpu, self.xyz_lookup2_gpu, self.image_gpu, np.int32(self.nlookup_calls), np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.width*self.height//self.nblocks+1,1))
+    def update_image(self):
+        if self.enable3d:
+            self.update_image_from_rays(self.image1_gpu, self.rays1)
+            self.update_image_from_rays(self.image2_gpu, self.rays2)
+        else:
+            self.update_image_from_rays(self.image_gpu, self.rays)
 
         self.nimages += 1
 
     def process_image(self):
-        self.gpu.kernels.process_image(np.int32(self.width*self.height), self.image_gpu, self.pixels_gpu, np.int32(self.nimages), block=(self.nblocks,1,1), grid=((self.width*self.height)//self.nblocks+1,1))
+        if self.enable3d:
+            self.gpu_funcs.process_image(np.int32(self.pixels1_gpu.size), self.image1_gpu, self.pixels1_gpu, np.int32(self.nimages), block=(self.nblocks,1,1), grid=((self.pixels1_gpu.size)//self.nblocks+1,1))
+            self.gpu_funcs.process_image(np.int32(self.pixels2_gpu.size), self.image2_gpu, self.pixels2_gpu, np.int32(self.nimages), block=(self.nblocks,1,1), grid=((self.pixels2_gpu.size)//self.nblocks+1,1))
+        else:
+            self.gpu_funcs.process_image(np.int32(self.pixels_gpu.size), self.image_gpu, self.pixels_gpu, np.int32(self.nimages), block=(self.nblocks,1,1), grid=((self.pixels_gpu.size)//self.nblocks+1,1))
 
     def screenshot(self, dir='', start=0):
         root, ext = 'screenshot', 'png'
@@ -238,8 +264,15 @@ class Camera(Thread):
         print 'image saved to %s' % filename
 
     def rotate(self, phi, n):
-        self.gpu.kernels.rotate(np.int32(self.pixels_gpu.size), self.origins_gpu, np.float32(phi), gpuarray.vec.make_float3(*n), block=(self.nblocks,1,1), grid=(self.pixels_gpu.size//self.nblocks+1,1))
-        self.gpu.kernels.rotate(np.int32(self.pixels_gpu.size), self.directions_gpu, np.float32(phi), gpuarray.vec.make_float3(*n), block=(self.nblocks,1,1), grid=(self.pixels_gpu.size//self.nblocks+1,1))
+        if self.enable3d:
+            self.rays1.rotate(phi, n)
+            self.rays2.rotate(phi, n)
+            self.scope_rays.rotate(phi, n)
+
+            self.point1 = rotate(self.point1, phi, n)
+            self.point2 = rotate(self.point2, phi, n)
+        else:
+            self.rays.rotate(phi, n)
 
         self.point = rotate(self.point, phi, n)
         self.axis1 = rotate(self.axis1, phi, n)
@@ -251,12 +284,16 @@ class Camera(Thread):
         self.update()
 
     def rotate_around_point(self, phi, n, point, redraw=True):
-        self.gpu.kernels.rotate_around_point(np.int32(self.origins_gpu.size), self.origins_gpu, np.float32(phi), gpuarray.vec.make_float3(*n), gpuarray.vec.make_float3(*point), block=(self.nblocks,1,1), grid=(self.origins_gpu.size//self.nblocks+1,1))
-        self.gpu.kernels.rotate(np.int32(self.directions_gpu.size), self.directions_gpu, np.float32(phi), gpuarray.vec.make_float3(*n), block=(self.nblocks,1,1), grid=(self.directions_gpu.size//self.nblocks+1,1))
-
         self.axis1 = rotate(self.axis1, phi, n)
         self.axis2 = rotate(self.axis2, phi, n)
 
+        if self.enable3d:
+            self.rays1.rotate_around_point(phi, n, point)
+            self.rays2.rotate_around_point(phi, n, point)
+            self.scope_rays.rotate_around_point(phi, n, point)
+        else:
+            self.rays.rotate_around_point(phi, n, point)
+
         if redraw:
             if self.render:
                 self.clear_image()
@@ -264,17 +301,25 @@ class Camera(Thread):
             self.update()
 
     def translate(self, v, redraw=True):
-        self.gpu.kernels.translate(np.int32(self.pixels_gpu.size), self.origins_gpu, gpuarray.vec.make_float3(*v), block=(self.nblocks,1,1), grid=(self.pixels_gpu.size//self.nblocks,1))
-
         self.point += v
 
+        if self.enable3d:
+            self.rays1.translate(v)
+            self.rays2.translate(v)
+            self.scope_rays.translate(v)
+
+            self.point1 += v
+            self.point2 += v
+        else:
+            self.rays.translate(v)
+
         if redraw:
             if self.render:
                 self.clear_image()
 
             self.update()
 
-    def update(self):
+    def update_pixels(self):
         if self.render:
             while self.nlookup_calls < 10:
                 self.update_xyz_lookup(self.source_position)
@@ -282,11 +327,69 @@ class Camera(Thread):
             self.process_image()
         else:
             if self.alpha:
-                self.gpu.kernels.ray_trace_alpha(np.int32(self.pixels_gpu.size), self.origins_gpu, self.directions_gpu, self.pixels_gpu, block=(self.nblocks,1,1), grid=(self.pixels_gpu.size//self.nblocks+1,1))
+                if self.enable3d:
+                    self.gpu_funcs.ray_trace_alpha(np.int32(self.rays1.pos.size), self.rays1.pos, self.rays1.dir, self.pixels1_gpu, block=(self.nblocks,1,1), grid=(self.rays2.pos.size//self.nblocks+1,1))
+                    self.gpu_funcs.ray_trace_alpha(np.int32(self.rays2.pos.size), self.rays2.pos, self.rays2.dir, self.pixels2_gpu, block=(self.nblocks,1,1), grid=(self.rays2.pos.size//self.nblocks+1,1))
+                else:
+                    self.gpu_funcs.ray_trace_alpha(np.int32(self.rays.pos.size), self.rays.pos, self.rays.dir, self.pixels_gpu, block=(self.nblocks,1,1), grid=(self.rays.pos.size//self.nblocks+1,1))
+            else:
+                if self.enable3d:
+                    self.gpu_funcs.ray_trace(np.int32(self.rays1.pos.size), self.rays1.pos, self.rays1.dir, self.pixels1_gpu, block=(self.nblocks,1,1), grid=(self.rays1.pos.size//self.nblocks+1,1))
+                    self.gpu_funcs.ray_trace(np.int32(self.rays2.pos.size), self.rays2.pos, self.rays2.dir, self.pixels2_gpu, block=(self.nblocks,1,1), grid=(self.rays2.pos.size//self.nblocks+1,1))
+                else:
+                    self.gpu_funcs.ray_trace(np.int32(self.rays.pos.size), self.rays.pos, self.rays.dir, self.pixels_gpu, block=(self.nblocks,1,1), grid=(self.rays.pos.size//self.nblocks+1,1))
+
+    def update(self):
+        if self.enable3d:
+            self.gpu_funcs.distance_to_mesh(np.int32(self.scope_rays.pos.size), self.scope_rays.pos, self.scope_rays.dir, self.distances_gpu, block=(self.nblocks,1,1), grid=(self.scope_rays.pos.size//self.nblocks,1))
+
+            baseline = ga.min(self.distances_gpu).get().item()
+
+            if baseline < 1e9:
+                d1 = self.point1 - self.point
+                v1 = d1/np.linalg.norm(d1)
+                v1 *= baseline/60 - np.linalg.norm(d1)
+
+                self.rays1.translate(v1)
+
+                self.point1 += v1
+
+                d2 = self.point2 - self.point
+                v2 = d2/np.linalg.norm(d2)
+                v2 *= baseline/60 - np.linalg.norm(d2)
+
+                self.rays2.translate(v2)
+
+                self.point2 += v2
+
+            direction = np.cross(self.axis1,self.axis2)
+            direction /= np.linalg.norm(direction)
+            direction1 = self.point + direction*baseline - self.point1
+            direction1 /= np.linalg.norm(direction1)
+
+            new_viewing_angle = np.arccos(direction1.dot(direction))
+
+            phi = new_viewing_angle - self.viewing_angle
+
+            self.rays1.rotate_around_point(phi, self.axis1, self.point1)
+            self.rays2.rotate_around_point(-phi, self.axis1, self.point2)
+
+            self.viewing_angle = new_viewing_angle
+
+        self.update_pixels()
+
+        if self.enable3d:
+            pixels1 = self.pixels1_gpu.get()
+            pixels2 = self.pixels2_gpu.get()
+
+            if self.green_magenta:
+                pixels = (pixels1 & 0x00ff00) | (pixels2 & 0xff00ff)
             else:
-                self.gpu.kernels.ray_trace(np.int32(self.pixels_gpu.size), self.origins_gpu, self.directions_gpu, self.pixels_gpu, block=(self.nblocks,1,1), grid=(self.pixels_gpu.size//self.nblocks+1,1))
+                pixels = (pixels1 & 0xff0000) | (pixels2 & 0x00ffff)
+        else:
+            pixels = self.pixels_gpu.get()
 
-        pygame.surfarray.blit_array(self.screen, self.pixels_gpu.get().reshape(self.size))
+        pygame.surfarray.blit_array(self.screen, pixels.reshape(self.size))
         if self.doom_mode:
             self.screen.blit(self.doom_hud, self.doom_rect)
         pygame.display.flip()
@@ -298,11 +401,11 @@ class Camera(Thread):
     def loadlayer(self, layer):
         try:
             layergeometry = self.bvh_layers[layer]
+            layergeometry.activate()
         except KeyError:
             layergeometry = build(bvh_mesh(self.geometry, layer),8)
-            self.bvh_layers[layer] = layergeometry
+            self.bvh_layers[layer] = gpu.GPUGeometry(self.gpu, layergeometry)
 
-        self.gpu.load_geometry(layergeometry)
         self.update()
 
     def process_event(self, event):
@@ -331,15 +434,15 @@ class Camera(Thread):
 
             length = np.linalg.norm(movement)
 
-            mouse_direction = movement[0]*self.axis1 + movement[1]*self.axis2
+            mouse_direction = movement[0]*self.axis2 - movement[1]*self.axis1
             mouse_direction /= np.linalg.norm(mouse_direction)
 
             if pygame.key.get_mods() & (KMOD_LSHIFT | KMOD_RSHIFT):
-                v = mouse_direction*self.scale*length/float(self.width)
+                v = -mouse_direction*self.scale*length/float(self.width)
                 self.translate(v)
             else:
                 phi = np.float32(2*np.pi*length/float(self.width))
-                n = rotate(mouse_direction, np.pi/2, -np.cross(self.axis1,self.axis2))
+                n = rotate(mouse_direction, np.pi/2, np.cross(self.axis1,self.axis2))
 
                 if pygame.key.get_mods() & KMOD_LCTRL:
                     self.rotate_around_point(phi, n, self.point)
@@ -348,11 +451,11 @@ class Camera(Thread):
 
         elif event.type == KEYDOWN:
             if event.key == K_a:
-                v = self.scale*self.axis1/10.0
+                v = self.scale*self.axis2/10.0
                 self.translate(v)
 
             elif event.key == K_d:
-                v = -self.scale*self.axis1/10.0
+                v = -self.scale*self.axis2/10.0
                 self.translate(v)
 
             elif event.key == K_w:
@@ -363,10 +466,6 @@ class Camera(Thread):
                 v = -self.scale*np.cross(self.axis1,self.axis2)/10.0
                 self.translate(v)
 
-            elif event.key == K_SPACE:
-                v = self.scale*self.axis2/10.0
-                self.translate(v)
-
             elif event.key == K_F6:
                 self.clear_xyz_lookup()
                 self.clear_image()
@@ -443,10 +542,10 @@ class Camera(Thread):
                 else:
                     accelerate_factor = 1.0
                 #print 'raw:', spnav_event
-                v1 = self.axis1
-                v2 = self.axis2
-                v3 = np.cross(self.axis1,self.axis2)
-
+                v1 = self.axis2
+                v2 = self.axis1
+                v3 = -np.cross(self.axis1,self.axis2)
+                
                 v = -v1*spnav_event.motion.x + v2*spnav_event.motion.y \
                     + v3*spnav_event.motion.z
                 v *= self.scale / 5000.0 * accelerate_factor
@@ -502,12 +601,10 @@ class Camera(Thread):
                 self.spnav = False
 
         self.update()
-        
+
         self.done = False
         self.clicked = False
 
-        #current_layer = 0
-
         while not self.done:
             self.clock.tick(20)
 
@@ -525,53 +622,44 @@ class Camera(Thread):
         if self.spnav:
             self.spnav_module.spnav_close()
 
-        del self.gpu
+        del self.gpu_instance
 
 class EventViewer(Camera):
     def __init__(self, geometry, filename, **kwargs):
         Camera.__init__(self, geometry, **kwargs)
-
-        import ROOT
-
-        self.f = ROOT.TFile(filename)
-        self.T = self.f.Get('T')
-        self.T.GetEntry(0)
-        self.nsolids = geometry.solid_id.max() + 1
+        self.rr = RootReader(filename)
 
     def color_hit_pmts(self):
-        self.gpu.reset_colors()
-
-        hit = np.empty(self.nsolids, np.int32)
-        t = np.empty(self.nsolids, np.float32)
-        q = np.empty(self.nsolids, np.float32)
-
-        # self.nsolids has a weird data type that PyROOT doesn't understand
-        self.T.ev.get_channels(int(self.nsolids), hit, t, q)
+        self.gpu_geometry.reset_colors()
 
-        # PyROOT prints warnings when we try to pass a bool array directly
-        # so we convert afterward
-        hit = hit.astype(np.bool)
+        hit = self.ev.channels.hit
+        t = self.ev.channels.t
+        q = self.ev.channels.q
 
         # Important: Compute range only with HIT channels
         solid_colors = map_to_color(q, range=(q[hit].min(),q[hit].max()))
-        self.gpu.color_solids(hit, solid_colors)
+        self.gpu_geometry.color_solids(hit, solid_colors)
         self.update()
 
     def process_event(self, event):
         if event.type == KEYDOWN:
             if event.key == K_PAGEUP:
-                entry = self.T.GetReadEntry()
-                if entry < self.T.GetEntries() - 1:
-                    self.T.GetEntry(entry+1)
+                try:
+                    self.ev = self.rr.next()
+                except StopIteration:
+                    pass
+                else:
                     self.color_hit_pmts()
-                    return
+                return
 
             elif event.key == K_PAGEDOWN:
-                entry = self.T.GetReadEntry()
-                if entry > 0:
-                    self.T.GetEntry(entry-1)
+                try:
+                    self.ev = self.rr.prev()
+                except StopIteration:
+                    pass
+                else:
                     self.color_hit_pmts()
-                    return
+                return
 
         Camera.process_event(self, event)
 
@@ -589,6 +677,10 @@ if __name__ == '__main__':
                       help='bits for z-ordering space axes', default=10)
     parser.add_option('-r', '--resolution', dest='resolution',
                       help='specify resolution', default='1024,576')
+    parser.add_option('--3d', action='store_true', dest='enable3d',
+                      help='enable 3d', default=False)
+    parser.add_option('--green', action='store_true', dest='green_magenta',
+                      help='3d with green and magenta lenses', default=False)
     parser.add_option('-i', dest='io_file', default=None)
     options, args = parser.parse_args()
 
@@ -619,9 +711,9 @@ if __name__ == '__main__':
 
     geometry = build(obj, options.bits)
     if options.io_file is not None:
-        camera = EventViewer(geometry, options.io_file, size=size)
+        camera = EventViewer(geometry, options.io_file, size=size, enable3d=options.enable3d, green_magenta=options.green_magenta)
     else:
-        camera = Camera(geometry, size)
+        camera = Camera(geometry, size, enable3d=options.enable3d, green_magenta=options.green_magenta)
 
     camera.start()
     camera.join()