merge with tony

1 files changed, 11 insertions, 21 deletions

diff --git a/geometry.py b/geometry.py
index ac71429..d80b730 100644
--- a/geometry.py
+++ b/geometry.py
@@ -49,8 +49,13 @@ class Mesh(object):
 
         self.vertices = np.array(unique_vertices)
 
-    def __getitem__(self, key):
-        return self.vertices[self.triangles[key]]
+    def assemble(self, key=slice(None), group=True):
+        if group:
+            vertex_indices = self.triangles[key]
+        else:
+            vertex_indices = self.triangles[key].flatten()
+
+        return self.vertices[vertex_indices]
 
     def __len__(self):
         return len(self.triangles)
@@ -167,15 +172,6 @@ def morton_order(mesh, bits):
     bits of the quantized center coordinates of each triangle. Each coordinate
     axis is quantized into 2**bits bins.
     """
-
-    #lower_bound = np.array([np.min(mesh[:,:,0]),
-    #                        np.min(mesh[:,:,1]),
-    #                        np.min(mesh[:,:,2])])
-
-    #upper_bound = np.array([np.max(mesh[:,:,0]),
-    #                        np.max(mesh[:,:,1]),
-    #                        np.max(mesh[:,:,2])])
-
     lower_bound, upper_bound = mesh.get_bounds()
 
     if bits <= 0 or bits > 10:
@@ -186,8 +182,7 @@ def morton_order(mesh, bits):
     def quantize(x):
         return np.uint32((x-lower_bound)*max_value/(upper_bound-lower_bound))
 
-    #mean_positions = quantize(np.mean(mesh, axis=1))
-    mean_positions = quantize(np.mean(mesh[:], axis=1))
+    mean_positions = quantize(np.mean(mesh.assemble(), axis=1))
 
     return interleave(mean_positions, bits)
 
@@ -288,14 +283,9 @@ class Geometry(object):
             i1 = np.searchsorted(zvalues_mesh, z)
             i2 = np.searchsorted(zvalues_mesh, z, side='right')
 
-            self.lower_bounds[i] = [np.min(self.mesh[i1:i2][:,:,0]),
-                                    np.min(self.mesh[i1:i2][:,:,1]),
-                                    np.min(self.mesh[i1:i2][:,:,2])]
-            
-            self.upper_bounds[i] = [np.max(self.mesh[i1:i2][:,:,0]),
-                                    np.max(self.mesh[i1:i2][:,:,1]),
-                                    np.max(self.mesh[i1:i2][:,:,2])]
-            
+            self.lower_bounds[i] = np.min(self.mesh.assemble(slice(i1,i2), group=False), axis=0)
+            self.upper_bounds[i] = np.max(self.mesh.assemble(slice(i1,i2), group=False), axis=0)
+
             self.node_map[i] = i1
             self.node_length[i] = i2-i1
             self.layers[i] = layer