add the ability to cache a geometry along with its bounding volume hierarchy. cached files are stored in $HOME/.chroma. fixed the timeit() decorator so that the decorated function is still able to pass back a return value.

5 files changed, 88 insertions, 21 deletions

diff --git a/camera.py b/camera.py
index d9b70c7..52a0add 100755
--- a/camera.py
+++ b/camera.py
@@ -5,6 +5,7 @@ from threading import Thread, Lock
 import os
 import sys
 from color import map_wavelength, map_to_color
+import matplotlib.pyplot as plt
 
 from gpu import GPU, CUDAFuncs
 from tools import timeit
@@ -431,7 +432,7 @@ class EventViewer(Camera):
             t[i] = channel.t
             q[i] = channel.q
 
-        self.gpu.color_solids(solid_ids, map_to_color(t, (t.min(), t.mean())))
+        self.gpu.color_solids(solid_ids, map_to_color(t, range=(t.min(),t.mean())))
         self.update()
 
     def process_event(self, event):
diff --git a/color/colormap.py b/color/colormap.py
index e5bd73e..6f3e056 100644
--- a/color/colormap.py
+++ b/color/colormap.py
@@ -1,13 +1,20 @@
 import numpy as np
 import matplotlib.cm as cm
 
-def map_to_color(a, range=None, map=cm.hsv):
+def map_to_color(a, range=None, map=cm.hsv, weights=None):
     a = np.asarray(a)
     if range is None:
         range = (a.min(), a.max())
 
     ax = (a - range[0])/(range[1]-range[0])
 
-    rgba = (map(ax)*255).astype(np.uint32)
+    frgba = map(ax)
+
+    if weights is not None:
+        frgba[:,0] *= weights
+        frgba[:,1] *= weights
+        frgba[:,2] *= weights
+
+    rgba = (frgba*255).astype(np.uint32)
 
     return rgba[:,0] << 16 | rgba[:,1] << 8 | rgba[:,2]
diff --git a/geometry.py b/geometry.py
index 0657a77..579e01d 100644
--- a/geometry.py
+++ b/geometry.py
@@ -1,7 +1,10 @@
 import sys
+import os
 import numpy as np
 from itertoolset import *
 from tools import timeit
+from hashlib import md5
+import cPickle as pickle
 
 # all material/surface properties are interpolated at these
 # wavelengths when they are sent to the gpu
@@ -227,43 +230,88 @@ class Geometry(object):
             
         self.mesh = Mesh(vertices, triangles)
 
-        zvalues_mesh = morton_order(self.mesh, bits)
-        reorder = np.argsort(zvalues_mesh)
-        zvalues_mesh = zvalues_mesh[reorder]
-
-        if (np.diff(zvalues_mesh) < 0).any():
-            raise Exception('zvalues_mesh out of order.')
+        self.colors = np.concatenate([solid.color for solid in self.solids])
 
-        self.mesh.triangles = self.mesh.triangles[reorder]
+        self.solid_id = np.concatenate([np.fromiter(repeat(i, len(solid.mesh.triangles)), dtype=np.uint32) for i, solid in enumerate(self.solids)])
 
         self.unique_materials = list(np.unique(np.concatenate([solid.unique_materials for solid in self.solids])))
 
         material_lookup = dict(zip(self.unique_materials, range(len(self.unique_materials))))
 
         self.material1_index = \
-            np.fromiter(imap(material_lookup.get, chain(*[solid.material1 for solid in self.solids])), dtype=np.int32)[reorder]
+            np.fromiter(imap(material_lookup.get, chain(*[solid.material1 for solid in self.solids])), dtype=np.int32)
 
         self.material2_index = \
-            np.fromiter(imap(material_lookup.get, chain(*[solid.material2 for solid in self.solids])), dtype=np.int32)[reorder]
+            np.fromiter(imap(material_lookup.get, chain(*[solid.material2 for solid in self.solids])), dtype=np.int32)
 
         self.unique_surfaces = list(np.unique(np.concatenate([solid.unique_surfaces for solid in self.solids])))
 
         surface_lookup = dict(zip(self.unique_surfaces, range(len(self.unique_surfaces))))
 
         self.surface_index = \
-            np.fromiter(imap(surface_lookup.get, chain(*[solid.surface for solid in self.solids])), dtype=np.int32)[reorder]
+            np.fromiter(imap(surface_lookup.get, chain(*[solid.surface for solid in self.solids])), dtype=np.int32)
 
         try:
             self.surface_index[self.surface_index == self.unique_surfaces.index(None)] = -1
         except ValueError:
             pass
 
-        self.colors = \
-            np.concatenate([solid.color for solid in self.solids])[reorder]
+        checksum = md5(str(bits))
+
+        checksum.update(self.mesh.vertices)
+        checksum.update(self.mesh.triangles)
+        checksum.update(self.colors)
+        checksum.update(self.solid_id)
+        checksum.update(self.material1_index)
+        checksum.update(self.material2_index)
+        checksum.update(self.surface_index)
+
+        for material in self.unique_materials:
+            checksum.update(material.refractive_index)
+            checksum.update(material.absorption_length)
+            checksum.update(material.scattering_length)
+
+        for surface in self.unique_surfaces:
+            if surface is not None:
+                checksum.update(surface.detect)
+                checksum.update(surface.absorb)
+                checksum.update(surface.reflect_diffuse)
+                checksum.update(surface.reflect_specular)
+            else:
+                checksum.update(str(None))
+
+        cache_dir = os.path.expanduser('~/.chroma')
+        cache_file = checksum.hexdigest()
+        cache_path = os.path.join(cache_dir, cache_file)
 
-        self.solid_id = \
-            np.concatenate([np.tile(i, len(solid.mesh)) for i, solid in \
-                                enumerate(self.solids)])[reorder]
+        try:
+            f = open(cache_path, 'rb')
+        except IOError:
+            pass
+        else:
+            print 'loading cache.'
+            data = pickle.load(f)
+            self.mesh.triangles = self.mesh.triangles[data['reorder']]
+            del data['reorder']
+            for key, value in data.items():
+                setattr(self, key, value)
+            f.close()
+            return
+
+        zvalues_mesh = morton_order(self.mesh, bits)
+        reorder = np.argsort(zvalues_mesh)
+        zvalues_mesh = zvalues_mesh[reorder]
+
+        if (np.diff(zvalues_mesh) < 0).any():
+            raise Exception('zvalues_mesh out of order.')
+
+        self.mesh.triangles = self.mesh.triangles[reorder]
+
+        self.material1_index = self.material1_index[reorder]
+        self.material2_index = self.material2_index[reorder]
+        self.surface_index = self.surface_index[reorder]
+        self.colors = self.colors[reorder]
+        self.solid_id = self.solid_id[reorder]
 
         unique_zvalues = np.unique(zvalues_mesh)
 
@@ -308,3 +356,14 @@ class Geometry(object):
 
             if unique_zvalues.size == 1:
                 break
+
+        if not os.path.exists(cache_dir):
+            os.makedirs(cache_dir)
+
+        f = open(cache_path, 'wb')
+        data = {}
+        for key in ['material1_index', 'material2_index', 'surface_index', 'colors', 'solid_id', 'lower_bounds', 'upper_bounds', 'node_map', 'node_map_end', 'layers', 'first_node']:
+            data[key] = getattr(self, key)
+        data['reorder'] = reorder
+        pickle.dump(data, f, -1)
+        f.close()
diff --git a/src/mesh.h b/src/mesh.h
index 7c148af..b4714c4 100644
--- a/src/mesh.h
+++ b/src/mesh.h
@@ -16,10 +16,9 @@ __device__ unsigned int g_first_node;
 texture<float4, 1, cudaReadModeElementType> upper_bounds;
 texture<float4, 1, cudaReadModeElementType> lower_bounds;
 
-/* map to child nodes/triangles and the number of child nodes/triangles */
+/* map to child node/triangle indices */
 texture<unsigned int, 1, cudaReadModeElementType> node_map;
 texture<unsigned int, 1, cudaReadModeElementType> node_map_end;
-//texture<unsigned int, 1, cudaReadModeElementType> node_length;
 
 __device__ float3 make_float3(const float4 &a)
 {
diff --git a/tools.py b/tools.py
index 2a44104..6835228 100644
--- a/tools.py
+++ b/tools.py
@@ -5,9 +5,10 @@ import datetime
 def timeit(func):
     def f(*args, **kwargs):
         t0 = time.time()
-        func(*args, **kwargs)
+        retval = func(*args, **kwargs)
         elapsed = time.time() - t0
         print '%s elapsed in %s().' % (datetime.timedelta(seconds=elapsed), func.__name__)
+        return retval
     return f
 
 def read_csv(filename):