#!/usr/bin/env python import numpy as np from itertools import product, count from threading import Thread, Lock import os import sys from chroma.color import map_wavelength, map_to_color from chroma.geometry import Mesh, Solid, Geometry import chroma.make as make import matplotlib.pyplot as plt from chroma import gpu from chroma.tools import timeit from chroma.transform import rotate from chroma.optics import vacuum, lambertian_surface import chroma.project as project from chroma.fileio.root import RootReader import pygame from pygame.locals import * from pycuda import gpuarray as ga from pycuda.characterize import sizeof import pycuda.driver as cuda from subprocess import call import shutil import tempfile import inspect def buildable(identifier): """ Create a decorator which tags a function as buildable and assigns the identifying string `identifier`. Example: >>> @buildable('my_sphere') >>> def build_my_sphere(): >>> g = Geometry() >>> g.add_solid(Solid(sphere(), vacuum, water)) >>> return g """ def tag_as_buildable(func): func.buildable = True func.identifier = identifier return func return tag_as_buildable def build(obj, bits): """Construct and build a geometry from `obj`.""" if inspect.isfunction(obj): try: if obj.buildable: obj = obj() except AttributeError: raise Exception('function %s is not buildable.' % obj.__name__) if isinstance(obj, Geometry): geometry = obj elif isinstance(obj, Solid): geometry = Geometry() geometry.add_solid(obj) elif isinstance(obj, Mesh): geometry = Geometry() geometry.add_solid(Solid(obj, vacuum, vacuum, surface=lambertian_surface, color=0x99ffffff)) else: raise Exception('cannot build type %s' % type(obj)) geometry.build(bits) return geometry def bvh_mesh(geometry, layer): lower_bounds = geometry.lower_bounds[geometry.layers == layer] upper_bounds = geometry.upper_bounds[geometry.layers == layer] if len(lower_bounds) == 0 or len(upper_bounds) == 0: raise Exception('no nodes at layer %i' % layer) dx, dy, dz = upper_bounds[0] - lower_bounds[0] center = np.mean([upper_bounds[0],lower_bounds[0]], axis=0) mesh = make.box(dx, dy, dz, center) for center, dx, dy, dz in zip(np.mean([lower_bounds,upper_bounds],axis=0),*zip(*upper_bounds-lower_bounds))[1:]: mesh += make.box(dx,dy,dz,center) return mesh def encode_movie(dir): root, ext = 'movie', 'avi' for i in count(): filename = '.'.join([root + str(i).zfill(5), ext]) if not os.path.exists(filename): break call(['mencoder', 'mf://' + dir + '/*.png', '-mf', 'fps=10', '-o', filename, '-ovc', 'xvid', '-xvidencopts', 'bitrate=3000']) shutil.rmtree(dir) print 'movie saved to %s.' % filename class Camera(Thread): def __init__(self, geometry, size=(800,600), device_id=None, enable3d=False, green_magenta=False, alpha_depth=10): Thread.__init__(self) self.geometry = geometry self.device_id = device_id self.size = size self.enable3d = enable3d self.green_magenta = green_magenta self.alpha_depth = alpha_depth self.unique_bvh_layers = np.unique(self.geometry.layers) self.currentlayer = None self.bvh_layers = {} try: import spnav as spnav_module self.spnav_module = spnav_module self.spnav = True except: self.spnav = False def init_gpu(self): self.gpu_instance = gpu.GPU(self.device_id) self.gpu_geometry = gpu.GPUGeometry(self.geometry) self.gpu_funcs = gpu.GPUFuncs(gpu.get_cu_module('mesh.h')) self.hybrid_funcs = gpu.GPUFuncs(gpu.get_cu_module('hybrid_render.cu')) self.gpu_geometries = [self.gpu_geometry] self.width, self.height = self.size self.npixels = self.width*self.height pygame.init() self.window = pygame.display.set_mode(self.size) self.screen = pygame.Surface(self.size, pygame.SRCALPHA) pygame.display.set_caption('') self.clock = pygame.time.Clock() self.doom_mode = False try: if self.width == 640: # SECRET DOOM MODE! print 'shotgun activated!' self.doom_hud = pygame.image.load('images/doomhud.png').convert_alpha() rect = self.doom_hud.get_rect() self.doom_rect = rect.move(0, self.height - rect.height) self.doom_mode = True except: pass lower_bound, upper_bound = self.geometry.mesh.get_bounds() self.scale = np.linalg.norm(upper_bound-lower_bound) self.nblocks = 64 self.point = np.array([0, -self.scale*1.0, (lower_bound[2]+upper_bound[2])/2]) self.axis1 = np.array([0,0,1], float) self.axis2 = np.array([1,0,0], float) self.film_width = 0.035 if self.enable3d: self.point1 = self.point-(self.scale/60,0,0) self.point2 = self.point+(self.scale/60,0,0) self.viewing_angle = 0.0 pos1, dir1 = project.from_film(self.point1, size=self.size, width=self.film_width) pos2, dir2 = project.from_film(self.point2, size=self.size, width=self.film_width) self.rays1 = gpu.GPURays(pos1, dir1) self.rays2 = gpu.GPURays(pos2, dir2) scope_pos, scope_dir = project.from_film(self.point, size=np.array(self.size)/4.0, width=self.film_width/4.0) self.scope_rays = gpu.GPURays(scope_pos, scope_dir) self.pixels1_gpu = ga.empty(self.width*self.height, dtype=np.int32) self.pixels2_gpu = ga.empty(self.width*self.height, dtype=np.int32) self.distances_gpu = ga.empty(self.scope_rays.pos.size, dtype=np.float32) else: pos, dir = project.from_film(self.point, size=self.size, width=self.film_width) self.rays = gpu.GPURays(pos, dir) self.distance_array = ga.empty(self.alpha_depth*self.rays.pos.size, dtype=np.float32) self.index_array = ga.empty(self.alpha_depth*self.rays.pos.size, dtype=np.uint32) self.n_array = ga.zeros(self.rays.pos.size, dtype=np.uint32) self.pixels_gpu = ga.empty(self.npixels, dtype=np.int32) self.movie = False self.movie_index = 0 self.movie_dir = None self.render = False @timeit def initialize_render(self): self.rng_states_gpu = gpu.get_rng_states(self.npixels) self.xyz_lookup1_gpu = ga.zeros(len(self.geometry.mesh.triangles), dtype=ga.vec.float3) self.xyz_lookup2_gpu = ga.zeros(len(self.geometry.mesh.triangles), dtype=ga.vec.float3) if self.enable3d: self.image1_gpu = ga.zeros(self.npixels, dtype=ga.vec.float3) self.image2_gpu = ga.zeros(self.npixels, dtype=ga.vec.float3) else: self.image_gpu = ga.zeros(self.npixels, dtype=ga.vec.float3) self.source_position = self.point self.nimages = 0 self.nlookup_calls = 0 self.max_steps = 10 def clear_xyz_lookup(self): self.xyz_lookup1_gpu.fill(ga.vec.make_float3(0.0,0.0,0.0)) self.xyz_lookup2_gpu.fill(ga.vec.make_float3(0.0,0.0,0.0)) self.nlookup_calls = 0 def update_xyz_lookup(self, source_position): for wavelength, rgb_tuple in \ zip([685.0, 545.0, 445.0],[(1,0,0),(0,1,0),(0,0,1)]): for i in range(self.xyz_lookup1_gpu.size//(self.npixels)+1): self.hybrid_funcs.update_xyz_lookup(np.int32(self.npixels), np.int32(self.xyz_lookup1_gpu.size), np.int32(i*self.npixels), ga.vec.make_float3(*source_position), self.rng_states_gpu, np.float32(wavelength), ga.vec.make_float3(*rgb_tuple), self.xyz_lookup1_gpu, self.xyz_lookup2_gpu, np.int32(self.max_steps), self.gpu_geometry.gpudata, block=(self.nblocks,1,1), grid=(self.npixels//self.nblocks+1,1)) self.nlookup_calls += 1 def clear_image(self): if self.enable3d: self.image1_gpu.fill(ga.vec.make_float3(0.0,0.0,0.0)) self.image2_gpu.fill(ga.vec.make_float3(0.0,0.0,0.0)) else: self.image_gpu.fill(ga.vec.make_float3(0.0,0.0,0.0)) self.nimages = 0 def update_image_from_rays(self, image, rays): for wavelength, rgb_tuple in \ zip([685.0, 545.0, 445.0],[(1,0,0),(0,1,0),(0,0,1)]): self.hybrid_funcs.update_xyz_image(np.int32(rays.pos.size), self.rng_states_gpu, rays.pos, rays.dir, np.float32(wavelength), ga.vec.make_float3(*rgb_tuple), self.xyz_lookup1_gpu, self.xyz_lookup2_gpu, image, np.int32(self.nlookup_calls), np.int32(self.max_steps), self.gpu_geometry.gpudata, block=(self.nblocks,1,1), grid=(rays.pos.size//self.nblocks+1,1)) def update_image(self): if self.enable3d: self.update_image_from_rays(self.image1_gpu, self.rays1) self.update_image_from_rays(self.image2_gpu, self.rays2) else: self.update_image_from_rays(self.image_gpu, self.rays) self.nimages += 1 def process_image(self): if self.enable3d: self.hybrid_funcs.process_image(np.int32(self.pixels1_gpu.size), self.image1_gpu, self.pixels1_gpu, np.int32(self.nimages), block=(self.nblocks,1,1), grid=((self.pixels1_gpu.size)//self.nblocks+1,1)) self.hybrid_funcs.process_image(np.int32(self.pixels2_gpu.size), self.image2_gpu, self.pixels2_gpu, np.int32(self.nimages), block=(self.nblocks,1,1), grid=((self.pixels2_gpu.size)//self.nblocks+1,1)) else: self.hybrid_funcs.process_image(np.int32(self.pixels_gpu.size), self.image_gpu, self.pixels_gpu, np.int32(self.nimages), block=(self.nblocks,1,1), grid=((self.pixels_gpu.size)//self.nblocks+1,1)) def screenshot(self, dir='', start=0): root, ext = 'screenshot', 'png' for i in count(start): filename = os.path.join(dir, '.'.join([root + str(i).zfill(5), ext])) if not os.path.exists(filename): break pygame.image.save(self.screen, filename) print 'image saved to %s' % filename def rotate(self, phi, n): if self.enable3d: self.rays1.rotate(phi, n) self.rays2.rotate(phi, n) self.scope_rays.rotate(phi, n) self.point1 = rotate(self.point1, phi, n) self.point2 = rotate(self.point2, phi, n) else: self.rays.rotate(phi, n) self.point = rotate(self.point, phi, n) self.axis1 = rotate(self.axis1, phi, n) self.axis2 = rotate(self.axis2, phi, n) if self.render: self.clear_image() self.update() def rotate_around_point(self, phi, n, point, redraw=True): self.axis1 = rotate(self.axis1, phi, n) self.axis2 = rotate(self.axis2, phi, n) if self.enable3d: self.rays1.rotate_around_point(phi, n, point) self.rays2.rotate_around_point(phi, n, point) self.scope_rays.rotate_around_point(phi, n, point) else: self.rays.rotate_around_point(phi, n, point) if redraw: if self.render: self.clear_image() self.update() def translate(self, v, redraw=True): self.point += v if self.enable3d: self.rays1.translate(v) self.rays2.translate(v) self.scope_rays.translate(v) self.point1 += v self.point2 += v else: self.rays.translate(v) if redraw: if self.render: self.clear_image() self.update() def update_pixels(self, gpu_geometry=None, keep_last_render=False): if gpu_geometry is None: gpu_geometry = self.gpu_geometry if self.render: while self.nlookup_calls < 10: self.update_xyz_lookup(self.source_position) self.update_image() self.process_image() else: if self.enable3d: self.rays1.render(gpu_geometry, self.pixels1_gpu, self.alpha_depth, keep_last_render) self.rays2.render(gpu_geometry, self.pixels2_gpu, self.alpha_depth, keep_last_render) else: self.rays.render(gpu_geometry, self.pixels_gpu, self.alpha_depth, keep_last_render) def update_viewing_angle(self): if self.enable3d: self.gpu_funcs.distance_to_mesh(np.int32(self.scope_rays.pos.size), self.scope_rays.pos, self.scope_rays.dir, self.gpu_geometry.gpudata, self.distances_gpu, block=(self.nblocks,1,1), grid=(self.scope_rays.pos.size//self.nblocks,1)) baseline = ga.min(self.distances_gpu).get().item() if baseline < 1e9: d1 = self.point1 - self.point v1 = d1/np.linalg.norm(d1) v1 *= baseline/60 - np.linalg.norm(d1) self.rays1.translate(v1) self.point1 += v1 d2 = self.point2 - self.point v2 = d2/np.linalg.norm(d2) v2 *= baseline/60 - np.linalg.norm(d2) self.rays2.translate(v2) self.point2 += v2 direction = np.cross(self.axis1,self.axis2) direction /= np.linalg.norm(direction) direction1 = self.point + direction*baseline - self.point1 direction1 /= np.linalg.norm(direction1) new_viewing_angle = np.arccos(direction1.dot(direction)) phi = new_viewing_angle - self.viewing_angle self.rays1.rotate_around_point(phi, self.axis1, self.point1) self.rays2.rotate_around_point(-phi, self.axis1, self.point2) self.viewing_angle = new_viewing_angle def update(self): if self.enable3d: self.update_viewing_angle() n = len(self.gpu_geometries) for i, gpu_geometry in enumerate(self.gpu_geometries): if i == 0: self.update_pixels(gpu_geometry) else: self.update_pixels(gpu_geometry, keep_last_render=True) if self.enable3d: pixels1 = self.pixels1_gpu.get() pixels2 = self.pixels2_gpu.get() if self.green_magenta: pixels = (pixels1 & 0x00ff00) | (pixels2 & 0xff00ff) else: pixels = (pixels1 & 0xff0000) | (pixels2 & 0x00ffff) alpha = ((0xff & (pixels1 >> 24)) + (0xff & (pixels2 >> 24)))/2 pixels |= (alpha << 24) else: pixels = self.pixels_gpu.get() pygame.surfarray.blit_array(self.screen, pixels.reshape(self.size)) if self.doom_mode: self.screen.blit(self.doom_hud, self.doom_rect) self.window.fill(0) self.window.blit(self.screen, (0,0)) pygame.display.flip() if self.movie: self.screenshot(self.movie_dir, self.movie_index) self.movie_index += 1 def loadlayer(self, layer): if layer is None: self.gpu_geometries = [self.gpu_geometry] else: try: gpu_geometry = self.bvh_layers[layer] except KeyError: geometry = build(bvh_mesh(self.geometry, layer), 8) gpu_geometry = gpu.GPUGeometry(geometry) self.bvh_layers[layer] = gpu_geometry self.gpu_geometries = [self.gpu_geometry, gpu_geometry] self.update() def process_event(self, event): if event.type == MOUSEBUTTONDOWN: if event.button == 4: v = self.scale*np.cross(self.axis1,self.axis2)/10.0 self.translate(v) elif event.button == 5: v = -self.scale*np.cross(self.axis1,self.axis2)/10.0 self.translate(v) elif event.button == 1: mouse_position = pygame.mouse.get_rel() self.clicked = True elif event.type == MOUSEBUTTONUP: if event.button == 1: self.clicked = False elif event.type == MOUSEMOTION and self.clicked: movement = np.array(pygame.mouse.get_rel()) if (movement == 0).all(): return length = np.linalg.norm(movement) mouse_direction = movement[0]*self.axis2 - movement[1]*self.axis1 mouse_direction /= np.linalg.norm(mouse_direction) if pygame.key.get_mods() & (KMOD_LSHIFT | KMOD_RSHIFT): v = -mouse_direction*self.scale*length/float(self.width) self.translate(v) else: phi = np.float32(2*np.pi*length/float(self.width)) n = rotate(mouse_direction, np.pi/2, np.cross(self.axis1,self.axis2)) if pygame.key.get_mods() & KMOD_LCTRL: self.rotate_around_point(phi, n, self.point) else: self.rotate(phi, n) elif event.type == KEYDOWN: if event.key == K_a: v = self.scale*self.axis2/10.0 self.translate(v) elif event.key == K_d: v = -self.scale*self.axis2/10.0 self.translate(v) elif event.key == K_w: v = self.scale*np.cross(self.axis1,self.axis2)/10.0 self.translate(v) elif event.key == K_s: v = -self.scale*np.cross(self.axis1,self.axis2)/10.0 self.translate(v) elif event.key == K_F6: self.clear_xyz_lookup() self.clear_image() self.source_position = self.point elif event.key == K_p: for i in range(100): self.update_xyz_lookup(self.point) self.source_position = self.point elif event.key == K_F11: pygame.display.toggle_fullscreen() elif event.key == K_ESCAPE: self.done = True return elif event.key == K_EQUALS: self.alpha_depth += 1 self.update() elif event.key == K_MINUS: if self.alpha_depth > 0: self.alpha_depth -= 1 self.update() elif event.key == K_PAGEDOWN: if self.currentlayer is None: self.currentlayer = self.unique_bvh_layers[-1] else: if self.currentlayer > np.min(self.unique_bvh_layers): self.currentlayer -= 1 else: self.currentlayer = None self.loadlayer(self.currentlayer) elif event.key == K_PAGEUP: if self.currentlayer is None: self.currentlayer = self.unique_bvh_layers[0] else: if self.currentlayer < np.max(self.unique_bvh_layers): self.currentlayer += 1 else: self.currentlayer = None self.loadlayer(self.currentlayer) elif event.key == K_F12: self.screenshot() elif event.key == K_F5: if not hasattr(self, 'rng_states_gpu'): self.initialize_render() self.render = not self.render self.clear_image() self.update() elif event.key == K_m: if self.movie: encode_movie(self.movie_dir) self.movie_dir = None self.movie = False else: self.movie_index = 0 self.movie_dir = tempfile.mkdtemp() self.movie = True elif event.type == pygame.SYSWMEVENT and self.spnav: # Space Navigator controls spnav_event = self.spnav_module.spnav_x11_event(event.event) if spnav_event is None: return if spnav_event.type == self.spnav_module.SPNAV_EVENT_MOTION: if pygame.key.get_mods() & (KMOD_LSHIFT | KMOD_RSHIFT): accelerate_factor = 2.0 else: accelerate_factor = 1.0 #print 'raw:', spnav_event v1 = self.axis2 v2 = self.axis1 v3 = -np.cross(self.axis1,self.axis2) v = -v1*spnav_event.motion.x + v2*spnav_event.motion.y \ + v3*spnav_event.motion.z v *= self.scale / 5000.0 * accelerate_factor #print 'translate:', v self.translate(v, redraw=False) axis = v1*spnav_event.motion.rx + -v2*spnav_event.motion.ry \ + -v3*spnav_event.motion.rz # Zero all but non-max values #axis[~(np.abs(axis) == np.abs(axis).max())] = 0 #axis[np.abs(axis) < 15] = 0 if (axis != 0).any(): axis = axis.astype(float) length = np.linalg.norm(axis) angle = length * 0.0001 * accelerate_factor axis /= length #print 'rotate:', angle, axis self.rotate_around_point(angle, axis, self.point, redraw=False) if self.render: self.clear_image() self.update() pygame.event.clear(pygame.SYSWMEVENT) elif spnav_event.type == self.spnav_module.SPNAV_EVENT_BUTTON: if spnav_event.button.bnum == 0 and spnav_event.button.press: if not hasattr(self, 'rng_states_gpu'): self.initialize_render() self.render = not self.render self.clear_image() self.update() pygame.event.clear(pygame.SYSWMEVENT) elif event.type == pygame.QUIT: self.done = True return def run(self): self.init_gpu() if self.spnav: try: wm_info = pygame.display.get_wm_info() self.spnav_module.spnav_x11_open(wm_info['display'], wm_info['window']) pygame.event.set_allowed(pygame.SYSWMEVENT) print 'Space Navigator support enabled.' except: self.spnav = False self.update() self.done = False self.clicked = False while not self.done: self.clock.tick(20) if self.render and not self.clicked and not pygame.event.peek(KEYDOWN): self.update() # Grab only last SYSWMEVENT (SpaceNav!) to avoid lagged controls for event in pygame.event.get(pygame.SYSWMEVENT)[-1:] + pygame.event.get(): self.process_event(event) if self.movie: encode_movie(self.movie_dir) pygame.display.quit() if self.spnav: self.spnav_module.spnav_close() del self.gpu_instance class EventViewer(Camera): def __init__(self, geometry, filename, **kwargs): Camera.__init__(self, geometry, **kwargs) self.rr = RootReader(filename) def render_particle_track(self): marker = Solid(make.cube(0.1), vacuum, vacuum) geometry = Geometry() for pos in self.ev.photons_beg.pos[::100]: geometry.add_solid(marker, displacement=pos) geometry.build(bits=11) gpu_geometry = gpu.GPUGeometry(geometry) self.gpu_geometries = [self.gpu_geometry, gpu_geometry] self.update() def color_hit_pmts(self): self.gpu_geometry.reset_colors() hit = self.ev.channels.hit t = self.ev.channels.t q = self.ev.channels.q # Important: Compute range only with HIT channels solid_colors = map_to_color(q, range=(q[hit].min(),q[hit].max())) self.gpu_geometry.color_solids(hit, solid_colors) self.update() def process_event(self, event): if event.type == KEYDOWN: if event.key == K_PAGEUP: try: self.ev = self.rr.next() except StopIteration: pass else: self.color_hit_pmts() self.render_particle_track() return elif event.key == K_PAGEDOWN: try: self.ev = self.rr.prev() except StopIteration: pass else: self.color_hit_pmts() self.render_particle_track() return Camera.process_event(self, event) def view(obj, size, **camera_kwargs): geometry = build(obj, 8) camera = Camera(geometry, size, **camera_kwargs) camera.start() camera.join() if __name__ == '__main__': import optparse import inspect import chroma.solids import chroma.detectors import chroma.scenes from chroma.stl import mesh_from_stl parser = optparse.OptionParser('%prog filename.stl') parser.add_option('-b', '--bits', type='int', dest='bits', help='bits for z-ordering space axes', default=10) parser.add_option('-r', '--resolution', dest='resolution', help='specify resolution', default='1024,576') parser.add_option('--3d', action='store_true', dest='enable3d', help='enable 3d', default=False) parser.add_option('--green', action='store_true', dest='green_magenta', help='3d with green and magenta lenses', default=False) parser.add_option('-i', dest='io_file', default=None) options, args = parser.parse_args() if len(args) < 1: sys.exit(parser.format_help()) size = [int(s) for s in options.resolution.split(',')] if os.path.exists(args[0]): root, ext = os.path.splitext(os.path.split(args[0])[1]) if ext.lower() in ('.stl', '.bz2'): obj = mesh_from_stl(args[0]) else: members = dict(inspect.getmembers(chroma.detectors) + inspect.getmembers(chroma.solids) + inspect.getmembers(chroma.scenes)) buildable_lookup = {} for member in members.values(): if inspect.isfunction(member) and \ hasattr(member, 'buildable') and member.buildable == True: buildable_lookup[member.identifier] = member if args[0] in buildable_lookup: obj = buildable_lookup[args[0]] else: raise Exception("can't find object %s" % args[0]) geometry = build(obj, options.bits) if options.io_file is not None: camera = EventViewer(geometry, options.io_file, size=size, enable3d=options.enable3d, green_magenta=options.green_magenta) else: camera = Camera(geometry, size, enable3d=options.enable3d, green_magenta=options.green_magenta) camera.start() camera.join()