1 files changed, 60 insertions, 0 deletions
diff --git a/track.py b/track.py
new file mode 100644
index 0000000..f73f2ed
--- /dev/null
+++ b/track.py
@@ -0,0 +1,60 @@
+import numpy as np
+import pycuda.driver as cuda
+from gputhread import GPUThread
+from detectors import LBNE
+from Queue import Queue
+from threadtest import create_job
+import matplotlib.pyplot as plt
+from itertoolset import roundrobin
+from color import map_wavelength
+from solids import r7081
+from geometry import Geometry
+
+nphotons = 100000
+
+jobs = Queue()
+output = Queue()
+
+#geometry = LBNE()
+geometry = Geometry()
+geometry.add_solid(r7081, displacement=(0,-1,0))
+geometry.build(bits=8)
+
+cuda.init()
+
+try:
+    gputhread = GPUThread(0, geometry, jobs, output, 64)
+    gputhread.start()
+
+    job = create_job((0,0,0), nphotons)
+
+    x = np.empty((nphotons, 10, 3))
+    for i in range(10):
+        print '%i' % i
+        x[:,i,0] = job.positions['x']
+        x[:,i,1] = job.positions['y']
+        x[:,i,2] = job.positions['z']
+
+        jobs.put(job)
+        jobs.join()
+
+        job = output.get()
+
+        for j in np.unique(job.states):
+            print 'state %2i, %i' % (j, len(job.states[job.states == j]))
+finally:
+    gputhread.stop()
+    gputhread.join()
+
+mask = job.states != 0
+
+rgb = (map_wavelength(job.wavelengths[mask])*255).astype(np.uint32)
+
+def format_hex_string(s):
+    return '#' + s.rstrip('L')[2:].zfill(6)
+
+colors = map(format_hex_string, map(hex, rgb[:,0] << 16 | rgb[:,1] << 8 | rgb[:,2]))
+
+plt.figure()
+plt.plot(*roundrobin(x[mask,:,0], x[mask,:,1], colors))
+plt.show()