add a GPU class to handle both the gpu context and module; since the geometry requires global device pointers, there should be a one to one correspondence between modules and contexts. the current plan is to perform all gpu operations within this class. also add a simple color map to display hit pmt charge and timing information.

4 files changed, 182 insertions, 13 deletions

diff --git a/camera.py b/camera.py
index 37dc072..12f3e6b 100644..100755
--- a/camera.py
+++ b/camera.py
@@ -1,3 +1,4 @@
+#!/usr/bin/env python
 import numpy as np
 from itertools import product, count
 from threading import Thread, Lock
@@ -71,12 +72,16 @@ def get_rays(position, size = (800, 600), film_size = (0.035, 0.024), focal_leng
     return grid, focal_point-grid
 
 class Camera(Thread):
-    def __init__(self, geometry, module, context, lock, size=(800,600)):
+    def __init__(self, geometry, module, context, lock=None, size=(800,600)):
         Thread.__init__(self)
         self.geometry = geometry
         self.module = module
         self.context = context
-        self.lock = lock
+
+        if lock is None:
+            self.lock = Lock()
+        else:
+            self.lock = lock
 
         self.size = size
         self.width, self.height = size
@@ -247,7 +252,7 @@ class Camera(Thread):
 
     def update(self):
         if self.render:
-            while self.nlookup_calls < 10:
+            while self.nlookup_calls < 100:
                 self.update_xyz_lookup(self.source_position)
             self.update_image()
             self.process_image()
@@ -355,7 +360,7 @@ class Camera(Thread):
                         self.source_position = self.point
 
                     if event.key == K_p:
-                        for i in range(10):
+                        for i in range(100):
                             self.update_xyz_lookup(self.point)
                         self.source_position = self.point
 
@@ -479,18 +484,14 @@ if __name__ == '__main__':
             raise Exception("can't find object %s" % args[0])
 
     from pycuda.tools import make_default_context
+    from gpu import *
 
-    cuda.init()
-    context = make_default_context()
-    print 'device %s' % context.get_device().name()
-
-    module = SourceModule(src.kernel, options=['-I' + src.dir], no_extern_c=True, cache_dir=False)
     geometry = build(obj, options.bits)
-    geometry.load(module)
 
     lock = Lock()
 
-    camera = Camera(geometry, module, context, lock, size)
-
-    context.pop()
+    gpu = GPU()
+    gpu.load_geometry(geometry)
+    gpu.context.pop()
+    camera = Camera(geometry, gpu.module, gpu.context, lock, size=size)
     camera.start()
diff --git a/color/__init__.py b/color/__init__.py
index c5fda6a..a8e7550 100644
--- a/color/__init__.py
+++ b/color/__init__.py
@@ -1 +1,2 @@
 from chromaticity import map_wavelength
+from colormap import map_to_color
diff --git a/color/colormap.py b/color/colormap.py
new file mode 100644
index 0000000..b3761ae
--- /dev/null
+++ b/color/colormap.py
@@ -0,0 +1,27 @@
+import numpy as np
+
+import matplotlib.pyplot as plt
+
+def map_to_color(a, range=None):
+    a = np.asarray(a)
+    if range is None:
+        range = (a.min(), a.max())
+
+    x = np.linspace(0, np.pi, 100)
+
+    yr = np.cos(x)**2
+    yr[x > np.pi/2] = 0.0
+    yg = np.sin(x)**2
+    yb = np.cos(x)**2
+    yb[x < np.pi/2] = 0.0
+
+    #plt.plot(x, yr, 'r-', x, yb, 'b-', x, yg, 'g-')
+    #plt.show()
+
+    ax = np.pi*(a - range[0])/(range[1]-range[0])
+
+    r = (np.interp(ax, x, yr)*255).astype(np.uint32)
+    g = (np.interp(ax, x, yg)*255).astype(np.uint32)
+    b = (np.interp(ax, x, yb)*255).astype(np.uint32)
+
+    return r << 16 | g << 8 | b
diff --git a/gpu.py b/gpu.py
new file mode 100644
index 0000000..e5fd5f1
--- /dev/null
+++ b/gpu.py
@@ -0,0 +1,140 @@
+import numpy as np
+import numpy.ma as ma
+from pycuda.tools import make_default_context
+from pycuda.compiler import SourceModule
+import pycuda.driver as cuda
+from pycuda import gpuarray
+from copy import copy
+from itertools import izip
+
+import src
+from geometry import standard_wavelengths
+from color import map_to_color
+
+cuda.init()
+
+def format_size(size):
+    if size < 1e3:
+        return '%.1f%s' % (size, ' ')
+    elif size < 1e6:
+        return '%.1f%s' % (size/1e3, 'K')
+    elif size < 1e9:
+        return '%.1f%s' % (size/1e6, 'M')
+    else:
+        return '%.1f%s' % (size/1e9, 'G')
+
+def format_array(name, array):
+    return '%-15s %6s %6s' % \
+        (name, format_size(len(array)), format_size(array.nbytes))
+
+class GPU(object):
+    def __init__(self):
+        self.context = make_default_context()
+        print 'device %s' % self.context.get_device().name()
+        self.module = SourceModule(src.kernel, options=['-I' + src.dir], no_extern_c=True, cache_dir=False)
+
+    def print_device_usage(self):
+        print 'device usage:'
+        print format_array('vertices', self.vertices_gpu)
+        print format_array('triangles', self.triangles_gpu)
+        print format_array('lower_bounds', self.lower_bounds_gpu)
+        print format_array('upper_bounds', self.upper_bounds_gpu)
+        print format_array('node_map', self.node_map_gpu)
+        print format_array('node_length', self.node_length_gpu)
+        print '%-15s %6s %6s' % ('total', '', format_size(self.vertices_gpu.nbytes + self.triangles_gpu.nbytes + self.lower_bounds_gpu.nbytes + self.upper_bounds_gpu.nbytes + self.node_map_gpu.nbytes + self.node_length_gpu.nbytes))
+
+    def load_geometry(self, geometry):
+        if not hasattr(geometry, 'mesh'):
+            geometry.build(bits=8)
+
+        set_wavelength_range = self.module.get_function('set_wavelength_range')
+        set_wavelength_range(np.float32(standard_wavelengths[0]), np.float32(standard_wavelengths[-1]), np.float32(standard_wavelengths[1]-standard_wavelengths[0]), np.uint32(standard_wavelengths.size), block=(1,1,1), grid=(1,1))
+
+        set_material = self.module.get_function('set_material')
+        self.materials = []
+        for i in range(len(geometry.unique_materials)):
+            material = copy(geometry.unique_materials[i])
+
+            if material is None:
+                raise Exception('one or more triangles is missing a material.')
+
+            material.refractive_index_gpu = gpuarray.to_gpu(np.interp(standard_wavelengths, material.refractive_index[:,0], material.refractive_index[:,1]).astype(np.float32))
+            material.absorption_length_gpu = gpuarray.to_gpu(np.interp(standard_wavelengths, material.absorption_length[:,0], material.absorption_length[:,1]).astype(np.float32))
+            material.scattering_length_gpu = gpuarray.to_gpu(np.interp(standard_wavelengths, material.scattering_length[:,0], material.scattering_length[:,1]).astype(np.float32))
+
+            set_material(np.int32(i), material.refractive_index_gpu, material.absorption_length_gpu, material.scattering_length_gpu, block=(1,1,1), grid=(1,1))
+
+            self.materials.append(material)
+
+        set_surface = self.module.get_function('set_surface')
+        self.surfaces = []
+        for i in range(len(geometry.unique_surfaces)):
+            surface = copy(geometry.unique_surfaces[i])
+
+            if surface is None:
+                continue
+
+            surface.detect_gpu = gpuarray.to_gpu(np.interp(standard_wavelengths, surface.detect[:,0], surface.detect[:,1]).astype(np.float32))
+            surface.absorb_gpu = gpuarray.to_gpu(np.interp(standard_wavelengths, surface.absorb[:,0], surface.absorb[:,1]).astype(np.float32))
+            surface.reflect_diffuse_gpu = gpuarray.to_gpu(np.interp(standard_wavelengths, surface.reflect_diffuse[:,0], surface.reflect_diffuse[:,1]).astype(np.float32))
+            surface.reflect_specular_gpu = gpuarray.to_gpu(np.interp(standard_wavelengths, surface.reflect_specular[:,0], surface.reflect_specular[:,1]).astype(np.float32))
+
+            set_surface(np.int32(i), surface.detect_gpu, surface.absorb_gpu, surface.reflect_diffuse_gpu, surface.reflect_specular_gpu, block=(1,1,1), grid=(1,1))
+
+            self.surfaces.append(surface)
+
+        self.vertices_gpu = gpuarray.to_gpu(geometry.mesh.vertices.astype(np.float32).view(gpuarray.vec.float3))
+
+        triangles = \
+            np.empty(len(geometry.mesh.triangles), dtype=gpuarray.vec.uint4)
+        triangles['x'] = geometry.mesh.triangles[:,0]
+        triangles['y'] = geometry.mesh.triangles[:,1]
+        triangles['z'] = geometry.mesh.triangles[:,2]
+        triangles['w'] = ((geometry.material1_index & 0xff) << 24) | ((geometry.material2_index & 0xff) << 16) | ((geometry.surface_index & 0xff) << 8)
+        self.triangles_gpu = gpuarray.to_gpu(triangles)
+
+        lower_bounds_float4 = np.empty(geometry.lower_bounds.shape[0], dtype=gpuarray.vec.float4)
+        lower_bounds_float4['x'] = geometry.lower_bounds[:,0]
+        lower_bounds_float4['y'] = geometry.lower_bounds[:,1]
+        lower_bounds_float4['z'] = geometry.lower_bounds[:,2]
+        self.lower_bounds_gpu = gpuarray.to_gpu(lower_bounds_float4)
+
+        upper_bounds_float4 = np.empty(geometry.upper_bounds.shape[0], dtype=gpuarray.vec.float4)
+        upper_bounds_float4['x'] = geometry.upper_bounds[:,0]
+        upper_bounds_float4['y'] = geometry.upper_bounds[:,1]
+        upper_bounds_float4['z'] = geometry.upper_bounds[:,2]
+        self.upper_bounds_gpu = gpuarray.to_gpu(upper_bounds_float4)
+
+        self.colors_gpu = gpuarray.to_gpu(geometry.colors.astype(np.uint32))
+        self.node_map_gpu = gpuarray.to_gpu(geometry.node_map.astype(np.uint32))
+        self.node_length_gpu = gpuarray.to_gpu(geometry.node_length.astype(np.uint32))
+        self.solid_id_map_gpu = gpuarray.to_gpu(geometry.solid_id.astype(np.uint32))
+
+        set_global_mesh_variables = self.module.get_function('set_global_mesh_variables')
+        set_global_mesh_variables(self.triangles_gpu, self.vertices_gpu, self.colors_gpu, np.uint32(geometry.node_map.size-1), np.uint32(geometry.first_node), block=(1,1,1), grid=(1,1))
+
+        self.lower_bounds_tex = self.module.get_texref('lower_bounds')
+        self.upper_bounds_tex = self.module.get_texref('upper_bounds')
+        self.node_map_tex = self.module.get_texref('node_map')
+        self.node_length_tex = self.module.get_texref('node_length')
+
+        self.lower_bounds_tex.set_address(self.lower_bounds_gpu.gpudata, self.lower_bounds_gpu.nbytes)
+        self.upper_bounds_tex.set_address(self.upper_bounds_gpu.gpudata, self.upper_bounds_gpu.nbytes)
+        self.node_map_tex.set_address(self.node_map_gpu.gpudata, self.node_map_gpu.nbytes)
+        self.node_length_tex.set_address(self.node_length_gpu.gpudata, self.node_length_gpu.nbytes)
+
+        self.lower_bounds_tex.set_format(cuda.array_format.FLOAT, 4)
+        self.upper_bounds_tex.set_format(cuda.array_format.FLOAT, 4)
+        self.node_map_tex.set_format(cuda.array_format.UNSIGNED_INT32, 1)
+        self.node_length_tex.set_format(cuda.array_format.UNSIGNED_INT32, 1)
+
+        self.print_device_usage()
+
+    def color_solids(self, solid_ids, colors):
+        solid_id_map = self.solid_id_map_gpu.get()
+        triangle_colors = self.colors_gpu.get()
+
+        for i, color in izip(solid_ids, colors):
+            triangle_colors[solid_id_map == i] = color
+
+        self.colors_gpu.set(triangle_colors)