diff options
| -rwxr-xr-x | camera.py | 117 | ||||
| -rw-r--r-- | color/chromaticity.py | 12 | ||||
| -rw-r--r-- | detectors/__init__.py | 9 | ||||
| -rw-r--r-- | detectors/lbne.py | 22 | ||||
| -rw-r--r-- | detectors/sno.py | 12 | ||||
| -rw-r--r-- | fileio/root.py | 2 | ||||
| -rw-r--r-- | generator/g4gen.py | 2 | ||||
| -rw-r--r-- | generator/vertex.py | 4 | ||||
| -rw-r--r-- | geometry.py | 6 | ||||
| -rw-r--r-- | gpu.py | 15 | ||||
| -rw-r--r-- | make.py | 20 | ||||
| -rw-r--r-- | models/__init__.py | 3 | ||||
| -rw-r--r-- | optics.py | 2 | ||||
| -rw-r--r-- | sample.py | 2 | ||||
| -rw-r--r-- | scenes/checkerboard.py | 16 | ||||
| -rw-r--r-- | solids/__init__.py | 21 | ||||
| -rw-r--r-- | solids/pmts.py | 14 | ||||
| -rw-r--r-- | src/__init__.py | 9 | ||||
| -rw-r--r-- | stl.py | 2 | ||||
| -rwxr-xr-x | view.py | 375 |
20 files changed, 179 insertions, 486 deletions
@@ -4,12 +4,15 @@ from itertools import product, count from threading import Thread, Lock import os import sys -from color import map_wavelength, map_to_color +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 gpu import GPU, CUDAFuncs -from tools import timeit -from transform import rotate +from chroma.gpu import GPU, CUDAFuncs +from chroma.tools import timeit +from chroma.transform import rotate +from chroma.optics import vacuum, lambertian_surface import pygame from pygame.locals import * @@ -22,6 +25,65 @@ from subprocess import call import shutil import tempfile +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(): @@ -59,6 +121,10 @@ class Camera(Thread): self.device_id = device_id self.size = size + self.unique_bvh_layers = np.unique(self.geometry.layers) + self.currentlayer = None + self.bvh_layers = {self.currentlayer : self.geometry} + try: import spnav as spnav_module self.spnav_module = spnav_module @@ -231,6 +297,16 @@ class Camera(Thread): self.screenshot(self.movie_dir, self.movie_index) self.movie_index += 1 + def loadlayer(self, layer): + try: + layergeometry = self.bvh_layers[layer] + except KeyError: + layergeometry = build(bvh_mesh(self.geometry, layer),8) + self.bvh_layers[layer] = layergeometry + + self.gpu.load_geometry(layergeometry) + self.update() + def process_event(self, event): if event.type == MOUSEBUTTONDOWN: if event.button == 4: @@ -310,6 +386,28 @@ class Camera(Thread): self.done = True return + 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() @@ -483,11 +581,10 @@ if __name__ == '__main__': import optparse import inspect - import solids - import detectors - import scenes - from stl import mesh_from_stl - from view import build + 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', @@ -511,7 +608,7 @@ if __name__ == '__main__': obj = mesh_from_stl(args[0]) else: - members = dict(inspect.getmembers(detectors) + inspect.getmembers(solids) + inspect.getmembers(scenes)) + members = dict(inspect.getmembers(chroma.detectors) + inspect.getmembers(chroma.solids) + inspect.getmembers(chroma.scenes)) buildable_lookup = {} for member in members.values(): diff --git a/color/chromaticity.py b/color/chromaticity.py index 859cc14..0076957 100644 --- a/color/chromaticity.py +++ b/color/chromaticity.py @@ -1,14 +1,8 @@ import numpy as np -import os -import sys +from os.path import realpath, dirname +from chroma.tools import read_csv -dir = os.path.split(os.path.realpath(__file__))[0] - -sys.path.append(dir + '/..') - -from tools import read_csv - -color_map = read_csv(dir + '/ciexyz64_1.csv') +color_map = read_csv(dirname(realpath(__file__)) + '/ciexyz64_1.csv') def map_wavelength(wavelength): r = np.interp(wavelength, color_map[:,0], color_map[:,1]) diff --git a/detectors/__init__.py b/detectors/__init__.py index 8550751..f812058 100644 --- a/detectors/__init__.py +++ b/detectors/__init__.py @@ -1,15 +1,8 @@ from lbne import build_lbne from sno import build_sno as build_sno_detector - -import os -import sys +from chroma.camera import buildable import inspect -dir = os.path.split(os.path.realpath(__file__))[0] -sys.path.append(dir + '/..') - -from view import buildable - # from LBNE document #3951 radius = 63.30/2 height = 76.60 diff --git a/detectors/lbne.py b/detectors/lbne.py index 8942cb2..ca13c2d 100644 --- a/detectors/lbne.py +++ b/detectors/lbne.py @@ -1,27 +1,21 @@ -import os -import sys import numpy as np - -dir = os.path.split(os.path.realpath(__file__))[0] -sys.path.append(dir + '/..') - -from geometry import * -from solids import build_12inch_pmt, build_12inch_pmt_with_lc, build_12inch_pmt_shell -from optics import * -from transform import rotate, make_rotation_matrix +from chroma.geometry import Solid, Geometry +import chroma.solids as solids +from chroma.optics import vacuum, water, water_wcsim, black_surface +from chroma.transform import rotate, make_rotation_matrix from itertools import product -from make import cylinder, segmented_cylinder +import chroma.make as make def build_lbne(radius, height, nstrings, pmts_per_string, endcap_spacing, physical_model=True): if physical_model: - pmt = build_12inch_pmt_with_lc(outer_material=water_wcsim) + pmt = solids.build_12inch_pmt_with_lc(outer_material=water_wcsim) else: - pmt = build_12inch_pmt_shell(outer_material=water_wcsim) + pmt = solids.build_12inch_pmt_shell(outer_material=water_wcsim) lbne = Geometry() # outer cylinder - cylinder_mesh = segmented_cylinder(radius, height+height/(pmts_per_string-1), nsteps=16*nstrings, nsegments=1200) + cylinder_mesh = make.segmented_cylinder(radius, height+height/(pmts_per_string-1), nsteps=16*nstrings, nsegments=1200) cylinder_mesh.vertices = rotate(cylinder_mesh.vertices, np.pi/2, (-1,0,0)) lbne.add_solid(Solid(cylinder_mesh, water_wcsim, vacuum, black_surface, 0xff0000ff)) diff --git a/detectors/sno.py b/detectors/sno.py index 3e748b4..d0334d3 100644 --- a/detectors/sno.py +++ b/detectors/sno.py @@ -1,12 +1,12 @@ import numpy as np import numpy.linalg import math -from make import rotate_extrude -from stl import mesh_from_stl -from geometry import * -from optics import * -from solids import build_8inch_pmt_with_lc -from transform import rotate, make_rotation_matrix +from chroma.make import rotate_extrude +from chroma.stl import mesh_from_stl +from chroma.geometry import * +from chroma.optics import * +from chroma.solids import build_8inch_pmt_with_lc +from chroma.transform import rotate, make_rotation_matrix import os def sno_vessel(sphere_radius, neck_radius, neck_top, nsteps=40): diff --git a/fileio/root.py b/fileio/root.py index 77c23ff..ae9e741 100644 --- a/fileio/root.py +++ b/fileio/root.py @@ -4,7 +4,7 @@ import os.path ROOT.gROOT.ProcessLine('.L '+os.path.join(os.path.dirname(__file__), 'root.C+g')) import ROOT -import event +import chroma.event as event class RootWriter(object): def __init__(self, filename): diff --git a/generator/g4gen.py b/generator/g4gen.py index cec53f7..b43ce92 100644 --- a/generator/g4gen.py +++ b/generator/g4gen.py @@ -4,7 +4,7 @@ import g4py.NISTmaterials import g4py.ParticleGun import pyublas import numpy as np -import event +import chroma.event as event try: import G4chroma diff --git a/generator/vertex.py b/generator/vertex.py index 1fa61cb..3998e6e 100644 --- a/generator/vertex.py +++ b/generator/vertex.py @@ -3,8 +3,8 @@ import numpy as np from math import sin, cos, sqrt import copy -import pi0 -from event import Event, Subtrack +import chroma.pi0 as pi0 +from chroma.event import Event, Subtrack # generator parts for use with gun() diff --git a/geometry.py b/geometry.py index 856ded3..6ea8994 100644 --- a/geometry.py +++ b/geometry.py @@ -1,8 +1,8 @@ import sys import os import numpy as np -from itertoolset import * -from tools import timeit, profile_if_possible +from chroma.itertoolset import * +from chroma.tools import timeit, profile_if_possible from hashlib import md5 import cPickle as pickle import gzip @@ -274,7 +274,7 @@ class Geometry(object): try: self.surface_index[self.surface_index == surface_lookup[None]] = -1 - except ValueError: + except KeyError: pass checksum = md5(str(bits)) @@ -7,11 +7,12 @@ import pycuda.driver as cuda from pycuda import gpuarray from copy import copy from itertools import izip -from tools import timeit +from chroma.tools import timeit +from os.path import dirname -import src -from geometry import standard_wavelengths -from color import map_to_color +import chroma.src +from chroma.geometry import standard_wavelengths +from chroma.color import map_to_color import sys import event @@ -97,15 +98,15 @@ class GPU(object): self.context.set_cache_config(cuda.func_cache.PREFER_L1) - cuda_options = ['-I' + src.dir, '--use_fast_math', '--ptxas-options=-v'] + cuda_options = ['-I' + dirname(chroma.src.__file__), '--use_fast_math', '--ptxas-options=-v'] - self.module = SourceModule(src.kernel, options=cuda_options, no_extern_c=True) + self.module = SourceModule(chroma.src.kernel, options=cuda_options, no_extern_c=True) self.geo_funcs = CUDAFuncs(self.module, ['set_wavelength_range', 'set_material', 'set_surface', 'set_global_mesh_variables', 'color_solids']) self.prop_funcs = CUDAFuncs(self.module, ['init_rng', 'propagate', 'swap']) self.nthread_per_block = 64 self.max_blocks = 1024 - self.daq_module = SourceModule(src.daq, options=cuda_options, no_extern_c=True) + self.daq_module = SourceModule(chroma.src.daq, options=cuda_options, no_extern_c=True) self.daq_funcs = CUDAFuncs(self.daq_module, ['reset_earliest_time_int', 'run_daq', 'convert_sortable_int_to_float']) @@ -1,12 +1,12 @@ import numpy as np -from geometry import Mesh -from transform import rotate -from itertoolset import * +from chroma.geometry import Mesh +from chroma.transform import rotate +from chroma.itertoolset import * def mesh_grid(grid): return np.vstack(zip(grid[:-1].flatten(),grid[1:].flatten(),np.roll(grid[1:],-1,1).flatten()) + zip(grid[:-1].flatten(),np.roll(grid[1:],-1,1).flatten(),np.roll(grid[:-1],-1,1).flatten())) -def linear_extrude(x1, y1, height, x2=None, y2=None): +def linear_extrude(x1, y1, height, x2=None, y2=None, center=None): """ Return the solid mesh formed by linearly extruding the polygon formed by the x and y points `x1` and `y1` by a distance `height`. If `x2` and `y2` @@ -41,6 +41,9 @@ def linear_extrude(x1, y1, height, x2=None, y2=None): vertices = np.fromiter(flatten(roundrobin(*vertex_iterators)), float) vertices = vertices.reshape((len(vertices)//3,3)) + if center is not None: + vertices += center + triangles = mesh_grid(np.arange(len(vertices)).reshape((len(x1),len(vertices)//len(x1))).transpose()[::-1]) return Mesh(vertices, triangles, remove_duplicate_vertices=True) @@ -65,8 +68,15 @@ def rotate_extrude(x, y, nsteps=64): return Mesh(vertices, triangles, remove_duplicate_vertices=True) -def cube(size=1): +def box(dx, dy, dz, center=(0,0,0)): + "Return a box with linear dimensions `dx`, `dy`, and `dz`." + return linear_extrude([-dx/2.0,dx/2.0,dx/2.0,-dx/2.0],[-dy/2.0,-dy/2.0,dy/2.0,dy/2.0],height=dz,center=center) + +def cube(size=1, height=None): "Return a cube mesh whose sides have length `size`." + if height is None: + height = size + return linear_extrude([-size/2.0,size/2.0,size/2.0,-size/2.0],[-size/2.0,-size/2.0,size/2.0,size/2.0], height=size) def cylinder(radius=1, height=2, radius2=None, nsteps=64): diff --git a/models/__init__.py b/models/__init__.py deleted file mode 100644 index f369c9b..0000000 --- a/models/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -import os - -dir = os.path.split(os.path.realpath(__file__))[0] @@ -1,5 +1,5 @@ import numpy as np -from geometry import Material, Surface +from chroma.geometry import Material, Surface vacuum = Material('vacuum') vacuum.set('refractive_index', 1.0) @@ -1,5 +1,5 @@ import numpy as np -from transform import rotate +from chroma.transform import rotate def uniform_sphere(size=None, dtype=np.double): """ diff --git a/scenes/checkerboard.py b/scenes/checkerboard.py index d2a9ec8..7e4f755 100644 --- a/scenes/checkerboard.py +++ b/scenes/checkerboard.py @@ -1,15 +1,9 @@ -import os -import sys - -dir = os.path.split(os.path.realpath(__file__))[0] -sys.path.append(dir + '/..') - import numpy as np -from itertoolset import * -from geometry import Mesh, Solid, Geometry -from optics import * -from make import sphere -from view import buildable +from chroma.itertoolset import * +from chroma.geometry import Mesh, Solid, Geometry +from chroma.optics import * +from chroma.make import sphere +from chroma.camera import buildable @buildable('checkerboard_scene') def build_checkerboard_scene(checkers_per_side=10, squares_per_checker=50): diff --git a/solids/__init__.py b/solids/__init__.py index a24ce9c..421b66f 100644 --- a/solids/__init__.py +++ b/solids/__init__.py @@ -1,23 +1,16 @@ import numpy as np - from pmts import build_pmt, build_light_collector, build_light_collector_from_file, build_pmt_shell - -import os -import sys - -dir = os.path.split(os.path.realpath(__file__))[0] -sys.path.append(dir + '/..') - -from optics import * -from view import buildable +from chroma.optics import * +from chroma.camera import buildable +from os.path import dirname @buildable('12inch_pmt') def build_12inch_pmt(outer_material=water, nsteps=16): - return build_pmt(dir + '/hamamatsu_12inch.txt', 0.003, outer_material, nsteps) + return build_pmt(dirname(__file__) + '/hamamatsu_12inch.txt', 0.003, outer_material, nsteps) @buildable('12inch_pmt_shell') def build_12inch_pmt_shell(outer_material=water, nsteps=16): - return build_pmt_shell(dir + '/hamamatsu_12inch.txt') + return build_pmt_shell(dirname(__file__) + '/hamamatsu_12inch.txt') # from Jelena Maricic lc_12inch_a = 0.16597 @@ -39,10 +32,10 @@ def build_12inch_pmt_with_lc_hd(outer_material=water, nsteps=128): @buildable('8inch_pmt') def build_8inch_pmt(outer_material=water, nsteps=24): - return build_pmt(dir + '/sno_pmt_reduced.txt', 0.003, outer_material, nsteps) + return build_pmt(dirname(__file__) + '/sno_pmt_reduced.txt', 0.003, outer_material, nsteps) @buildable('8inch_pmt_with_lc') def build_8inch_pmt_with_lc(outer_material=water, nsteps=24): pmt = build_8inch_pmt(outer_material, nsteps) - lc = build_light_collector_from_file(dir + '/sno_cone.txt', outer_material, nsteps) + lc = build_light_collector_from_file(dirname(__file__) + '/sno_cone.txt', outer_material, nsteps) return pmt + lc diff --git a/solids/pmts.py b/solids/pmts.py index fcbc8fc..43aed3f 100644 --- a/solids/pmts.py +++ b/solids/pmts.py @@ -1,14 +1,8 @@ -import os -import sys - -dir = os.path.split(os.path.realpath(__file__))[0] -sys.path.append(dir + '/..') - import numpy as np -from geometry import Solid -from make import rotate_extrude -from optics import * -from tools import read_csv, offset +from chroma.geometry import Solid +from chroma.make import rotate_extrude +from chroma.optics import * +from chroma.tools import read_csv, offset def get_lc_profile(radii, a, b, d, rmin, rmax): c = -b*np.sqrt(1 - (rmin-d)**2/a**2) diff --git a/src/__init__.py b/src/__init__.py index 865d612..f2f5215 100644 --- a/src/__init__.py +++ b/src/__init__.py @@ -1,6 +1,7 @@ -import os +from os.path import dirname -dir = os.path.split(os.path.realpath(__file__))[0] +with open(dirname(__file__) + '/kernel.cu') as f: + kernel = f.read() -kernel = open(dir + '/kernel.cu').read() -daq = open(dir + '/daq.cu').read() +with open(dirname(__file__) + '/daq.cu') as f: + daq = f.read() @@ -1,7 +1,7 @@ import numpy as np import string import struct -from geometry import Mesh +from chroma.geometry import Mesh import bz2 def mesh_from_stl(filename): diff --git a/view.py b/view.py deleted file mode 100755 index 7ae6a27..0000000 --- a/view.py +++ /dev/null @@ -1,375 +0,0 @@ -#!/usr/bin/env python -import os -import sys -import time -import numpy as np -import inspect - -import pygame -from pygame.locals import * - -import src -from camera import get_rays -from geometry import Mesh, Solid, Geometry -from transform import rotate -from optics import * -from gpu import * - -#from pycuda import autoinit -from pycuda.compiler import SourceModule -from pycuda import gpuarray -import pycuda.driver as cuda - -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 screenshot(screen, name='', dir='', index=0): - """Take a screenshot of `screen`.""" - if name == '': - root, ext = 'screenshot', 'png' - else: - root, ext = name, 'png' - - i = index - - filename = os.path.join(dir,'.'.join([root + str(i).zfill(4), ext])) - while os.path.exists(filename): - filename = os.path.join(dir,'.'.join([root + str(i).zfill(4), ext])) - i += 1 - - pygame.image.save(screen, filename) - print 'image saved to %s' % filename - -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 box(lower_bound, upper_bound): - """ - Return a mesh of the box defined by the opposing corners `lower_bound` - and `upper_bound`. - """ - dx, dy, dz = upper_bound - lower_bound - - vertices = np.array([lower_bound, - lower_bound + (dx,0,0), - lower_bound + (dx,dy,0), - lower_bound + (0,dy,0), - lower_bound + (0,0,dz), - lower_bound + (dx,0,dz), - lower_bound + (dx,dy,dz), - lower_bound + (0,dy,dz)]) - - triangles = np.empty((12,3), dtype=np.int32) - # bottom - triangles[0] = (1,3,2) - triangles[1] = (1,4,3) - # top - triangles[2] = (5,7,6) - triangles[3] = (5,8,7) - # left - triangles[4] = (5,1,2) - triangles[5] = (5,2,6) - # right - triangles[6] = (3,4,8) - triangles[7] = (3,8,7) - # front - triangles[8] = (2,3,7) - triangles[9] = (2,7,6) - # back - triangles[10] = (1,5,8) - triangles[11] = (1,8,4) - - triangles -= 1 - - return Mesh(vertices, triangles) - -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) - - mesh = box(lower_bounds[0], upper_bounds[0]) - - for lower_bound, upper_bound in zip(lower_bounds, upper_bounds)[1:]: - mesh += box(lower_bound, upper_bound) - - return mesh - -def view(viewable, size=(800,600), name='', bits=8, load_bvh=False): - """ - Render `viewable` in a pygame window. - - Movement: - - zoom: scroll the mouse wheel - - rotate: click and drag the mouse - - move: shift+click and drag the mouse - """ - - gpu = GPU() - - geometry = build(viewable, bits) - - if load_bvh: - print 'loading bvh...' - bvhg = [] - - bvhg.append(geometry) - - for layer in sorted(np.unique(geometry.layers)): - print 'building layer %i' % layer - bvhg.append(build(bvh_mesh(geometry, layer), bits)) - - lower_bound, upper_bound = geometry.mesh.get_bounds() - - scale = np.linalg.norm(upper_bound-lower_bound) - - #from pycuda import autoinit - - #print 'device %s' % autoinit.device.name() - - module = gpu.module - - #module = SourceModule(src.kernel, options=['-I' + src.dir], no_extern_c=True)#, cache_dir=False) - #geometry.load(module) - gpu.load_geometry(geometry) - cuda_raytrace = module.get_function('ray_trace') - cuda_rotate = module.get_function('rotate') - cuda_translate = module.get_function('translate') - - pygame.init() - width, height = size - screen = pygame.display.set_mode(size) - pygame.display.set_caption(name) - - #camera = Camera(size) - - diagonal = np.linalg.norm(upper_bound-lower_bound) - - point = np.array([0, diagonal*1.75, (lower_bound[2]+upper_bound[2])/2]) - axis1 = np.array([1,0,0], dtype=np.double) - axis2 = np.array([0,0,1], dtype=np.double) - - #camera.position(point) - - origins, directions = get_rays(point, size) - - origins_gpu = gpuarray.to_gpu(origins.astype(np.float32).view(gpuarray.vec.float3)) - - directions_gpu = gpuarray.to_gpu(directions.astype(np.float32).view(gpuarray.vec.float3)) - - pixels_gpu = gpuarray.zeros(width*height, dtype=np.int32) - - nblocks = 64 - - def update(): - """Render the mesh and display to screen.""" - t0 = time.time() - cuda_raytrace(np.int32(pixels_gpu.size), origins_gpu, directions_gpu, pixels_gpu, block=(nblocks,1,1), grid=(pixels_gpu.size//nblocks+1,1)) - cuda.Context.synchronize() - elapsed = time.time() - t0 - - #print 'elapsed %f sec' % elapsed - - pygame.surfarray.blit_array(screen, pixels_gpu.get().reshape(size)) - pygame.display.flip() - - update() - - done = False - clicked = False - shift = False - - current_layer = 0 - - while not done: - for event in pygame.event.get(): - if event.type == MOUSEBUTTONDOWN: - if event.button == 4: - v = scale*np.cross(axis1,axis2)/10.0 - - cuda_translate(np.int32(pixels_gpu.size), origins_gpu, gpuarray.vec.make_float3(*v), block=(nblocks,1,1), grid=(pixels_gpu.size//nblocks+1,1)) - - point += v - - update() - - if event.button == 5: - v = -scale*np.cross(axis1,axis2)/10.0 - - cuda_translate(np.int32(pixels_gpu.size), origins_gpu, gpuarray.vec.make_float3(*v), block=(nblocks,1,1), grid=(pixels_gpu.size//nblocks+1,1)) - - point += v - - update() - - if event.button == 1: - clicked = True - mouse_position = pygame.mouse.get_rel() - - if event.type == MOUSEBUTTONUP: - if event.button == 1: - clicked = False - - if event.type == MOUSEMOTION and clicked: - movement = np.array(pygame.mouse.get_rel()) - - if (movement == 0).all(): - continue - - length = np.linalg.norm(movement) - - mouse_direction = movement[0]*axis1 + movement[1]*axis2 - mouse_direction /= np.linalg.norm(mouse_direction) - - if shift: - v = mouse_direction*scale*length/float(width) - - cuda_translate(np.int32(pixels_gpu.size), origins_gpu, gpuarray.vec.make_float3(*v), block=(nblocks,1,1), grid=(pixels_gpu.size//nblocks+1,1)) - - point += v - - update() - else: - phi = np.float32(2*np.pi*length/float(width)) - n = rotate(mouse_direction, np.pi/2, \ - -np.cross(axis1,axis2)) - - cuda_rotate(np.int32(pixels_gpu.size), origins_gpu, phi, gpuarray.vec.make_float3(*n), block=(nblocks,1,1), grid=(pixels_gpu.size//nblocks+1,1)) - - cuda_rotate(np.int32(pixels_gpu.size), directions_gpu, phi, gpuarray.vec.make_float3(*n), block=(nblocks,1,1), grid=(pixels_gpu.size//nblocks+1,1)) - - point = rotate(point, phi, n) - axis1 = rotate(axis1, phi, n) - axis2 = rotate(axis2, phi, n) - - update() - - if event.type == KEYDOWN: - if event.key == K_LSHIFT or event.key == K_RSHIFT: - shift = True - - if event.key == K_ESCAPE: - done = True - break - - if event.key == K_PAGEUP and load_bvh: - try: - if current_layer+1 >= len(bvhg): - raise IndexError - - geometry = bvhg[current_layer+1] - current_layer += 1 - - gpu.load_geometry(geometry) - #geometry.load(module, color=True) - update() - except IndexError: - print 'no further layers to view' - - if event.key == K_PAGEDOWN and load_bvh: - try: - if current_layer-1 < 0: - raise IndexError - - geometry = bvhg[current_layer-1] - current_layer -= 1 - - gpu.load_geometry(geometry) - #geometry.load(module, color=True) - update() - except IndexError: - print 'no further layers to view' - - if event.key == K_F12: - screenshot(screen, name) - - if event.type == KEYUP: - if event.key == K_LSHIFT or event.key == K_RSHIFT: - shift = False - - if event.type == pygame.QUIT: - done = True - break - - pygame.display.quit() - -if __name__ == '__main__': - import optparse - - from stl import mesh_from_stl - - import solids - import detectors - import scenes - - parser = optparse.OptionParser('%prog filename.stl') - parser.add_option('-b', '--bits', type='int', dest='bits', - help='bits for z-ordering space axes', default=8) - parser.add_option('-l', action='store_true', dest='load_bvh', - help='load bounding volumes', default=False) - parser.add_option('-r', '--resolution', dest='resolution', - help='specify resolution', default='800,600') - 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]): - head, tail = os.path.split(args[0]) - root, ext = os.path.splitext(tail) - - if ext.lower() in ('.stl', '.bz2'): - view(mesh_from_stl(args[0]), size, root, options.bits, options.load_bvh) - else: - members = dict(inspect.getmembers(detectors) + inspect.getmembers(solids) + inspect.getmembers(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: - view(buildable_lookup[args[0]], size, args[0], options.bits, options.load_bvh) - else: - sys.exit("couldn't find object %s" % args[0]) |