speedup Geometry.build()

1 files changed, 12 insertions, 18 deletions

diff --git a/geometry.py b/geometry.py
index 37367e5..a98b5cd 100644
--- a/geometry.py
+++ b/geometry.py
@@ -212,6 +212,7 @@ class Geometry(object):
         return len(self.solids)-1
 
     @timeit
+    #@profile
     def build(self, bits=8, shift=3):
         offsets = [ (0,0) ]
         for solid in self.solids:
@@ -269,34 +270,27 @@ class Geometry(object):
                                 enumerate(self.solids)])[reorder]
 
         unique_zvalues = np.unique(zvalues_mesh)
-        zvalues = np.empty(unique_zvalues.size, dtype=np.uint64)
 
         self.lower_bounds = np.empty((unique_zvalues.size,3), dtype=np.float32)
         self.upper_bounds = np.empty((unique_zvalues.size,3), dtype=np.float32)
-        self.node_map = np.empty(unique_zvalues.size, dtype=np.uint32)
-        self.node_length = np.empty(unique_zvalues.size, dtype=np.uint32)
-        self.layers = np.empty(unique_zvalues.size, dtype=np.uint32)
 
-        layer = 0
+        assembled_mesh = self.mesh.assemble(group=False)
+        zi1_array = np.searchsorted(zvalues_mesh, unique_zvalues)
+        zi2_array = np.searchsorted(zvalues_mesh, unique_zvalues, side='right')
 
-        for i, z in enumerate(unique_zvalues):
-            i1 = np.searchsorted(zvalues_mesh, z)
-            i2 = np.searchsorted(zvalues_mesh, z, side='right')
-
-            self.lower_bounds[i] = self.mesh.assemble(slice(i1,i2), group=False).min(axis=0)
-            self.upper_bounds[i] = self.mesh.assemble(slice(i1,i2), group=False).max(axis=0)
-
-            self.node_map[i] = i1
-            self.node_length[i] = i2-i1
-            self.layers[i] = layer
-
-            zvalues[i] = z
+        for i, (zi1, zi2) in enumerate(zip(zi1_array,zi2_array)):
+            self.lower_bounds[i] = assembled_mesh[zi1*3:zi2*3].min(0)
+            self.upper_bounds[i] = assembled_mesh[zi1*3:zi2*3].max(0)
 
+        zvalues = np.array(unique_zvalues, copy=True)
+        self.node_map = np.array(zi1_array, copy=True)
+        self.node_length = zi2_array - zi1_array
+        self.layers = np.zeros(unique_zvalues.size, dtype=np.uint32)
         self.first_node = unique_zvalues.size
 
         begin_last_layer = 0
 
-        layer += 1
+        layer = 1
 
         while True:
             bit_shifted_zvalues = zvalues >> shift