diff options
| -rw-r--r-- | geometry.py | 103 |
1 files changed, 65 insertions, 38 deletions
diff --git a/geometry.py b/geometry.py index b08de54..856ded3 100644 --- a/geometry.py +++ b/geometry.py @@ -2,7 +2,7 @@ import sys import os import numpy as np from itertoolset import * -from tools import timeit +from tools import timeit, profile_if_possible from hashlib import md5 import cPickle as pickle import gzip @@ -36,14 +36,33 @@ class Mesh(object): self.remove_duplicate_vertices() def get_bounds(self): + "Return the lower and upper bounds for the mesh as a tuple." return np.min(self.vertices, axis=0), np.max(self.vertices, axis=0) def remove_duplicate_vertices(self): + "Remove any duplicate vertices in the mesh." + # view the vertices as a structured array in order to identify unique + # rows, i.e. unique vertices unique_vertices, inverse = np.unique(self.vertices.view([('', self.vertices.dtype)]*self.vertices.shape[1]), return_inverse=True) + # turn the structured vertex array back into a normal array self.vertices = unique_vertices.view(self.vertices.dtype).reshape((unique_vertices.shape[0], self.vertices.shape[1])) + # remap the triangles to the unique vertices self.triangles = np.vectorize(lambda i: inverse[i])(self.triangles) def assemble(self, key=slice(None), group=True): + """ + Return an assembled triangle mesh; i.e. return the vertex positions + of every triangle. If `group` is True, the array returned will have + an extra axis denoting the triangle number; i.e. if the mesh contains + N triangles, the returned array will have the shape (N,3,3). If `group` + is False, return just the vertex positions without any grouping; in + this case the grouping is implied (i.e. elements [0:3] are the first + triangle, [3:6] the second, and so on. + + The `key` argument is a slice object if you just want to assemble + a certain range of the triangles, i.e. to get the assembled mesh for + triangles 3 through 6, key = slice(3,7). + """ if group: vertex_indices = self.triangles[key] else: @@ -51,9 +70,6 @@ class Mesh(object): return self.vertices[vertex_indices] - def __len__(self): - return len(self.triangles) - def __add__(self, other): return Mesh(np.concatenate((self.vertices, other.vertices)), np.concatenate((self.triangles, other.triangles + len(self.vertices)))) @@ -62,41 +78,38 @@ class Solid(object): self.mesh = mesh if np.iterable(material1): - if len(material1) != len(mesh): + if len(material1) != len(mesh.triangles): raise ValueError('shape mismatch') self.material1 = np.array(material1, dtype=np.object) else: - self.material1 = np.tile(material1, len(self.mesh)) + self.material1 = np.tile(material1, len(self.mesh.triangles)) if np.iterable(material2): - if len(material2) != len(mesh): + if len(material2) != len(mesh.triangles): raise ValueError('shape mismatch') self.material2 = np.array(material2, dtype=np.object) else: - self.material2 = np.tile(material2, len(self.mesh)) + self.material2 = np.tile(material2, len(self.mesh.triangles)) if np.iterable(surface): - if len(surface) != len(mesh): + if len(surface) != len(mesh.triangles): raise ValueError('shape mismatch') self.surface = np.array(surface, dtype=np.object) else: - self.surface = np.tile(surface, len(self.mesh)) + self.surface = np.tile(surface, len(self.mesh.triangles)) if np.iterable(color): - if len(color) != len(mesh): + if len(color) != len(mesh.triangles): raise ValueError('shape mismatch') self.color = np.array(color, dtype=np.uint32) else: - self.color = np.tile(color, len(self.mesh)).astype(np.uint32) + self.color = np.tile(color, len(self.mesh.triangles)).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) - def __add__(self, other): return Solid(self.mesh + other.mesh, np.concatenate((self.material1, other.material1)), np.concatenate((self.material2, other.material2)), np.concatenate((self.surface, other.surface)), np.concatenate((self.color, other.color))) @@ -187,17 +200,23 @@ class Geometry(object): self.solid_displacements = [] def add_solid(self, solid, rotation=np.identity(3), displacement=(0,0,0)): + """ + Add the solid `solid` to the geometry. When building the final triangle + mesh, `solid` will be placed by rotating it with the rotation matrix + `rotation` and displacing it by the vector `displacement`. + """ + rotation = np.asarray(rotation, dtype=np.float32) if rotation.shape != (3,3): - raise ValueError('shape mismatch') + raise ValueError('rotation matrix has the wrong shape.') self.solid_rotations.append(rotation.astype(np.float32)) displacement = np.asarray(displacement, dtype=np.float32) if displacement.shape != (3,): - raise ValueError('shape mismatch') + raise ValueError('displacement vector has the wrong shape.') self.solid_displacements.append(displacement) @@ -207,21 +226,32 @@ class Geometry(object): @timeit def build(self, bits=8, shift=3): - offsets = [ (0,0) ] - for solid in self.solids: - offsets.append( (offsets[-1][0] + len(solid.mesh.vertices), - offsets[-1][1] + len(solid.mesh.triangles)) ) - - vertices = np.zeros(shape=(offsets[-1][0], 3), dtype=np.float32) - triangles = np.zeros(shape=(offsets[-1][1],3), dtype=np.int32) - - for i, (solid, (vertex_offset, triangle_offset)) in \ - enumerate(zip(self.solids, offsets[:-1])): - triangles[triangle_offset:triangle_offset+len(solid.mesh.triangles),:] = \ - solid.mesh.triangles + vertex_offset - vertices[vertex_offset:vertex_offset+len(solid.mesh.vertices),:] = \ + """ + Build the bounding volume hierarchy, material/surface code arrays, and + color array for this geometry. If the bounding volume hierarchy is + cached, load the cache instead of rebuilding, else build and cache it. + + Args: + - bits: int, *optional* + The number of bits to quantize each linear dimension with when + morton ordering the triangle centers for building the bounding + volume hierarchy. Defaults to 8. + - shift: int, *optional* + The number of bits to shift the zvalue of each node when + building the next layer of the bounding volume hierarchy. + Defaults to 3. + """ + 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]) + + vertices = np.empty((nv[-1],3), dtype=np.float32) + triangles = np.empty((nt[-1],3), dtype=np.uint32) + + for i, solid in enumerate(self.solids): + vertices[nv[i]:nv[i+1]] = \ np.inner(solid.mesh.vertices, self.solid_rotations[i]) + self.solid_displacements[i] - + triangles[nt[i]:nt[i+1]] = solid.mesh.triangles + nv[i] + self.mesh = Mesh(vertices, triangles) self.colors = np.concatenate([solid.color for solid in self.solids]) @@ -232,21 +262,18 @@ class Geometry(object): material_lookup = dict(zip(self.unique_materials, range(len(self.unique_materials)))) - self.material1_index = \ - np.fromiter(imap(material_lookup.get, chain(*[solid.material1 for solid in self.solids])), dtype=np.int32) + self.material1_index = np.fromiter(imap(material_lookup.get, chain(*[solid.material1 for solid in self.solids])), dtype=np.int32) - self.material2_index = \ - np.fromiter(imap(material_lookup.get, chain(*[solid.material2 for solid in self.solids])), dtype=np.int32) + self.material2_index = np.fromiter(imap(material_lookup.get, chain(*[solid.material2 for solid in self.solids])), dtype=np.int32) 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, chain(*[solid.surface for solid in self.solids])), dtype=np.int32) + self.surface_index = np.fromiter(imap(surface_lookup.get, chain(*[solid.surface for solid in self.solids])), dtype=np.int32) try: - self.surface_index[self.surface_index == self.unique_surfaces.index(None)] = -1 + self.surface_index[self.surface_index == surface_lookup[None]] = -1 except ValueError: pass |