diff options
| -rw-r--r-- | geometry.py | 29 |
1 files changed, 14 insertions, 15 deletions
diff --git a/geometry.py b/geometry.py index 3181176..704fdf6 100644 --- a/geometry.py +++ b/geometry.py @@ -1,8 +1,8 @@ import numpy as np import pycuda.driver as cuda from pycuda import gpuarray -from itertools import imap -from itertoolset import unique_everseen +from itertools import * +from itertoolset import * # all material/surface properties are interpolated at these # wavelengths when they are sent to the gpu @@ -85,6 +85,11 @@ class Solid(object): else: self.color = np.tile(color, len(self.mesh)).astype(np.uint32) + self.unique_materials = \ + np.unique(np.concatenate([self.material1, self.material2])) + + self.unique_surfaces = np.unique(self.surface) + def __len__(self): return len(self.mesh) @@ -224,28 +229,22 @@ class Geometry(object): self.mesh.triangles = self.mesh.triangles[reorder] - material1 = np.concatenate([solid.material1 for solid in self.solids]) - material2 = np.concatenate([solid.material2 for solid in self.solids]) - - self.materials = \ - list(np.unique(np.concatenate([material1, material2]))) + self.unique_materials = list(np.unique(np.concatenate([solid.unique_materials for solid in self.solids]))) - material_lookup = dict(zip(self.materials, range(len(self.materials)))) + material_lookup = dict(zip(self.unique_materials, range(len(self.unique_materials)))) self.material1_index = \ - np.fromiter(imap(material_lookup.get, material1[reorder]), dtype=np.int32) + np.fromiter(imap(material_lookup.get, chain(*[solid.material1 for solid in self.solids])), dtype=np.int32)[reorder] self.material2_index = \ - np.fromiter(imap(material_lookup.get, material2[reorder]), dtype=np.int32) - - surfaces = np.concatenate([solid.surface for solid in self.solids]) + np.fromiter(imap(material_lookup.get, chain(*[solid.material2 for solid in self.solids])), dtype=np.int32)[reorder] - self.unique_surfaces = list(np.unique(surfaces)) + 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, surfaces[reorder]), dtype=np.int32) + np.fromiter(imap(surface_lookup.get, chain(*[solid.surface for solid in self.solids])), dtype=np.int32)[reorder] self.surface_index[self.surface_index == self.unique_surfaces.index(None)] = -1 @@ -340,7 +339,7 @@ class Geometry(object): set_wavelength_range(np.float32(standard_wavelengths[0]), np.float32(standard_wavelengths[-1]), np.float32(standard_wavelengths[1]-standard_wavelengths[0]), np.uint32(standard_wavelengths.size), block=(1,1,1), grid=(1,1)) set_material = module.get_function('set_material') - for i, material in enumerate(self.materials): + for i, material in enumerate(self.unique_materials): if material is None: if color is False: raise Exception('one or more triangles is missing a material.') |