import numpy as np import pycuda.driver as cuda from pycuda import gpuarray # 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): 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 def build(self): return self.vertices[self.triangles] 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]))) self.material1_index = np.empty(len(self.mesh), dtype=np.int32) self.material2_index = np.empty(len(self.mesh), 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 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 surface = np.concatenate([solid.surface for solid in self.solids]) self.surfaces = list(np.unique(surface)) self.surface_index = np.empty(len(self.mesh), 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.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: 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.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)) self.material1_index_tex = module.get_texref('material1_lookup') self.material2_index_tex = module.get_texref('material2_lookup') self.surface_index_tex = module.get_texref('surface_lookup') self.material1_index_gpu = cuda.to_device(self.material1_index) self.material2_index_gpu = cuda.to_device(self.material2_index) self.surface_index_gpu = cuda.to_device(self.surface_index) self.material1_index_tex.set_address(self.material1_index_gpu, self.material1_index.nbytes) self.material2_index_tex.set_address(self.material2_index_gpu, self.material2_index.nbytes) self.surface_index_tex.set_address(self.surface_index_gpu, self.surface_index.nbytes) self.material1_index_tex.set_format(cuda.array_format.SIGNED_INT32, 1) self.material2_index_tex.set_format(cuda.array_format.SIGNED_INT32, 1) self.surface_index_tex.set_format(cuda.array_format.SIGNED_INT32, 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 << 24) | (self.material2_index << 16) | (self.surface_index << 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) set_pointer = module.get_function('set_pointer') set_pointer(self.triangles_gpu, block=(1,1,1), grid=(1,1)) vertices_tex = module.get_texref('vertices') 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') vertices_tex.set_address(self.vertices_gpu, vertices.nbytes) 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) vertices_tex.set_format(cuda.array_format.FLOAT, 4) 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 [vertices_tex, lower_bounds_tex, upper_bounds_tex, node_map_tex, node_length_tex]