diff options
Diffstat (limited to 'photon_map.py')
| -rw-r--r-- | photon_map.py | 139 |
1 files changed, 118 insertions, 21 deletions
diff --git a/photon_map.py b/photon_map.py index cc07db0..4c6110e 100644 --- a/photon_map.py +++ b/photon_map.py @@ -1,8 +1,10 @@ +import sys +import time import numpy as np from solid import Solid from geometry import Geometry from mesh import mesh_from_stl -from sample import uniform_sphere +from sample import uniform_sphere, flashlight from pycuda import autoinit from pycuda.compiler import SourceModule from pycuda import gpuarray @@ -10,45 +12,82 @@ import pycuda.driver as cuda import src import models from materials import * +from transform import * +from itertoolset import roundrobin +from view import view +import matplotlib.pyplot as plt print 'device %s' % autoinit.device.name() +nphotons, nblocks = 1000000, 64 + +start_position = (5.0,0.0,0.0) + +gpu_kwargs = {'block': (nblocks,1,1), 'grid': (nphotons/nblocks+1,1)} + lens_mesh = mesh_from_stl(models.dir + '/lens.stl') -lens_solid = Solid(0, lens_mesh, material1=pmt_glass, material2=pmt_vacuum) +lens_mesh.vertices /= 1000 +lens_solid = Solid(0, lens_mesh, material1=acrylic_sno, material2=vacuum, color=0x00ff0000) + +box_mesh = mesh_from_stl(models.dir + '/box.stl') +box_mesh.vertices /= 1000 +box_solid = Solid(0, box_mesh, material1=vacuum, material2=vacuum, surface=black_surface, color=0x000000ff) film_mesh = mesh_from_stl(models.dir + '/film.stl') -film_solid = Solid(0, film_mesh, displacement=(-5,0,0), material1=h2o, material2=pmt_vacuum) +film_mesh.vertices /= 10 +film_solid = Solid(1, film_mesh, material1=vacuum, material2=vacuum, surface=black_surface, displacement=(-0.3,0,0), color=0x0000ff00) + +pmt_mesh1 = mesh_from_stl(models.dir + '/hamamatsu_12inch.stl') +pmt_mesh1.vertices /= 1000 +pmt_mesh2 = mesh_from_stl(models.dir + '/hamamatsu_12inch.stl') +pmt_mesh2.vertices /= 1000 +pmt_mesh2.vertices *= 0.95 +pmt_solid1 = Solid(3, pmt_mesh1, displacement=(2.0,0.0,0.0), material1=glass, material2=vacuum) +pmt_solid2 = Solid(4, pmt_mesh2, displacement=(2.0,0.0,0.0), material1=vacuum, material2=glass, surface=lambertian_surface) sphere_mesh = mesh_from_stl(models.dir + '/sphere.stl') -sphere_solid = Solid(0, sphere_mesh, displacement=(10,0,0), material1=h2o, material2=pmt_vacuum, surface=aluminum) +sphere_mesh.vertices *= 100 +sphere_solid = Solid(2, sphere_mesh, material1=vacuum, material2=vacuum, surface=black_surface) -geometry = Geometry() +geometry = Geometry()#[lens_solid1, lens_solid2, film_solid, sphere_solid]) geometry.add_solid(lens_solid) geometry.add_solid(film_solid) -geometry.add_solid(sphere_solid) -geometry.build(bits=8) +geometry.add_solid(box_solid) +#geometry.add_solid(pmt_solid1) +#geometry.add_solid(pmt_solid2) +#geometry.add_solid(sphere_solid) + +#from detectors import LBNE +#geometry = LBNE() + +texrefs = geometry.build(bits=8) + +#view(geometry) module = SourceModule(src.kernel, options=['-I' + src.dir], no_extern_c=True, cache_dir=False) geometry.load(module) -cuda_propagate = module.get_function('propagate') +propagate = module.get_function('propagate') +init_rng = module.get_function('init_rng') -nphotons = 100 +init_rng(np.int32(nphotons), np.int32(0), np.int32(0), **gpu_kwargs) -positions = np.tile((10,5,0), nphotons).reshape(nphotons,3) +positions = np.tile(start_position, nphotons).reshape(nphotons,3) positions_float3 = np.empty(positions.shape[0], dtype=gpuarray.vec.float3) positions_float3['x'] = positions[:,0] positions_float3['y'] = positions[:,1] positions_float3['z'] = positions[:,2] positions_gpu = cuda.to_device(positions_float3) -directions = uniform_sphere(nphotons) +directions = flashlight(np.pi/256, (-1,0,0), size=nphotons) +#directions = flashlight(np.pi/8, (0,-1,0), size=nphotons) +#directions = uniform_sphere(nphotons) directions_float3 = np.empty(directions.shape[0], dtype=gpuarray.vec.float3) directions_float3['x'] = directions[:,0] directions_float3['y'] = directions[:,1] directions_float3['z'] = directions[:,2] directions_gpu = cuda.to_device(directions_float3) -wavelengths = np.random.uniform(200, 800, nphotons).astype(np.float32) +wavelengths = np.random.uniform(400, 700, nphotons).astype(np.float32) wavelengths_gpu = cuda.to_device(wavelengths) times = np.tile(0, nphotons).astype(np.float32) @@ -61,26 +100,84 @@ polarizations_float3['y'] = polarizations[:,1] polarizations_float3['z'] = polarizations[:,2] polarizations_gpu = cuda.to_device(polarizations_float3) -last_hit_triangles = np.empty(nphotons, dtype=np.int32) +last_hit_triangles = -np.ones(nphotons, dtype=np.int32) last_hit_triangles_gpu = cuda.to_device(last_hit_triangles) -states = np.empty(nphotons, dtype=np.int32) +states = -np.ones(nphotons, dtype=np.int32) states_gpu = cuda.to_device(states) -nblocks = 64 +x = np.empty((nphotons, 10, 3)) -gpu_kwargs = {'block': (nblocks,1,1), 'grid': (nphotons/nblocks+1,1)} +for i in range(10): + x[:,i,0] = positions_float3['x'] + x[:,i,1] = positions_float3['y'] + x[:,i,2] = positions_float3['z'] -cuda_propagate(np.int32(nphotons), positions_gpu, directions_gpu, wavelengths_gpu, polarizations_gpu, times_gpu, states_gpu, last_hit_triangles_gpu, np.int32(geometry.node_map.size-1), np.int32(geometry.first_node), **gpu_kwargs) + t0 = time.time() + propagate(np.int32(nphotons), positions_gpu, directions_gpu, wavelengths_gpu, polarizations_gpu, times_gpu, states_gpu, last_hit_triangles_gpu, np.int32(geometry.node_map.size-1), np.int32(geometry.first_node), np.int32(1), **gpu_kwargs) + cuda.Context.synchronize() + elapsed = time.time() - t0 + print 'elapsed %f sec; %f photons/sec ' % (elapsed, nphotons/elapsed) -cuda.memcpy_dtoh(states, states_gpu) + cuda.memcpy_dtoh(positions_float3, positions_gpu) + cuda.memcpy_dtoh(directions_float3, directions_gpu) + cuda.memcpy_dtoh(states, states_gpu) + cuda.memcpy_dtoh(last_hit_triangles, last_hit_triangles_gpu) -print states + print np.unique(states) +from chromaticity import map_wavelength +#mask = states != 0 +mask = geometry.solid_id[last_hit_triangles] == 1 +print 'mask length = ', len(states[mask]) -#from view import view +rgb_colors = map_wavelength(wavelengths[mask]) +colors = rgb_colors*255 +colors = colors.astype(np.uint32) +colors = colors[:,0] << 16 | colors[:,1] << 8 | colors[:,2] -#view(geometry) +def format_hex_string(s): + return '#' + s.rstrip('L')[2:].zfill(6) + +colors = map(format_hex_string, map(hex, colors)) + +plt.figure() +plt.plot(*roundrobin(x[mask,:,0], x[mask,:,1], colors)) +plt.show() + +size = (600, 600) + +y = np.linspace(np.min(film_mesh.vertices[:,1]), np.max(film_mesh.vertices[:,1]), size[0]) +z = np.linspace(np.min(film_mesh.vertices[:,2]), np.max(film_mesh.vertices[:,2]), size[1]) + +film_positions = np.empty((len(positions_float3[mask]), 3)) +film_positions[:,0] = positions_float3[mask]['x'] +film_positions[:,1] = positions_float3[mask]['y'] +film_positions[:,2] = positions_float3[mask]['z'] + +film_positions -= film_solid.displacement +film_positions = np.inner(film_positions, np.linalg.inv(film_solid.rotation)) + +pixels = np.zeros((size[0], size[1], 3)) +for position, rgb_color in zip(film_positions, rgb_colors): + ybin = np.digitize(np.array([position[1]]), y)[0] + zbin = np.digitize(np.array([position[2]]), z)[0] + + try: + pixels[ybin,zbin] += rgb_color + except IndexError: + continue + +import pygame + +pixels = ((pixels/np.max(pixels))*255).astype(np.uint32) +pixels = pixels[:,:,0] << 16 | pixels[:,:,1] << 8 | pixels[:,:,2] + +pygame.init() +screen = pygame.display.set_mode(size) +pygame.surfarray.blit_array(screen, pixels) +pygame.display.flip() +raw_input() |