merge

1 files changed, 21 insertions, 111 deletions

diff --git a/sim.py b/sim.py
index 7c9720a..2e8f2f9 100755
--- a/sim.py
+++ b/sim.py
@@ -1,19 +1,12 @@
 #!/usr/bin/env python
-import sys
 import time
 import os
 import numpy as np
-import itertools
-import threading
 
-from chroma import detectors
-from chroma import optics
 from chroma import generator
 from chroma import gpu
-from chroma.fileio import root
-from chroma import tools
 from chroma import event
-from chroma.itertoolset import peek, repeat_copy, repeating_iterator
+from chroma import itertoolset
 
 def pick_seed():
     """Returns a seed for a random number generator selected using
@@ -21,7 +14,8 @@ def pick_seed():
     return int(time.time()) ^ (os.getpid() << 16)
 
 class Simulation(object):
-    def __init__(self, detector, seed=None, cuda_device=None, geant4_processes=4, bvh_bits=11, use_cache=True, nthreads_per_block=64, max_blocks=1024):
+    def __init__(self, detector, seed=None, cuda_device=None,
+                 geant4_processes=4, nthreads_per_block=64, max_blocks=1024):
         self.detector = detector
 
         self.nthreads_per_block = nthreads_per_block
@@ -32,7 +26,6 @@ class Simulation(object):
         else:
             self.seed = seed
 
-        print 'RNG seed: %i' % self.seed
         # We have three generators to seed: numpy.random, GEANT4, and CURAND.
         # The latter two are done below.
         np.random.seed(self.seed)
@@ -42,24 +35,24 @@ class Simulation(object):
         else:
             self.photon_generator = None
 
+        self.context = gpu.create_cuda_context(cuda_device)
+
         if not hasattr(detector, 'mesh'):
             # need to build geometry
-            print 'Creating BVH with %i bits...' % bvh_bits
-            detector.build(bits=bvh_bits, use_cache=use_cache)
-        
-        self.gpu = gpu.GPU(cuda_device)
+            detector.build()
+
         self.gpu_geometry = gpu.GPUGeometry(detector)
         self.gpu_daq = gpu.GPUDaq(self.gpu_geometry, max(self.detector.pmtids))
         self.gpu_pdf = gpu.GPUPDF()
 
-        print 'Initializing random numbers generators...'
         self.rng_states = gpu.get_rng_states(self.nthreads_per_block*self.max_blocks, seed=self.seed)
 
         self.pdf_config = None
 
-    def simulate(self, iterable, keep_photons_beg=False, keep_photons_end=False, run_daq=True, max_steps=10):
+    def simulate(self, iterable, keep_photons_beg=False,
+                 keep_photons_end=False, run_daq=True, max_steps=10):
         try:
-            first_element, iterable = peek(iterable)
+            first_element, iterable = itertoolset.peek(iterable)
         except TypeError:
             first_element, iterable = iterable, [iterable]
 
@@ -68,13 +61,16 @@ class Simulation(object):
         elif isinstance(first_element, event.Photons):
             iterable = (event.Event(photons_beg=x) for x in iterable)
         elif isinstance(first_element, event.Vertex):
-            iterable = (event.Event(vertices=[vertex]) for vertex in iterable)
+            iterable = (event.Event(primary_vertex=vertex, vertices=[vertex]) for vertex in iterable)
             iterable = self.photon_generator.generate_events(iterable)
 
         for ev in iterable:
             gpu_photons = gpu.GPUPhotons(ev.photons_beg)
 
-            gpu_photons.propagate(self.gpu_geometry, self.rng_states, nthreads_per_block=self.nthreads_per_block, max_blocks=self.max_blocks, max_steps=max_steps)
+            gpu_photons.propagate(self.gpu_geometry, self.rng_states,
+                                  nthreads_per_block=self.nthreads_per_block,
+                                  max_blocks=self.max_blocks,
+                                  max_steps=max_steps)
 
             ev.nphotons = len(ev.photons_beg.pos)
 
@@ -93,7 +89,7 @@ class Simulation(object):
     def create_pdf(self, iterable, tbins, trange, qbins, qrange, nreps=1):
         """Returns tuple: 1D array of channel hit counts, 3D array of
         (channel, time, charge) pdfs."""
-        first_element, iterable = peek(iterable)
+        first_element, iterable = itertoolset.peek(iterable)
 
         if isinstance(first_element, event.Event):
             iterable = self.photon_generator.generate_events(iterable)
@@ -102,12 +98,12 @@ class Simulation(object):
         if pdf_config != self.pdf_config:
             self.pdf_config = pdf_config
             self.gpu_pdf.setup_pdf(max(self.detector.pmtids), tbins, trange,
-                                      qbins, qrange)
+                                   qbins, qrange)
         else:
             self.gpu_pdf.clear_pdf()
 
         if nreps > 1:
-            iterable = repeating_iterator(iterable, nreps)
+            iterable = itertoolset.repeating_iterator(iterable, nreps)
 
         for ev in iterable:
             gpu_photons = gpu.GPUPhotons(ev.photons_beg)
@@ -132,7 +128,7 @@ class Simulation(object):
                                     min_bin_content=min_bin_content,
                                     time_only=True)
 
-        first_element, iterable = peek(iterable)
+        first_element, iterable = itertoolset.peek(iterable)
 
         if isinstance(first_element, event.Event):
             iterable = self.photon_generator.generate_events(iterable)
@@ -149,91 +145,5 @@ class Simulation(object):
         
         return self.gpu_pdf.get_pdf_eval()
 
-@tools.profile_if_possible
-def main():
-    import optparse
-
-    parser = optparse.OptionParser('%prog filename')
-    parser.add_option('-b', type='int', dest='nbits', default=11)
-    parser.add_option('-j', type='int', dest='device', default=None)
-    parser.add_option('-s', type='int', dest='seed', default=None,
-                      help='Set random number generator seed')
-    parser.add_option('-g', type='int', dest='ngenerators', default=4,
-                      help='Number of GEANT4 generator processes')
-    parser.add_option('--detector', type='string', dest='detector',
-                      default='microlbne')
-    parser.add_option('--nevents', type='int', dest='nevents', default=100)    
-    parser.add_option('--particle', type='string', dest='particle',
-                      help='particle name', default='e-')
-    parser.add_option('--ke', type='float', dest='ke',
-                      help='particle kinetic energy in MeV', default=100.0)
-    parser.add_option('--pos', type='string', dest='pos',
-                      help='particle vertex origin.', default='(0,0,0)')
-    parser.add_option('--dir', type='string', dest='dir',
-                      help='particle vertex direction.', default='(1,0,0)')
-    parser.add_option('--save-photon-beg', action='store_true', 
-                      dest='save_photons_beg', default=False,
-                      help='Save initial photon vertices to disk')
-    parser.add_option('--save-photon-end', action='store_true', 
-                      dest='save_photons_end', default=False,
-                      help='Save final photon vertices to disk')
-
-    options, args = parser.parse_args()
-
-    if len(args) < 1:
-        sys.exit(parser.format_help())
-    else:
-        output_filename = args[0]
-
-    if options.nevents <= 0:
-        sys.exit('--nevents must be greater than 0!')
-
-    pos = np.array(eval(options.pos), dtype=float)
-    dir = np.array(eval(options.dir), dtype=float)
-
-    print 'Loading detector %s...' % options.detector
-    sys.stdout.flush()
-    detector = detectors.find(options.detector)
-
-    print 'Creating particle vertex generator...'
-    sys.stdout.flush()
-    if options.particle == 'pi0':
-        ev_vertex_iter = itertools.islice(generator.vertex.pi0_gun(itertools.repeat(pos), itertools.repeat(dir), itertools.repeat(options.ke)), options.nevents)
-    else:
-        vertex = event.Vertex(options.particle, pos, dir, None, options.ke)
-        ev_vertex_iter = (event.Event(i, vertex, [vertex]) for i, vertex in zip(range(options.nevents), repeat_copy(vertex)))
-
-    # Initializing simulation
-    simulation = Simulation(detector=detector, seed=options.seed, cuda_device=options.device, geant4_processes=options.ngenerators, bvh_bits=options.nbits)
-
-    # Create output file
-    writer = root.RootWriter(output_filename)
-
-    # Preheat generator
-    ev_iter = simulation.simulate(ev_vertex_iter, keep_photons_beg=options.save_photons_beg, keep_photons_end=options.save_photons_end)
-
-    print 'Starting simulation...'
-
-    nphotons = 0
-    t0 = time.time()
-
-    for i, ev in enumerate(ev_iter):
-        print "\rEvent: %i" % (i+1),
-        sys.stdout.flush()
-
-        assert ev.nphotons > 0, 'Geant4 generated event with no photons!'
-        nphotons += ev.nphotons
-
-        writer.write_event(ev)
-    print
-
-    elapsed = time.time() - t0
-
-    writer.close()
-
-    print '%f elapsed, %1.1f events/sec, %1.0f photons/sec.' % \
-        (elapsed, options.nevents/elapsed, nphotons/elapsed)
-
-if __name__ == '__main__':
-    tools.enable_debug_on_crash()
-    main()
+    def __del__(self):
+        self.context.pop()