added rotate_extrude() and a number of functions to produce meshes of useful solid shapes. tried to speed up geometry.build() by replacing some for loops with np.fromiter() and imap().

1 files changed, 31 insertions, 26 deletions

diff --git a/geometry.py b/geometry.py
index 04be460..d13fc9f 100644
--- a/geometry.py
+++ b/geometry.py
@@ -1,13 +1,15 @@
 import numpy as np
 import pycuda.driver as cuda
 from pycuda import gpuarray
+from itertools import imap
+from itertoolset import unique_everseen
 
 # all material/surface properties are interpolated at these
 # wavelengths when they are sent to the gpu
 standard_wavelengths = np.arange(200, 810, 20).astype(np.float32)
 
 class Mesh(object):
-    def __init__(self, vertices, triangles):
+    def __init__(self, vertices, triangles, remove_duplicate_vertices=False):
         vertices = np.asarray(vertices, dtype=np.float32)
         triangles = np.asarray(triangles, dtype=np.int32)
 
@@ -27,8 +29,20 @@ class Mesh(object):
         self.vertices = vertices
         self.triangles = triangles
 
-    def build(self):
-        return self.vertices[self.triangles]
+        if remove_duplicate_vertices:
+            self.remove_duplicate_vertices()
+
+    def remove_duplicate_vertices(self):
+        unique_vertices = list(unique_everseen(map(tuple, self.vertices)))
+
+        if len(unique_vertices) == len(self.vertices):
+            return
+
+        for i in range(len(self.triangles)):
+            for j in range(3):
+                self.triangles[i,j] = unique_vertices.index(tuple(self.vertices[self.triangles[i,j]]))
+
+        self.vertices = np.array(unique_vertices)
 
     def __getitem__(self, key):
         return self.vertices[self.triangles[key]]
@@ -216,32 +230,24 @@ class Geometry(object):
         self.materials = \
             list(np.unique(np.concatenate([material1, material2])))
 
-        self.material1_index = np.empty(len(self.mesh), dtype=np.int32)
-        self.material2_index = np.empty(len(self.mesh), dtype=np.int32)
+        material_lookup = dict(zip(self.materials, range(len(self.materials))))
+
+        self.material1_index = \
+            np.fromiter(imap(material_lookup.get, material1[reorder]), dtype=np.int32)
 
-        for i, material in enumerate(material1[reorder]):
-            if material is not None:
-                self.material1_index[i] = self.materials.index(material)
-            else:
-                self.material1_index[i] = -1
+        self.material2_index = \
+            np.fromiter(imap(material_lookup.get, material2[reorder]), dtype=np.int32)
 
-        for i, material in enumerate(material2[reorder]):
-            if material is not None:
-                self.material2_index[i] = self.materials.index(material)
-            else:
-                self.material2_index[i] = -1
+        surfaces = np.concatenate([solid.surface for solid in self.solids])
 
-        surface = np.concatenate([solid.surface for solid in self.solids])
+        self.unique_surfaces = list(np.unique(surfaces))
 
-        self.surfaces = list(np.unique(surface))
+        surface_lookup = dict(zip(self.unique_surfaces, range(len(self.unique_surfaces))))
 
-        self.surface_index = np.empty(len(self.mesh), dtype=np.int32)
+        self.surface_index = \
+            np.fromiter(imap(surface_lookup.get, surfaces[reorder]), dtype=np.int32)
 
-        for i, surface in enumerate(surface[reorder]):
-            if surface is not None:
-                self.surface_index[i] = self.surfaces.index(surface)
-            else:
-                self.surface_index[i] = -1
+        self.surface_index[self.surface_index == self.unique_surfaces.index(None)] = -1
 
         self.colors = \
             np.concatenate([solid.color for solid in self.solids])[reorder]
@@ -337,8 +343,7 @@ class Geometry(object):
         for i, material in enumerate(self.materials):
             if material is None:
                 if color is False:
-                    print 'warning: some triangles have no associated material.'
-
+                    raise Exception('one or more triangles is missing a material.')
                 continue
 
             refractive_index = np.interp(standard_wavelengths, material.refractive_index[:,0], material.refractive_index[:,1]).astype(np.float32)
@@ -352,7 +357,7 @@ class Geometry(object):
             set_material(np.int32(i), material.refractive_index_gpu, material.absorption_length_gpu, material.scattering_length_gpu, block=(1,1,1), grid=(1,1))
 
         set_surface = module.get_function('set_surface')
-        for i, surface in enumerate(self.surfaces):
+        for i, surface in enumerate(self.unique_surfaces):
             if surface is None:
                 continue