1 files changed, 183 insertions, 0 deletions
diff --git a/photon_map.py b/photon_map.py
new file mode 100644
index 0000000..4c6110e
--- /dev/null
+++ b/photon_map.py
@@ -0,0 +1,183 @@
+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, flashlight
+from pycuda import autoinit
+from pycuda.compiler import SourceModule
+from pycuda import gpuarray
+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_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_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_mesh.vertices *= 100
+sphere_solid = Solid(2, sphere_mesh, material1=vacuum, material2=vacuum, surface=black_surface)
+
+geometry = Geometry()#[lens_solid1, lens_solid2, film_solid, sphere_solid])
+geometry.add_solid(lens_solid)
+geometry.add_solid(film_solid)
+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)
+propagate = module.get_function('propagate')
+init_rng = module.get_function('init_rng')
+
+init_rng(np.int32(nphotons), np.int32(0), np.int32(0), **gpu_kwargs)
+
+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 = 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(400, 700, nphotons).astype(np.float32)
+wavelengths_gpu = cuda.to_device(wavelengths)
+
+times = np.tile(0, nphotons).astype(np.float32)
+times_gpu = cuda.to_device(times)
+
+polarizations = uniform_sphere(nphotons)
+polarizations_float3 = np.empty(polarizations.shape[0], dtype=gpuarray.vec.float3)
+polarizations_float3['x'] = polarizations[:,0]
+polarizations_float3['y'] = polarizations[:,1]
+polarizations_float3['z'] = polarizations[:,2]
+polarizations_gpu = cuda.to_device(polarizations_float3)
+
+last_hit_triangles = -np.ones(nphotons, dtype=np.int32)
+last_hit_triangles_gpu = cuda.to_device(last_hit_triangles)
+
+states = -np.ones(nphotons, dtype=np.int32)
+states_gpu = cuda.to_device(states)
+
+x = np.empty((nphotons, 10, 3))
+
+for i in range(10):
+    x[:,i,0] = positions_float3['x']
+    x[:,i,1] = positions_float3['y']
+    x[:,i,2] = positions_float3['z']
+
+    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(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 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])
+
+rgb_colors = map_wavelength(wavelengths[mask])
+colors = rgb_colors*255
+colors = colors.astype(np.uint32)
+colors = colors[:,0] << 16 | colors[:,1] << 8 | colors[:,2]
+
+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()