Add flag to Simulation() and Geometry.build() to ignore the BVH cache.

2 files changed, 40 insertions, 32 deletions

diff --git a/geometry.py b/geometry.py
index 6ea8994..b2cb9c0 100644
--- a/geometry.py
+++ b/geometry.py
@@ -225,7 +225,7 @@ class Geometry(object):
         return len(self.solids)-1
 
     @timeit
-    def build(self, bits=8, shift=3):
+    def build(self, bits=8, shift=3, use_cache=True):
         """
         Build the bounding volume hierarchy, material/surface code arrays, and
         color array for this geometry. If the bounding volume hierarchy is
@@ -240,6 +240,11 @@ class Geometry(object):
                 The number of bits to shift the zvalue of each node when
                 building the next layer of the bounding volume hierarchy.
                 Defaults to 3.
+            - use_cache: bool, *optional*
+                If true, the on-disk cache in ~/.chroma/ will be checked for
+                a previously built version of this geometry, otherwise the
+                BVH will be computed and saved to the cache.  If false,
+                the cache is ignored and also not updated.
         """
         nv = np.cumsum([0] + [len(solid.mesh.vertices) for solid in self.solids])
         nt = np.cumsum([0] + [len(solid.mesh.triangles) for solid in self.solids])
@@ -286,26 +291,27 @@ class Geometry(object):
         cache_file = checksum.hexdigest()
         cache_path = os.path.join(cache_dir, cache_file)
 
-        try:
-            f = gzip.GzipFile(cache_path, 'rb')
-        except IOError:
-            pass
-        else:
-            print 'loading cache.'
-            data = pickle.load(f)
-
-            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]
-
-            for key, value in data.iteritems():
-                setattr(self, key, value)
-            f.close()
-            return
+        if use_cache:
+            try:
+                f = gzip.GzipFile(cache_path, 'rb')
+            except IOError:
+                pass
+            else:
+                print 'loading cache.'
+                data = pickle.load(f)
+
+                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]
+
+                for key, value in data.iteritems():
+                    setattr(self, key, value)
+                f.close()
+                return
 
         zvalues_mesh = morton_order(self.mesh, bits)
         reorder = np.argsort(zvalues_mesh)
@@ -366,14 +372,15 @@ class Geometry(object):
             if unique_zvalues.size == 1:
                 break
 
-        print >>sys.stderr, 'Writing BVH to cache directory...'
+        if use_cache:
+            print >>sys.stderr, 'Writing BVH to cache directory...'
 
-        if not os.path.exists(cache_dir):
-            os.makedirs(cache_dir)
+            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']:
-                data[key] = getattr(self, key)
-            data['reorder'] = reorder
-            pickle.dump(data, f, -1)
+            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']:
+                    data[key] = getattr(self, key)
+                data['reorder'] = reorder
+                pickle.dump(data, f, -1)
diff --git a/sim.py b/sim.py
index 2d772bb..fce559e 100755
--- a/sim.py
+++ b/sim.py
@@ -23,7 +23,8 @@ def pick_seed():
 
 class Simulation(object):
     def __init__(self, detector, detector_material, 
-                 seed=None, cuda_device=None, geant4_processes=4, bvh_bits=11):
+                 seed=None, cuda_device=None, geant4_processes=4, bvh_bits=11,
+                 use_cache=True):
         self.detector = detector
 
         if seed is None:
@@ -44,7 +45,7 @@ class Simulation(object):
             self.photon_generator = None
 
         print >>sys.stderr, 'Creating BVH with %d bits...' % (bvh_bits)
-        detector.build(bits=bvh_bits)
+        detector.build(bits=bvh_bits, use_cache=use_cache)
         
         print >>sys.stderr, 'Initializing GPU...'
         self.gpu_worker = gpu.GPU(cuda_device)