added some more documentation and a more accurate miniature version of lbne

1 files changed, 59 insertions, 109 deletions

diff --git a/detectors/lbne.py b/detectors/lbne.py
index cac737a..54e11a9 100644
--- a/detectors/lbne.py
+++ b/detectors/lbne.py
@@ -1,130 +1,80 @@
-import os
 import numpy as np
-import pickle
-from chroma import *
 from copy import deepcopy
+from chroma import layout
+from chroma.stl import read_stl
+from chroma.transform import rotate
+from chroma.geometry import Geometry, Solid
+from chroma.materials import glass, h2o
+from chroma.photon import uniform_sphere
+from chroma.gpu import GPU
+from itertools import product
 from histogram import *
 
-models_directory = os.path.split(os.path.realpath(__file__))[0] + '/../models'
+endcap_spacing = .485
 
-strings = 20
-pmts_per_string = 10
-radius = 10.0
-height = 20.0
+radius = 25.0/10.0
+height = 50.0/10.0
 
-grid_spacing = height/pmts_per_string
+nstrings = 324//10
+pmts_per_string = 102//10
 
-block_size = 64
-
-
-class LBNE(geometry.Geometry):
+class LBNE(Geometry):
+    """Miniature version of the LBNE water Cerenkov detector geometry."""
     def __init__(self):
         super(LBNE, self).__init__()
 
-        pmt_mesh = stl.read_stl(models_directory + '/hamamatsu_12inch.stl')
-        pmt_mesh /= 1000.0
-
-        apmt = geometry.Solid(pmt_mesh, materials.glass, materials.h2o)
-
-        self.pmt_index = []
-        self.pmt_local_axes = []
-        self.pmt_positions = []
-
-        self.pmt_hits = []
+        pmt_mesh = read_stl(layout.models + '/hamamatsu_12inch.stl')/1000.0
+        pmt_solid = Solid(pmt_mesh, glass, h2o)
 
+        # construct the barrel
         for i in range(pmts_per_string):
-            for j in range(strings):
-                pmt = deepcopy(apmt)
-                pmt.mesh += (-radius,0,i*(height/pmts_per_string))
-                pmt.mesh = transform.rotate(pmt.mesh, j*2*np.pi/strings, (0,0,1))
+            for j in range(nstrings):
+                pmt = deepcopy(pmt_solid)
+                pmt.mesh += (-radius,0,-height/2+i*height/(pmts_per_string-1))
+                pmt.mesh = rotate(pmt.mesh, j*2*np.pi/nstrings, (0,0,1))
                 self.add_solid(pmt)
-                self.pmt_hits.append(Histogram(10000, (-0.5, 9999.5)))
-
-        for x in np.arange(-radius, radius, grid_spacing):
-            for y in np.arange(-radius, radius, grid_spacing):
-                if np.sqrt(x**2+y**2) <= radius:
-                    pmt = deepcopy(apmt)
-                    pmt.mesh = transform.rotate(pmt.mesh, np.pi/2, (0,1,0))
-                    pmt.mesh += (x,y,0)
-                    self.add_solid(pmt)
-                    self.pmt_hits.append(Histogram(10000, (-0.5, 9999.5)))
-
-        for x in np.arange(-radius, radius, grid_spacing):
-            for y in np.arange(-radius, radius, grid_spacing):
-                if np.sqrt(x**2+y**2) <= radius:
-                    pmt = deepcopy(apmt)
-                    pmt.mesh = transform.rotate(pmt.mesh, -np.pi/2, (0,1,0))
-                    pmt.mesh += (x,y,height)
-                    self.add_solid(pmt)
-                    self.pmt_hits.append(Histogram(10000, (-0.5, 9999.5)))
-
-        self.build(bits=4)
-
-        self.npmts = len(self.pmt_hits)
-
-        self.gpu = gpu.GPU()
-        self.gpu.load_geometry(self)
-
-    def throw_photon_bomb(self, z_position, nphotons=100000):
-        origin = np.zeros((nphotons,3)) + (0,0,z_position)
-
-        direction = photon.uniform_sphere(nphotons)
-
-        origin_gpu = gpu.cuda.to_device(gpu.make_vector(origin))
-        direction_gpu = gpu.cuda.to_device(gpu.make_vector(direction))
-
-        pixels = np.empty(nphotons, dtype=np.int32)
-        states = np.empty(nphotons, dtype=np.int32)
-
-        pixels_gpu = gpu.cuda.to_device(pixels)
-        states_gpu = gpu.cuda.to_device(states)
 
-        gpu_kwargs = {'block': (block_size,1,1), 'grid': (nphotons//block_size+1,1)}
-        self.gpu.call(np.int32(nphotons), origin_gpu, direction_gpu, np.int32(self.first_leaf), states_gpu, pixels_gpu, **gpu_kwargs)
-
-        gpu.cuda.memcpy_dtoh(states, states_gpu)
-
-        pmt_indices = self.solid_index[states[(states != -1)]]
-
-        bin_count = np.bincount(pmt_indices)
-
-        bin_count = np.append(bin_count, np.zeros(self.npmts-bin_count.size))
-
-        return bin_count
-
-    def generate_event(self, z_position):
-        self.bin_count = self.throw_photon_bomb(z_position)
-
-    def get_likelihood(self, z_position, calls=1000):
-        if not hasattr(self, 'bin_count'):
-            raise Exception('must call generate_event() first')
-
-        for pmt_hit in self.pmt_hits:
-            pmt_hit.reset()
-
-        for i in range(calls):
-            print 'throwing bomb %i' % i
-            bin_count = self.throw_photon_bomb(z_position)
-
-            for i, count in enumerate(bin_count):
-                self.pmt_hits[i].fill(count)
+        # construct the top endcap
+        for x, y in np.array(tuple(product(\
+                    np.arange(-radius+0.075, radius, endcap_spacing),
+                    np.arange(-radius+0.075, radius, endcap_spacing)))):
+            if np.sqrt(x**2 + y**2) <= radius:
+                pmt = deepcopy(pmt_solid)
+                pmt.mesh = rotate(pmt.mesh, -np.pi/2, (0,1,0))
+                pmt.mesh += (x,y,+height/2+height/(pmts_per_string-1)/2)
+                self.add_solid(pmt)
 
-        for pmt_hit in self.pmt_hits:
-            pmt_hit.normalize()
+        # construct the bottom endcap
+        for x, y in np.array(tuple(product(\
+                    np.arange(-radius+0.075, radius, endcap_spacing),
+                    np.arange(-radius+0.075, radius, endcap_spacing)))):
+            if np.sqrt(x**2 + y**2) <= radius:
+                pmt = deepcopy(pmt_solid)
+                pmt.mesh = rotate(pmt.mesh, +np.pi/2, (0,1,0))
+                pmt.mesh += (x,y,-height/2-height/(pmts_per_string-1)/2)
+                self.add_solid(pmt)
 
-        likelihood = 0.0
-        for i in range(self.npmts):
-            probability = self.pmt_hits[i].eval(self.bin_count[i])
+    def load_geometry(self):
+        self.gpu = GPU()
+        self.texrefs = self.gpu.load_geometry(self)
+        self.propagate = self.gpu.get_function('propagate')
 
-            if probability == 0.0:
-                print 'calculating likelihood from pmt %i' % i
-                print 'bin count =', self.bin_count[i]
-                print self.pmt_hits[i].hist
+if __name__ == '__main__':
+    import sys
+    import optparse
+    import pickle
 
-            likelihood -= np.log(self.pmt_hits[i].eval(self.bin_count[i]))
+    parser = optparse.OptionParser('prog filename')
+    parser.add_option('-b', '--bits', type='int', dest='bits',
+                      help='bits for z-ordering space axes', default=6)
+    options, args = parser.parse_args()
 
-        return likelihood
+    if len(args) < 1:
+        sys.exit(parser.format_help())
 
-if __name__ == '__main__':
     lbne = LBNE()
-    view.view(lbne)
+    lbne.build(bits=options.bits)
+
+    f = open(args[0], 'wb')
+    pickle.dump(lbne, f)
+    f.close()