diff options
| -rw-r--r-- | detectors/lbne.py | 3 | ||||
| -rw-r--r-- | geometry.py | 57 | ||||
| -rw-r--r-- | itertoolset.py | 19 | ||||
| -rw-r--r-- | make.py | 90 | ||||
| -rwxr-xr-x | view.py | 37 |
5 files changed, 124 insertions, 82 deletions
diff --git a/detectors/lbne.py b/detectors/lbne.py index c98fa42..9f74975 100644 --- a/detectors/lbne.py +++ b/detectors/lbne.py @@ -25,7 +25,8 @@ class LBNE(Geometry): # outer cylinder cylinder_mesh = \ - make.cylinder(radius, height+height/(pmts_per_string-1)) + make.cylinder(radius, radius, height+height/(pmts_per_string-1)) + cylinder_mesh.vertices = rotate(cylinder_mesh.vertices, np.pi/2, (-1,0,0)) cylinder_solid = Solid(cylinder_mesh, lightwater_sno, vacuum, black_surface, 0x0000ff) self.add_solid(cylinder_solid) diff --git a/geometry.py b/geometry.py index 111006f..3181176 100644 --- a/geometry.py +++ b/geometry.py @@ -1,13 +1,15 @@ 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): + def __init__(self, vertices, triangles, remove_duplicate_vertices=False): vertices = np.asarray(vertices, dtype=np.float32) triangles = np.asarray(triangles, dtype=np.int32) @@ -27,8 +29,20 @@ class Mesh(object): self.vertices = vertices self.triangles = triangles - def build(self): - return self.vertices[self.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]] @@ -216,32 +230,24 @@ class Geometry(object): 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) + material_lookup = dict(zip(self.materials, range(len(self.materials)))) + + self.material1_index = \ + np.fromiter(imap(material_lookup.get, material1[reorder]), 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 + self.material2_index = \ + np.fromiter(imap(material_lookup.get, material2[reorder]), dtype=np.int32) - 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 + surfaces = np.concatenate([solid.surface for solid in self.solids]) - surface = np.concatenate([solid.surface for solid in self.solids]) + self.unique_surfaces = list(np.unique(surfaces)) - self.surfaces = list(np.unique(surface)) + surface_lookup = dict(zip(self.unique_surfaces, range(len(self.unique_surfaces)))) - self.surface_index = np.empty(len(self.mesh), dtype=np.int32) + self.surface_index = \ + np.fromiter(imap(surface_lookup.get, surfaces[reorder]), 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.surface_index[self.surface_index == self.unique_surfaces.index(None)] = -1 self.colors = \ np.concatenate([solid.color for solid in self.solids])[reorder] @@ -337,8 +343,7 @@ class Geometry(object): for i, material in enumerate(self.materials): if material is None: if color is False: - print 'warning: some triangles have no associated material.' - + 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) @@ -352,7 +357,7 @@ class Geometry(object): 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): + for i, surface in enumerate(self.unique_surfaces): if surface is None: continue diff --git a/itertoolset.py b/itertoolset.py index 13bb896..ae4be8a 100644 --- a/itertoolset.py +++ b/itertoolset.py @@ -1,4 +1,4 @@ -from itertools import cycle, islice +from itertools import * def roundrobin(*iterables): """roundrobin('ABC', 'D', 'EF') --> A D E B F C""" @@ -11,3 +11,20 @@ def roundrobin(*iterables): except StopIteration: pending -= 1 nexts = cycle(islice(nexts, pending)) + +def unique_everseen(iterable, key=None): + "List unique elements, preserving order. Remember all elements ever seen." + # unique_everseen('AAAABBBCCDAABBB') --> A B C D + # unique_everseen('ABBCcAD', str.lower) --> A B C D + seen = set() + seen_add = seen.add + if key is None: + for element in ifilterfalse(seen.__contains__, iterable): + seen_add(element) + yield element + else: + for element in iterable: + k = key(element) + if k not in seen: + seen_add(k) + yield element @@ -1,43 +1,57 @@ import numpy as np from geometry import Mesh +from transform import rotate + +def rotate_extrude(x, y, theta=np.pi/32): + x, y, = np.asarray(x), np.asarray(y) + + if len(x.shape) != 1 or len(y.shape) != 1 or x.size != y.size: + raise ValueError('shape mismatch') + + points = np.zeros((len(x),3), dtype=np.float32) + points[:,0] = x + points[:,1] = y -def cylinder(radius=1.0, height=1.0, theta=np.pi/32): angles = np.arange(0, 2*np.pi, theta) - vertices = np.empty((2*len(angles) + 2, 3), dtype=np.float32) - vertices[:2*len(angles),0] = radius*np.tile(np.cos(angles),2) - vertices[:2*len(angles),1] = radius*np.tile(np.sin(angles),2) - vertices[:len(angles),2] = height/2.0 - vertices[len(angles):2*len(angles),2] = -height/2.0 - vertices[-2] = [0,0,height/2.0] - vertices[-1] = [0,0,-height/2.0] - - triangles = np.zeros((len(angles)*4, 3), dtype=np.uint32) - - # top - triangles[:len(angles),0] = np.arange(len(angles)) - triangles[:len(angles),2] = np.roll(np.arange(len(angles)),1) - triangles[:len(angles),1] = len(vertices) - 2 - - # bottom - triangles[len(angles):2*len(angles),0] = \ - np.arange(len(angles),2*len(angles)) - triangles[len(angles):2*len(angles),1] = \ - np.roll(np.arange(len(angles),2*len(angles)),1) - triangles[len(angles):2*len(angles),2] = len(vertices) - 1 - - # side - triangles[2*len(angles):3*len(angles),0] = np.arange(len(angles)) - triangles[2*len(angles):3*len(angles),2] = \ - np.arange(len(angles),2*len(angles)) - triangles[2*len(angles):3*len(angles),1] = \ - np.roll(np.arange(len(angles),2*len(angles)),1) - - # side - triangles[3*len(angles):4*len(angles),0] = np.arange(len(angles)) - triangles[3*len(angles):4*len(angles),2] = \ - np.roll(np.arange(len(angles),2*len(angles)),1) - triangles[3*len(angles):4*len(angles),1] = \ - np.roll(np.arange(len(angles)),1) - - return Mesh(vertices, triangles) + vertices = np.empty((len(points)*len(angles), 3), dtype=np.float32) + triangles = np.zeros((len(angles)*(len(points)-1)*2, 3), dtype=np.int32) + + step = len(points) - 1 + + for i, angle in enumerate(angles): + this_slice = i*len(points) + vertices[this_slice:this_slice+len(points)] = rotate(points, angle, (0,-1,0)) + + start = 2*i*step + next_slice = ((i+1) % len(angles))*len(points) + + triangles[start:start+step,0] = np.arange(this_slice, this_slice+len(points)-1) + triangles[start:start+step,1] = np.arange(next_slice, next_slice+len(points)-1) + triangles[start:start+step,2] = np.arange(this_slice+1, this_slice+len(points)) + + triangles[start+step:start+2*step,0] = np.arange(next_slice, next_slice+len(points)-1) + triangles[start+step:start+2*step,2] = np.arange(this_slice+1, this_slice+len(points)) + triangles[start+step:start+2*step,1] = np.arange(next_slice+1, next_slice+len(points)) + + return Mesh(vertices, triangles, remove_duplicate_vertices=True) + +def cube(size=1): + a = np.sqrt(size**2/2.0) + x = [0,a,a,0] + y = [-size/2.0,-size/2.0,size/2.0,size/2.0] + return rotate_extrude(x, y, np.pi/2) + +def cylinder(radius1=1, radius2=1, height=2, theta=np.pi/32): + x = [0,radius1,radius2,0] + y = [-height/2.0, -height/2.0, height/2.0, height/2.0] + return rotate_extrude(x, y, theta) + +def sphere(radius=1, theta=np.pi/32): + profile_angles = np.arange(-np.pi/2, np.pi/2+theta, theta) + return rotate_extrude(np.cos(profile_angles), np.sin(profile_angles), theta) + +def torus(radius=1, offset=3, phi=np.pi/32, theta=np.pi/32): + profile_angles = np.arange(0, 2*np.pi+phi, phi) + x, y = np.cos(profile_angles), np.sin(profile_angles) + return rotate_extrude(x + offset, y) @@ -1,5 +1,6 @@ #!/usr/bin/env python import os +import sys import numpy as np import pygame @@ -16,9 +17,9 @@ from pycuda import autoinit from pycuda.compiler import SourceModule from pycuda import gpuarray -def view(geometry, name=''): +def view(viewable, name='', bits=8): """ - Render `geometry` in a pygame window. + Render `viewable` in a pygame window. Movement: - zoom: scroll the mouse wheel @@ -26,6 +27,20 @@ def view(geometry, name=''): - move: shift+click and drag the mouse """ + if isinstance(viewable, Geometry): + geometry = viewable + geometry.build(bits) + elif isinstance(viewable, Solid): + geometry = Geometry() + geometry.add_solid(viewable) + geometry.build(bits) + elif isinstance(viewable, Mesh): + geometry = Geometry() + geometry.add_solid(Solid(viewable)) + geometry.build(bits) + else: + sys.exit("can't display %s" % args[0]) + lower_bound = np.array([np.min(geometry.mesh[:][:,:,0]), np.min(geometry.mesh[:][:,:,1]), np.min(geometry.mesh[:][:,:,2])]) @@ -192,12 +207,12 @@ def view(geometry, name=''): if event.key == K_LSHIFT or event.key == K_RSHIFT: shift = False + pygame.display.quit() + if __name__ == '__main__': - import os - import sys import optparse - from stl import * + from stl import mesh_from_stl import solids import detectors @@ -224,16 +239,6 @@ if __name__ == '__main__': members = dict(inspect.getmembers(detectors) + inspect.getmembers(solids)) if args[0] in members: - if isinstance(members[args[0]], Geometry): - geometry = members[args[0]] - geometry.build(options.bits) - view(geometry, args[0]) - elif isinstance(members[args[0]], Solid): - geometry = Geometry() - geometry.add_solid(members[args[0]]) - geometry.build(options.bits) - view(geometry, args[0]) - else: - sys.exit("can't display %s" % args[0]) + view(members[args[0]], args[0], options.bits) else: sys.exit("couldn't find object %s" % args[0]) |