diff options
Diffstat (limited to 'geometry.py')
| -rw-r--r-- | geometry.py | 434 |
1 files changed, 434 insertions, 0 deletions
diff --git a/geometry.py b/geometry.py new file mode 100644 index 0000000..3181176 --- /dev/null +++ b/geometry.py @@ -0,0 +1,434 @@ +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, remove_duplicate_vertices=False): + vertices = np.asarray(vertices, dtype=np.float32) + triangles = np.asarray(triangles, dtype=np.int32) + + if len(vertices.shape) != 2 or vertices.shape[1] != 3: + raise ValueError('shape mismatch') + + if len(triangles.shape) != 2 or triangles.shape[1] != 3: + raise ValueError('shape mismatch') + + if (triangles < 0).any(): + raise ValueError('indices in `triangles` must be positive.') + + if (triangles >= len(vertices)).any(): + raise ValueError('indices in `triangles` must be less than the ' + 'length of the vertex array.') + + self.vertices = vertices + self.triangles = 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]] + + 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)))) + +class Solid(object): + def __init__(self, mesh, material1=None, material2=None, surface=None, color=0xffffffff): + self.mesh = mesh + + if np.iterable(material1): + if len(material1) != len(mesh): + raise ValueError('shape mismatch') + self.material1 = np.array(material1, dtype=np.object) + else: + self.material1 = np.tile(material1, len(self.mesh)) + + if np.iterable(material2): + if len(material2) != len(mesh): + raise ValueError('shape mismatch') + self.material2 = np.array(material2, dtype=np.object) + else: + self.material2 = np.tile(material2, len(self.mesh)) + + if np.iterable(surface): + if len(surface) != len(mesh): + raise ValueError('shape mismatch') + self.surface = np.array(surface, dtype=np.object) + else: + self.surface = np.tile(surface, len(self.mesh)) + + if np.iterable(color): + if len(color) != len(mesh): + raise ValueError('shape mismatch') + self.color = np.array(color, dtype=np.uint32) + else: + self.color = np.tile(color, len(self.mesh)).astype(np.uint32) + + 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))) + +class Material(object): + """Material optical properties.""" + def __init__(self, name='none'): + self.name = name + + self.refractive_index = None + self.absorption_length = None + self.scattering_length = None + + def set(self, name, value, wavelengths=standard_wavelengths): + if np.iterable(value): + if len(value) != len(wavelengths): + raise ValueError('shape mismatch') + else: + value = np.tile(value, len(wavelengths)) + + self.__dict__[name] = np.array(zip(wavelengths, value), dtype=np.float32) + +class Surface(object): + """Surface optical properties.""" + def __init__(self, name='none'): + self.name = name + + self.absorption = None + self.reflection_diffuse = None + self.reflection_specular = None + + def set(self, name, value, wavelengths=standard_wavelengths): + if np.iterable(value): + if len(value) != len(wavelengths): + raise ValueError('shape mismatch') + else: + value = np.tile(value, len(wavelengths)) + + self.__dict__[name] = np.array(zip(wavelengths, value), dtype=np.float32) + +def interleave(arr, bits): + """ + Interleave the bits of quantized three-dimensional points in space. + + Example + >>> interleave(np.identity(3, dtype=np.int)) + array([4, 2, 1], dtype=uint64) + """ + if len(arr.shape) != 2 or arr.shape[1] != 3: + raise Exception('shape mismatch') + + z = np.zeros(arr.shape[0], dtype=np.uint32) + for i in range(bits): + z |= (arr[:,2] & 1 << i) << (2*i) | \ + (arr[:,1] & 1 << i) << (2*i+1) | \ + (arr[:,0] & 1 << i) << (2*i+2) + return z + +def morton_order(mesh, bits): + """ + Return a list of zvalues for triangles in `mesh` by interleaving the + 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])]) + + if bits <= 0 or bits > 12: + raise Exception('number of bits must be in the range (0,12].') + + max_value = 2**bits - 1 + + def quantize(x): + return np.uint32((x-lower_bound)*max_value/(upper_bound-lower_bound)) + + mean_positions = quantize(np.mean(mesh, axis=1)) + + return interleave(mean_positions, bits) + +class Geometry(object): + def __init__(self): + self.solids = [] + self.solid_rotations = [] + self.solid_displacements = [] + + def add_solid(self, solid, rotation=np.identity(3), displacement=(0,0,0)): + rotation = np.asarray(rotation, dtype=np.float32) + + if rotation.shape != (3,3): + raise ValueError('shape mismatch') + + self.solid_rotations.append(rotation.astype(np.float32)) + + displacement = np.asarray(displacement, dtype=np.float32) + + if displacement.shape != (3,): + raise ValueError('shape mismatch') + + self.solid_displacements.append(displacement) + + self.solids.append(solid) + + return len(self.solids)-1 + + def build(self, bits=8): + 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),:] = \ + np.inner(solid.mesh.vertices, self.solid_rotations[i]) + self.solid_displacements[i] + + self.mesh = Mesh(vertices, triangles) + + zvalues_mesh = morton_order(self.mesh[:], bits) + reorder = np.argsort(zvalues_mesh) + zvalues_mesh = zvalues_mesh[reorder] + + if (np.diff(zvalues_mesh) < 0).any(): + raise Exception('zvalues_mesh out of order.') + + 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]))) + + material_lookup = dict(zip(self.materials, range(len(self.materials)))) + + self.material1_index = \ + np.fromiter(imap(material_lookup.get, material1[reorder]), dtype=np.int32) + + self.material2_index = \ + np.fromiter(imap(material_lookup.get, material2[reorder]), dtype=np.int32) + + surfaces = np.concatenate([solid.surface for solid in self.solids]) + + self.unique_surfaces = list(np.unique(surfaces)) + + 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) + + self.surface_index[self.surface_index == self.unique_surfaces.index(None)] = -1 + + self.colors = \ + np.concatenate([solid.color for solid in self.solids])[reorder] + + self.solid_id = \ + np.concatenate([np.tile(i, len(solid.mesh)) for i, solid in \ + 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) + + 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] = [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.node_map[i] = i1 + self.node_length[i] = i2-i1 + + zvalues[i] = z + + self.first_node = unique_zvalues.size + + begin_last_layer = 0 + + while True: + bit_shifted_zvalues = zvalues >> 1 + unique_bit_shifted_zvalues = np.unique(bit_shifted_zvalues) + zvalues = np.empty(unique_bit_shifted_zvalues.size, dtype=np.uint64) + + self.lower_bounds.resize(\ + (self.lower_bounds.shape[0]+unique_bit_shifted_zvalues.size,3)) + self.upper_bounds.resize(\ + (self.upper_bounds.shape[0]+unique_bit_shifted_zvalues.size,3)) + + self.node_map.resize(\ + self.node_map.size+unique_bit_shifted_zvalues.size) + self.node_length.resize(\ + self.node_length.size+unique_bit_shifted_zvalues.size) + + for i, z in enumerate(unique_bit_shifted_zvalues): + i1 = np.searchsorted(bit_shifted_zvalues, z) + \ + begin_last_layer + i2 = np.searchsorted(bit_shifted_zvalues, z, side='right') + \ + begin_last_layer + + zvalues[i] = z + + i += begin_last_layer + bit_shifted_zvalues.size + + self.lower_bounds[i] = \ + [np.min(self.lower_bounds[i1:i2,0]), + np.min(self.lower_bounds[i1:i2,1]), + np.min(self.lower_bounds[i1:i2,2])] + + self.upper_bounds[i] = \ + [np.max(self.upper_bounds[i1:i2,0]), + np.max(self.upper_bounds[i1:i2,1]), + np.max(self.upper_bounds[i1:i2,2])] + + self.node_map[i] = i1 + self.node_length[i] = i2-i1 + + begin_last_layer += bit_shifted_zvalues.size + + if unique_bit_shifted_zvalues.size == 1: + break + + def load(self, module, color=False): + """ + Load the bounding volume hierarchy onto the GPU module, + bind it to the appropriate textures, and return a list + of the texture references. + """ + + set_wavelength_range = module.get_function('set_wavelength_range') + + 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): + if material is None: + if color is False: + 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) + absorption_length = np.interp(standard_wavelengths, material.absorption_length[:,0], material.absorption_length[:,1]).astype(np.float32) + scattering_length = np.interp(standard_wavelengths, material.scattering_length[:,0], material.scattering_length[:,1]).astype(np.float32) + + material.refractive_index_gpu = cuda.to_device(refractive_index) + material.absorption_length_gpu = cuda.to_device(absorption_length) + material.scattering_length_gpu = cuda.to_device(scattering_length) + + 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.unique_surfaces): + if surface is None: + continue + + absorption = np.interp(standard_wavelengths, surface.absorption[:,0], surface.absorption[:,1]).astype(np.float32) + reflection_diffuse = np.interp(standard_wavelengths, surface.reflection_diffuse[:,0], surface.reflection_diffuse[:,1]).astype(np.float32) + reflection_specular = np.interp(standard_wavelengths, surface.reflection_specular[:,0], surface.reflection_specular[:,1]).astype(np.float32) + + surface.absorption_gpu = cuda.to_device(absorption) + surface.reflection_diffuse_gpu = cuda.to_device(reflection_diffuse) + surface.reflection_specular_gpu = cuda.to_device(reflection_specular) + + set_surface(np.int32(i), surface.absorption_gpu, surface.reflection_diffuse_gpu, surface.reflection_specular_gpu, block=(1,1,1), grid=(1,1)) + + vertices = np.empty(len(self.mesh.vertices), dtype=gpuarray.vec.float4) + vertices['x'] = self.mesh.vertices[:,0] + vertices['y'] = self.mesh.vertices[:,1] + vertices['z'] = self.mesh.vertices[:,2] + + triangles = \ + np.empty(len(self.mesh.triangles), dtype=gpuarray.vec.uint4) + triangles['x'] = self.mesh.triangles[:,0] + triangles['y'] = self.mesh.triangles[:,1] + triangles['z'] = self.mesh.triangles[:,2] + + if color: + triangles['w'] = self.colors + else: + triangles['w'] = ((self.material1_index & 0xff) << 24) | ((self.material2_index & 0xff) << 16) | ((self.surface_index & 0xff) << 8) + + lower_bounds = np.empty(self.lower_bounds.shape[0], dtype=gpuarray.vec.float4) + lower_bounds['x'] = self.lower_bounds[:,0] + lower_bounds['y'] = self.lower_bounds[:,1] + lower_bounds['z'] = self.lower_bounds[:,2] + + upper_bounds = np.empty(self.upper_bounds.shape[0], dtype=gpuarray.vec.float4) + upper_bounds['x'] = self.upper_bounds[:,0] + upper_bounds['y'] = self.upper_bounds[:,1] + upper_bounds['z'] = self.upper_bounds[:,2] + + self.vertices_gpu = cuda.to_device(vertices) + self.triangles_gpu = cuda.to_device(triangles) + self.lower_bounds_gpu = cuda.to_device(lower_bounds) + self.upper_bounds_gpu = cuda.to_device(upper_bounds) + self.node_map_gpu = cuda.to_device(self.node_map) + self.node_length_gpu = cuda.to_device(self.node_length) + + print 'Device usage:' + print 'vertices:', vertices.nbytes + print 'triangles:', triangles.nbytes + print 'lower_bounds:', lower_bounds.nbytes + print 'upper_bounds:', upper_bounds.nbytes + print 'node_map:', self.node_map.nbytes + print 'node_length:', self.node_length.nbytes + + set_pointer = module.get_function('set_pointer') + set_pointer(self.triangles_gpu, self.vertices_gpu, + block=(1,1,1), grid=(1,1)) + + lower_bounds_tex = module.get_texref('lower_bounds') + upper_bounds_tex = module.get_texref('upper_bounds') + node_map_tex = module.get_texref('node_map') + node_length_tex = module.get_texref('node_length') + + lower_bounds_tex.set_address(self.lower_bounds_gpu, lower_bounds.nbytes) + upper_bounds_tex.set_address(self.upper_bounds_gpu, upper_bounds.nbytes) + node_map_tex.set_address(self.node_map_gpu, self.node_map.nbytes) + node_length_tex.set_address(self.node_length_gpu, self.node_length.nbytes) + + lower_bounds_tex.set_format(cuda.array_format.FLOAT, 4) + upper_bounds_tex.set_format(cuda.array_format.FLOAT, 4) + node_map_tex.set_format(cuda.array_format.UNSIGNED_INT32, 1) + node_length_tex.set_format(cuda.array_format.UNSIGNED_INT32, 1) + + return [lower_bounds_tex, upper_bounds_tex, node_map_tex, node_length_tex] |