import numpy as np import numpy.ma as ma import pycuda.driver as cuda from pycuda import gpuarray from mesh import Mesh 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, solids=[]): self.solids = solids def add_solid(self, solid): self.solids.append(solid) def build(self, bits=8): vertices = [] triangles = [] for solid in self.solids: triangles.extend(solid.mesh.triangles + len(vertices)) vertices.extend(np.inner(solid.mesh.vertices, solid.rotation) + \ solid.displacement) 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.uint32) self.material2_index = np.empty(len(self.mesh), dtype=np.uint32) for i, material in enumerate(material1[reorder]): self.material1_index[i] = self.materials.index(material) for i, material in enumerate(material2[reorder]): self.material2_index[i] = self.materials.index(material) surface1 = np.concatenate([solid.surface1 for solid in self.solids]) surface2 = np.concatenate([solid.surface2 for solid in self.solids]) self.surfaces = list(np.unique(np.concatenate([surface1, surface2]))) self.surface1_index = np.empty(len(self.mesh), dtype=np.uint32) self.surface2_index = np.empty(len(self.mesh), dtype=np.uint32) for i, surface in enumerate(surface1[reorder]): self.surface1_index[i] = self.surfaces.index(surface) for i, surface in enumerate(surface2[reorder]): self.surface2_index[i] = self.surfaces.index(surface) self.colors = np.concatenate([solid.color for solid in self.solids]) self.solid_id = np.concatenate([np.tile(solid.id, len(solid.mesh)) \ for solid in 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. """ 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 << 24) | \ (self.material2 << 16) | (self.surface1 << 8) | self.surface2 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) vertices_tex = module.get_texref('vertices') triangles_tex = module.get_texref('triangles') 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) triangles_tex.set_address(self.triangles_gpu, triangles.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) triangles_tex.set_format(cuda.array_format.UNSIGNED_INT32, 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, triangles_tex, lower_bounds_tex, upper_bounds_tex, node_map_tex, node_length_tex]