geometry on the GPU is now a struct created in the GPUGeometry class. coding style for cuda code is now compliant with python PEP 7 -- Style Guide for C Code.

1 files changed, 84 insertions, 36 deletions

diff --git a/camera.py b/camera.py
index ea76f28..1197e8c 100755
--- a/camera.py
+++ b/camera.py
@@ -102,17 +102,18 @@ def encode_movie(dir):
     print 'movie saved to %s.' % filename
 
 class Camera(Thread):
-    def __init__(self, geometry, size=(800,600), device_id=None, enable3d=False, green_magenta=False):
+    def __init__(self, geometry, size=(800,600), device_id=None, enable3d=False, green_magenta=False, alpha_depth=10):
         Thread.__init__(self)
         self.geometry = geometry
         self.device_id = device_id
         self.size = size
         self.enable3d = enable3d
         self.green_magenta = green_magenta
+        self.alpha_depth = alpha_depth
 
         self.unique_bvh_layers = np.unique(self.geometry.layers)
         self.currentlayer = None
-        self.bvh_layers = {self.currentlayer : self.geometry}
+        self.bvh_layers = {}
 
         try: 
             import spnav as spnav_module
@@ -123,15 +124,19 @@ class Camera(Thread):
 
     def init_gpu(self):
         self.gpu_instance = gpu.GPU(self.device_id)
-        self.gpu_geometry = gpu.GPUGeometry(self.gpu_instance, self.geometry)
-        self.gpu_funcs = gpu.GPUFuncs(self.gpu_instance.module)
+        self.gpu_geometry = gpu.GPUGeometry(self.geometry)
+        self.gpu_funcs = gpu.GPUFuncs(gpu.get_cu_module('mesh.h'))
+        self.hybrid_funcs = gpu.GPUFuncs(gpu.get_cu_module('hybrid_render.cu'))
+
+        self.gpu_geometries = [self.gpu_geometry]
 
         self.width, self.height = self.size
 
         self.npixels = self.width*self.height
 
         pygame.init()
-        self.screen = pygame.display.set_mode(self.size)
+        self.window = pygame.display.set_mode(self.size)
+        self.screen = pygame.Surface(self.size, pygame.SRCALPHA)
         pygame.display.set_caption('')
         self.clock = pygame.time.Clock()
 
@@ -182,9 +187,12 @@ class Camera(Thread):
 
             self.rays = gpu.GPURays(pos, dir)
 
+            self.distance_array = ga.empty(self.alpha_depth*self.rays.pos.size, dtype=np.float32)
+            self.index_array = ga.empty(self.alpha_depth*self.rays.pos.size, dtype=np.uint32)
+            self.n_array = ga.zeros(self.rays.pos.size, dtype=np.uint32)
+
             self.pixels_gpu = ga.empty(self.npixels, dtype=np.int32)
 
-        self.alpha = True
         self.movie = False
         self.movie_index = 0
         self.movie_dir = None
@@ -218,7 +226,7 @@ class Camera(Thread):
         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.hybrid_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), self.gpu_geometry.gpudata, block=(self.nblocks,1,1), grid=(self.npixels//self.nblocks+1,1))
 
         self.nlookup_calls += 1
 
@@ -234,7 +242,7 @@ class Camera(Thread):
     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.hybrid_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), self.gpu_geometry.gpudata, block=(self.nblocks,1,1), grid=(rays.pos.size//self.nblocks+1,1))
 
     def update_image(self):
         if self.enable3d:
@@ -247,10 +255,10 @@ class Camera(Thread):
 
     def process_image(self):
         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))
+            self.hybrid_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.hybrid_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))
+            self.hybrid_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'
@@ -320,29 +328,28 @@ class Camera(Thread):
 
             self.update()
 
-    def update_pixels(self):
+    def update_pixels(self, gpu_geometry=None, keep_last_render=False):
+        if gpu_geometry is None:
+            gpu_geometry = self.gpu_geometry
+
         if self.render:
             while self.nlookup_calls < 10:
                 self.update_xyz_lookup(self.source_position)
             self.update_image()
             self.process_image()
         else:
-            if self.alpha:
-                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))
+            if self.enable3d:
+                self.rays1.render(gpu_geometry, self.pixels1_gpu,
+                                  self.alpha_depth, keep_last_render)
+                self.rays2.render(gpu_geometry, self.pixels2_gpu,
+                                  self.alpha_depth, keep_last_render)
             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))
+                self.rays.render(gpu_geometry, self.pixels_gpu,
+                                 self.alpha_depth, keep_last_render)
 
-    def update(self):
+    def update_viewing_angle(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))
+            self.gpu_funcs.distance_to_mesh(np.int32(self.scope_rays.pos.size), self.scope_rays.pos, self.scope_rays.dir, self.gpu_geometry.gpudata, 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()
 
@@ -377,7 +384,16 @@ class Camera(Thread):
 
             self.viewing_angle = new_viewing_angle
 
-        self.update_pixels()
+    def update(self):
+        if self.enable3d:
+            self.update_viewing_angle()
+
+        n = len(self.gpu_geometries)
+        for i, gpu_geometry in enumerate(self.gpu_geometries):
+            if i == 0:
+                self.update_pixels(gpu_geometry)
+            else:
+                self.update_pixels(gpu_geometry, keep_last_render=True)
 
         if self.enable3d:
             pixels1 = self.pixels1_gpu.get()
@@ -387,12 +403,18 @@ class Camera(Thread):
                 pixels = (pixels1 & 0x00ff00) | (pixels2 & 0xff00ff)
             else:
                 pixels = (pixels1 & 0xff0000) | (pixels2 & 0x00ffff)
+
+            alpha = ((0xff & (pixels1 >> 24)) + (0xff & (pixels2 >> 24)))/2
+
+            pixels |= (alpha << 24)
         else:
             pixels = self.pixels_gpu.get()
 
         pygame.surfarray.blit_array(self.screen, pixels.reshape(self.size))
         if self.doom_mode:
             self.screen.blit(self.doom_hud, self.doom_rect)
+        self.window.fill(0)
+        self.window.blit(self.screen, (0,0))
         pygame.display.flip()
 
         if self.movie:
@@ -400,12 +422,17 @@ class Camera(Thread):
             self.movie_index += 1
 
     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] = gpu.GPUGeometry(self.gpu, layergeometry)
+        if layer is None:
+            self.gpu_geometries = [self.gpu_geometry]
+        else:
+            try:
+                gpu_geometry = self.bvh_layers[layer]
+            except KeyError:
+                geometry = build(bvh_mesh(self.geometry, layer), 8)
+                gpu_geometry = gpu.GPUGeometry(geometry)
+                self.bvh_layers[layer] = gpu_geometry
+
+            self.gpu_geometries = [self.gpu_geometry, gpu_geometry]
 
         self.update()
 
@@ -484,6 +511,15 @@ class Camera(Thread):
                 self.done = True
                 return
 
+            elif event.key == K_EQUALS:
+                self.alpha_depth += 1
+                self.update()
+
+            elif event.key == K_MINUS:
+                if self.alpha_depth > 0:
+                    self.alpha_depth -= 1
+                    self.update()
+
             elif event.key == K_PAGEDOWN:
                 if self.currentlayer is None:
                     self.currentlayer = self.unique_bvh_layers[-1]
@@ -517,10 +553,6 @@ class Camera(Thread):
                 self.clear_image()
                 self.update()
 
-            elif event.key == K_F8:
-                self.alpha = not self.alpha
-                self.update()
-
             elif event.key == K_m:
                 if self.movie:
                     encode_movie(self.movie_dir)
@@ -630,6 +662,20 @@ class EventViewer(Camera):
         Camera.__init__(self, geometry, **kwargs)
         self.rr = RootReader(filename)
 
+    def render_particle_track(self):
+        marker = Solid(make.cube(0.1), vacuum, vacuum)
+
+        geometry = Geometry()
+        for pos in self.ev.photons_beg.pos[::100]:
+            geometry.add_solid(marker, displacement=pos)
+
+        geometry.build(bits=11)
+        gpu_geometry = gpu.GPUGeometry(geometry)
+
+        self.gpu_geometries = [self.gpu_geometry, gpu_geometry]
+
+        self.update()
+
     def color_hit_pmts(self):
         self.gpu_geometry.reset_colors()
 
@@ -651,6 +697,7 @@ class EventViewer(Camera):
                     pass
                 else:
                     self.color_hit_pmts()
+                    self.render_particle_track()
                 return
 
             elif event.key == K_PAGEDOWN:
@@ -660,6 +707,7 @@ class EventViewer(Camera):
                     pass
                 else:
                     self.color_hit_pmts()
+                    self.render_particle_track()
                 return
 
         Camera.process_event(self, event)