Speed up Geometry.build() by a large factor when loading from cache.

A bunch of small tricks have been applied to reduce the amount of time
required to build an already cached geometry:

  * Replace uses of fromiter() on long sequences with code that operates
    on bigger arrays.
  * Use memoization on the Solids to more efficiently map materials
    to material codes when a solid is repeated (as is the case in all
    our detectors)
  * Use numpy.take() instead of fancy indexing on big arrays.  I
    learned about this trick from:
        http://wesmckinney.com/blog/?p=215

Also, switched over to compressed npz files for storing cache information.
They take the same size as the gzipped pickle files, but load 30% faster.

3 files changed, 69 insertions, 22 deletions

diff --git a/chroma/geometry.py b/chroma/geometry.py
index 8c97a74..80b83fc 100644
--- a/chroma/geometry.py
+++ b/chroma/geometry.py
@@ -7,7 +7,7 @@ import numpy as np
 import time
 
 from chroma.itertoolset import *
-from chroma.tools import timeit, profile_if_possible
+from chroma.tools import timeit, profile_if_possible, filled_array
 from chroma.log import logger
 
 # all material/surface properties are interpolated at these
@@ -77,6 +77,31 @@ class Mesh(object):
     def __add__(self, other):
         return Mesh(np.concatenate((self.vertices, other.vertices)), np.concatenate((self.triangles, other.triangles + len(self.vertices))))
 
+def memoize_method_with_dictionary_arg(func):
+    def lookup(*args):
+        # based on function by Michele Simionato
+        # http://www.phyast.pitt.edu/~micheles/python/
+        # Modified to work for class method with dictionary argument
+
+        assert len(args) == 2
+        # create hashable arguments by replacing dictionaries with tuples of items
+        dict_items = args[1].items()
+        dict_items.sort()
+        hashable_args = (args[0], tuple(dict_items))
+        try:
+            return func._memoize_dic[hashable_args]
+        except AttributeError:
+            # _memoize_dic doesn't exist yet.
+
+            result = func(*args)
+            func._memoize_dic = {hashable_args: result}
+            return result
+        except KeyError:
+            result = func(*args)
+            func._memoize_dic[hashable_args] = result
+            return result
+    return lookup
+
 class Solid(object):
     """Solid object attaches materials, surfaces, and colors to each triangle
     in a Mesh object."""
@@ -119,6 +144,18 @@ class Solid(object):
     def __add__(self, other):
         return Solid(self.mesh + other.mesh, np.concatenate((self.material1, other.material1)), np.concatenate((self.material2, other.material2)), np.concatenate((self.surface, other.surface)), np.concatenate((self.color, other.color)))
 
+    @memoize_method_with_dictionary_arg
+    def material1_indices(self, material_lookup):
+        return np.fromiter(imap(material_lookup.get, self.material1), dtype=np.int32, count=len(self.material1))
+
+    @memoize_method_with_dictionary_arg
+    def material2_indices(self, material_lookup):
+        return np.fromiter(imap(material_lookup.get, self.material2), dtype=np.int32, count=len(self.material2))
+
+    @memoize_method_with_dictionary_arg
+    def surface_indices(self, surface_lookup):
+        return np.fromiter(imap(surface_lookup.get, self.surface), dtype=np.int32, count=len(self.surface))
+
 class Material(object):
     """Material optical properties."""
     def __init__(self, name='none'):
@@ -247,6 +284,7 @@ class Geometry(object):
 
         return len(self.solids)-1
 
+    @profile_if_possible
     def build(self, bits=11, shift=3, use_cache=True):
         """
         Build the bounding volume hierarchy, material/surface code arrays, and
@@ -285,25 +323,26 @@ class Geometry(object):
                 np.inner(solid.mesh.vertices, self.solid_rotations[i]) + self.solid_displacements[i]
             triangles[nt[i]:nt[i+1]] = solid.mesh.triangles + nv[i]
 
-        self.mesh = Mesh(vertices, triangles)
+        # Different solids are very unlikely to share vertices, so this goes much faster
+        self.mesh = Mesh(vertices, triangles, remove_duplicate_vertices=False)
 
         self.colors = np.concatenate([solid.color for solid in self.solids])
 
-        self.solid_id = np.concatenate([np.fromiter(repeat(i, len(solid.mesh.triangles)), dtype=np.uint32) for i, solid in enumerate(self.solids)])
+        self.solid_id = np.concatenate([filled_array(i, shape=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)
+        self.material1_index = np.concatenate([solid.material1_indices(material_lookup) for solid in self.solids])
 
-        self.material2_index = np.fromiter(imap(material_lookup.get, chain(*[solid.material2 for solid in self.solids])), dtype=np.int32)
+        self.material2_index = np.concatenate([solid.material2_indices(material_lookup) for solid in self.solids])
 
         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)
+        self.surface_index = np.concatenate([solid.surface_indices(surface_lookup) for solid in self.solids])
 
         try:
             self.surface_index[self.surface_index == surface_lookup[None]] = -1
@@ -316,29 +355,30 @@ class Geometry(object):
         checksum.update(self.mesh.triangles)
 
         cache_dir = os.path.expanduser('~/.chroma')
-        cache_file = checksum.hexdigest()
+        cache_file = checksum.hexdigest()+'.npz'
         cache_path = os.path.join(cache_dir, cache_file)
 
         if use_cache:
             try:
-                f = gzip.GzipFile(cache_path, 'rb')
+                npz_file = np.load(cache_path)
             except IOError:
                 pass
             else:
                 logger.info('Loading BVH from cache.')
-                data = pickle.load(f)
+                data = dict(npz_file)
 
+                # take() is faster than fancy indexing by 5x!
+                # tip from http://wesmckinney.com/blog/?p=215
                 reorder = data.pop('reorder')
-                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]
+                self.mesh.triangles = self.mesh.triangles.take(reorder, axis=0)
+                self.material1_index = self.material1_index.take(reorder, axis=0)
+                self.material2_index = self.material2_index.take(reorder, axis=0)
+                self.surface_index = self.surface_index.take(reorder, axis=0)
+                self.colors = self.colors.take(reorder, axis=0)
+                self.solid_id = self.solid_id.take(reorder, axis=0)
 
                 for key, value in data.iteritems():
                     setattr(self, key, value)
-                f.close()
 
                 logger.info('  nodes: %d' % len(self.upper_bounds))
                 return
@@ -422,9 +462,8 @@ class Geometry(object):
             if not os.path.exists(cache_dir):
                 os.makedirs(cache_dir)
 
-            with gzip.GzipFile(cache_path, 'wb', compresslevel=1) as f:
-                data = {}
-                for key in ['lower_bounds', 'upper_bounds', 'node_map', 'node_map_end', 'layers', 'first_node', 'start_node']:
-                    data[key] = getattr(self, key)
-                data['reorder'] = reorder
-                pickle.dump(data, f, -1)
+            data = {}
+            for key in ['lower_bounds', 'upper_bounds', 'node_map', 'node_map_end', 'layers', 'first_node', 'start_node']:
+                data[key] = getattr(self, key)
+            data['reorder'] = reorder
+            np.savez_compressed(cache_path, **data)
diff --git a/chroma/sim.py b/chroma/sim.py
index f00f536..72fdf3e 100644
--- a/chroma/sim.py
+++ b/chroma/sim.py
@@ -7,6 +7,7 @@ from chroma import generator
 from chroma import gpu
 from chroma import event
 from chroma import itertoolset
+from chroma.tools import profile_if_possible
 
 def pick_seed():
     """Returns a seed for a random number generator selected using
@@ -14,6 +15,7 @@ def pick_seed():
     return int(time.time()) ^ (os.getpid() << 16)
 
 class Simulation(object):
+    @profile_if_possible
     def __init__(self, detector, seed=None, cuda_device=None,
                  geant4_processes=4, nthreads_per_block=64, max_blocks=1024):
         self.detector = detector
diff --git a/chroma/tools.py b/chroma/tools.py
index 1840fae..831a356 100644
--- a/chroma/tools.py
+++ b/chroma/tools.py
@@ -4,6 +4,12 @@ import datetime
 import sys
 import math
 
+def filled_array(value, shape, dtype):
+    '''Create a numpy array of given `shape` and `dtype` filled with the scalar `value`.'''
+    a = np.empty(shape=shape, dtype=dtype)
+    a.fill(value)
+    return a
+
 def ufloat_to_str(x):
     msd = -int(math.floor(math.log10(x.std_dev())))
     return '%.*f +/- %.*f' % (msd, round(x.nominal_value, msd),