import numpy as np import numpy.ma as ma import pycuda.driver as cuda from pycuda import gpuarray 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 Solid(object): """ Object which stores a closed triangle mesh associated with a physically distinct object. Args: - mesh, array A closed triangle mesh. - material1, Material The material inside within the mesh. - material2, Material The material outside the mesh. - surface1, Surface, The surface on the inside of the mesh. - surface2, Surface, The surface on the outside of the mesh. .. warning:: It is possible to define logically inconsistent geometries unless you are careful. For example, solid A may define its inside material as water but contain solid B which defines its outside material as air. In this case, a photon traveling out of solid B will reflect/refract assuming it's going into air, but upon reaching solid A's boundary will calculate attenuation factors assuming it just traveled through water. """ def __init__(self, mesh, material1, material2, \ surface1=None, surface2=None): if len(mesh.shape) != 3 or mesh.shape[1] != 3 or mesh.shape[2] != 3: raise Exception('shape mismatch; mesh must be a triangle mesh') self.mesh = mesh self.material1 = material1 self.material2 = material2 self.surface1 = surface1 self.surface2 = surface2 def __len__(self): return self.mesh.shape[0] class Geometry(object): """Object which stores the global mesh for a geometry.""" def __init__(self): self.solids = [] self.materials = [] self.surfaces = [] def add_solid(self, solid): """Add a solid to the geometry.""" self.solids.append(solid) if solid.material1 not in self.materials: self.materials.append(solid.material1) if solid.material2 not in self.materials: self.materials.append(solid.material2) if solid.surface1 not in self.surfaces: self.surfaces.append(solid.surface1) if solid.surface2 not in self.surfaces: self.surfaces.append(solid.surface1) def build(self, bits=8): """Build the bounding volume hierarchy of the geometry.""" mesh = np.concatenate([solid.mesh for solid in self.solids]) # lookup solid/material/surface index from triangle index solid_index = \ np.concatenate([np.tile(self.solids.index(solid), \ len(solid)) for solid in self.solids]) material1_index = \ np.concatenate([np.tile(self.materials.index(solid.material1), \ len(solid)) for solid in self.solids]) material2_index = \ np.concatenate([np.tile(self.materials.index(solid.material2), \ len(solid)) for solid in self.solids]) surface1_index = \ np.concatenate([np.tile(self.surfaces.index(solid.surface1), \ len(solid)) for solid in self.solids]) surface2_index = \ np.concatenate([np.tile(self.surfaces.index(solid.surface2), \ len(solid)) for solid in self.solids]) zvalues_mesh = morton_order(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 = mesh[reorder] self.solid_index = solid_index[reorder] self.material1_index = material1_index[reorder] self.material2_index = material2_index[reorder] self.surface1_index = surface1_index[reorder] self.surface2_index = surface2_index[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): """ Load the bounding volume hierarchy onto the GPU module, bind it to the appropriate textures, and return a list of the texture references. """ mesh = np.empty(self.mesh.shape[0]*3, dtype=gpuarray.vec.float4) mesh['x'] = self.mesh[:,:,0].flatten() mesh['y'] = self.mesh[:,:,1].flatten() mesh['z'] = self.mesh[:,:,2].flatten() 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.mesh_gpu = cuda.to_device(mesh) 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) mesh_tex = module.get_texref('mesh') 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') mesh_tex.set_address(self.mesh_gpu, mesh.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) mesh_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 [mesh_tex, lower_bounds_tex, upper_bounds_tex, node_map_tex, node_length_tex]