Refactor sim.py into a reusable Simulation class that is called by the

main function().  Also cleanup more event data structure names and add
an nphoton value that is preserved even if you prune off all the
actual photon vertices.

5 files changed, 112 insertions, 74 deletions

diff --git a/event.py b/event.py
index bf7c007..5b93f17 100644
--- a/event.py
+++ b/event.py
@@ -46,6 +46,7 @@ class Event(object):
         self.gen_total_energy = gen_total_energy
 
         self.subtracks = []
+        self.nphoton = 0
         self.photon_start = None
         self.photon_stop = None
         self.channels = None
diff --git a/fileio/root.C b/fileio/root.C
index 8f7d0a9..7052a71 100644
--- a/fileio/root.C
+++ b/fileio/root.C
@@ -4,10 +4,10 @@
 #include <string>
 
 struct Photon {
-  double t;
-  TVector3 pos;
-  TVector3 dir;
-  TVector3 pol;
+  double time;
+  TVector3 position;
+  TVector3 direction;
+  TVector3 polarization;
   double wavelength; // nm
   unsigned int history;
   int last_hit_triangle;
@@ -15,19 +15,20 @@ struct Photon {
 
 struct Track {
   std::string particle;
-  double t;
-  TVector3 pos;
-  TVector3 dir;
+  double time;
+  TVector3 position;
+  TVector3 direction;
   double start_time;
   double total_energy;
 };
 
 struct MC {
   std::string particle;
-  TVector3 gen_pos;
-  TVector3 gen_dir;
-  double gen_total_e;
+  TVector3 gen_position;
+  TVector3 gen_direction;
+  double gen_total_energy;
 
+  int nphoton;
   
   std::vector<Track> subtrack;
   std::vector<Photon> photon_start;
@@ -36,10 +37,10 @@ struct MC {
 };
 
 struct Channel {
-  Channel() : channel_id(-1), t(-9999.0), q(-9999.0) { };
+  Channel() : channel_id(-1), time(-9999.0), charge(-9999.0) { };
   int channel_id;
-  double t;
-  double q;
+  double time;
+  double charge;
   unsigned int mc_history;
 };
   
@@ -65,8 +66,8 @@ struct Event {
 
       if (channel_id < nentries) {
 	hit[channel_id] = 1;
-	time[channel_id] = channel[i].t;
-	charge[channel_id] = channel[i].q;
+	time[channel_id] = channel[i].time;
+	charge[channel_id] = channel[i].charge;
       }
     }
   }
@@ -83,10 +84,10 @@ void fill_photons(Event *ev, bool start,
   
   for (unsigned int i=0; i < nphotons; i++) {
     Photon &photon = photons[i];
-    photon.t = t0[i];
-    photon.pos.SetXYZ(pos[3*i], pos[3*i + 1], pos[3*i + 2]);
-    photon.dir.SetXYZ(dir[3*i], dir[3*i + 1], dir[3*i + 2]);
-    photon.pol.SetXYZ(pol[3*i], pol[3*i + 1], pol[3*i + 2]);
+    photon.time = t0[i];
+    photon.position.SetXYZ(pos[3*i], pos[3*i + 1], pos[3*i + 2]);
+    photon.direction.SetXYZ(dir[3*i], dir[3*i + 1], dir[3*i + 2]);
+    photon.polarization.SetXYZ(pol[3*i], pol[3*i + 1], pol[3*i + 2]);
     photon.wavelength = wavelength[i];
     if (histories)
       photon.history = histories[i];
@@ -101,18 +102,18 @@ void fill_photons(Event *ev, bool start,
 }
 		  
 
-void fill_hits(Event *ev, unsigned int nchannels, float *t,
-	       float *q, unsigned int *history)
+void fill_hits(Event *ev, unsigned int nchannels, float *time,
+	       float *charge, unsigned int *history)
 {
   ev->channel.resize(0);
   ev->nhit = 0;
   Channel ch;
   for (unsigned int i=0; i < nchannels; i++) {
-    if (t[i] < 1e8) {
+    if (time[i] < 1e8) {
       ev->nhit++;
       ch.channel_id = i;
-      ch.t = t[i];
-      ch.q = q[i];
+      ch.time = time[i];
+      ch.charge = charge[i];
       ch.mc_history = history[i];
       ev->channel.push_back(ch);
     }
diff --git a/fileio/root.py b/fileio/root.py
index 5bebda0..77c23ff 100644
--- a/fileio/root.py
+++ b/fileio/root.py
@@ -20,9 +20,10 @@ class RootWriter(object):
         self.ev.event_id = pyev.event_id
 
         self.ev.mc.particle = pyev.particle_name
-        self.ev.mc.gen_pos.SetXYZ(*pyev.gen_position)
-        self.ev.mc.gen_dir.SetXYZ(*pyev.gen_direction)
+        self.ev.mc.gen_position.SetXYZ(*pyev.gen_position)
+        self.ev.mc.gen_direction.SetXYZ(*pyev.gen_direction)
         self.ev.mc.gen_total_energy = pyev.gen_total_energy
+        self.ev.mc.nphoton = pyev.nphoton
 
         if pyev.photon_start is not None:
             photons = pyev.photon_start
@@ -46,8 +47,8 @@ class RootWriter(object):
             self.ev.mc.subtrack.resize(len(pyev.subtracks))
             for i, subtrack in enumerate(pyev.subtracks):
                 self.ev.mc.subtrack[i].name = subtrack.particle_name
-                self.ev.mc.subtrack[i].pos.SetXYZ(*subtrack.position)
-                self.ev.mc.subtrack[i].dir.SetXYZ(*subtrack.direction)
+                self.ev.mc.subtrack[i].position.SetXYZ(*subtrack.position)
+                self.ev.mc.subtrack[i].direction.SetXYZ(*subtrack.direction)
                 self.ev.mc.subtrack[i].start_time = subtrack.start_time
                 self.ev.mc.subtrack[i].total_energy = subtrack.total_energy
 
diff --git a/generator/g4gen.py b/generator/g4gen.py
index d4e79a1..4416c45 100644
--- a/generator/g4gen.py
+++ b/generator/g4gen.py
@@ -111,8 +111,8 @@ class G4Generator(object):
         else:
             # Create temporary subtrack for single primary particle
             subtracks = [event.Subtrack(particle_name=ev.particle_name,
-                                        position=ev.gen_pos,
-                                        direction=ev.gen_dir,
+                                        position=ev.gen_position,
+                                        direction=ev.gen_direction,
                                         start_time=0.0,
                                         total_energy=ev.gen_total_energy)]
 
diff --git a/sim.py b/sim.py
index a7dc220..9e9ef18 100755
--- a/sim.py
+++ b/sim.py
@@ -9,8 +9,10 @@ import detectors
 import optics
 import generator
 from generator import constant
+import itertools
 import gpu
 from fileio import root
+
 from tools import profile_if_possible, enable_debug_on_crash
 
 def pick_seed():
@@ -18,12 +20,77 @@ def pick_seed():
     a mixture of the current time and the current process ID.'''
     return int(time.time()) ^ (os.getpid() << 16)
 
+class Simulation(object):
+    def __init__(self, detector, detector_material, 
+                 seed=None, cuda_device=None, geant4_processes=4, bvh_bits=11):
+        self.detector = detector
+
+        if seed is None:
+            self.seed = pick_seed()
+        else:
+            self.seed = seed
+        print >>sys.stderr, 'RNG seed:', self.seed
+        # We have three generators to seed: numpy.random, GEANT4, and CURAND.
+        # The latter two are done below
+        np.random.seed(self.seed)
+
+        self.detector_material = detector_material
+        if geant4_processes > 0:
+            self.photon_generator = generator.photon.G4ParallelGenerator(geant4_processes, 
+                                                                         detector_material,
+                                                                         base_seed=self.seed)
+        else:
+            self.photon_generator = None
+
+        print >>sys.stderr, 'Creating BVH for detector "%s" with %d bits...' % (self.detector, bvh_bits)
+        detector.build(bits=bvh_bits)
+        
+        print >>sys.stderr, 'Initializing GPU...'
+        self.gpu_worker = gpu.GPU(cuda_device)
+        
+        print >>sys.stderr, 'Loading detector onto GPU...'
+        self.gpu_worker.load_geometry(detector)
+
+        print >>sys.stderr, 'Initializing random numbers generators...'
+        self.gpu_worker.setup_propagate(seed=self.seed)
+        self.gpu_worker.setup_daq(max(self.detector.pmtids))
+
+    def simulate(self, nevents, vertex_generator, keep_photon_start=False, keep_photon_stop=False,
+                 run_daq=True):
+        return self.simulate_photons(nevents, self.photon_generator.generate_events(nevents, vertex_generator),
+                                     keep_photon_start=keep_photon_start, keep_photon_stop=keep_photon_stop,
+                                     run_daq=run_daq)
+
+    def simulate_photons(self, nevents, photon_generator, keep_photon_start=False, keep_photon_stop=False,
+                         run_daq=True):
+        for ev in itertools.islice(photon_generator, nevents):
+            self.gpu_worker.load_photons(ev.photon_start)
+            self.gpu_worker.propagate()
+            self.gpu_worker.run_daq()
+            
+            ev.nphoton = len(ev.photon_start.positions)
+            
+            if not keep_photon_start:
+                ev.photon_start = None
+
+            if keep_photon_stop:
+                ev.photon_stop = self.gpu_worker.get_photons()
+
+            if run_daq:
+                ev.hits = self.gpu_worker.get_hits()        
+
+            yield ev
+
+    def propagate_photons(self, photons, max_steps=10):
+        self.gpu_worker.load_photons(photons)
+        self.gpu_worker.propagate()
+        return self.gpu_worker.get_photons()
+
 @profile_if_possible
 def main():
     parser = optparse.OptionParser('%prog')
     parser.add_option('-b', type='int', dest='nbits', default=10)
     parser.add_option('-j', type='int', dest='device', default=None)
-    parser.add_option('-n', type='int', dest='nblocks', default=64)
     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,
@@ -55,13 +122,6 @@ def main():
     position = np.array(eval(options.pos), dtype=float)
     direction = np.array(eval(options.dir), dtype=float)
     detector = detectors.find(options.detector)
-    if options.seed is None:
-        options.seed = pick_seed()
-
-    print >>sys.stderr, 'RNG seed:', options.seed
-    # We have three generators to seed: numpy.random, GEANT4, and CURAND.
-    # The latter two are done below
-    np.random.seed(options.seed)
 
     print >>sys.stderr, 'Creating generator...'
     if options.particle == 'pi0':
@@ -73,27 +133,12 @@ def main():
                                                          position=constant(position),
                                                          direction=constant(direction),
                                                          total_energy=constant(options.energy))
-    detector_material = optics.water_wcsim
-    photon_generator = generator.photon.G4ParallelGenerator(options.ngenerators, detector_material,
-                                                           base_seed=options.seed)
-    print >>sys.stderr, 'WARNING: ASSUMING DETECTOR IS WCSIM WATER!!'
-
-    # Do this now so we can get ahead of the photon propagation
-    print >>sys.stderr, 'Starting GEANT4 generators...'
-    event_iterator = photon_generator.generate_events(options.nevents, vertex_generator)
-
-    print >>sys.stderr, 'Creating BVH for detector "%s" with %d bits...' % (options.detector, options.nbits)
-    detector.build(bits=options.nbits)
-
-    print >>sys.stderr, 'Initializing GPU...'
-    gpu_worker = gpu.GPU(options.device)
-    
-    print >>sys.stderr, 'Loading detector onto GPU...'
-    gpu_worker.load_geometry(detector)
-    
-    print >>sys.stderr, 'Initializing random numbers generators...'
-    gpu_worker.setup_propagate(seed=options.seed)
-    gpu_worker.setup_daq(max(detector.pmtids))
+
+    # Initializing simulation
+    print >>sys.stderr, 'WARNING: ASSUMING DETECTOR IS WCSIM WATER!!'    
+    simulation = Simulation(detector=detector, detector_material= optics.water_wcsim,
+                            seed=options.seed, cuda_device=options.device, 
+                            geant4_processes=options.ngenerators, bvh_bits=options.nbits)
 
     # Create output file
     writer = root.RootWriter(output_filename)
@@ -102,22 +147,12 @@ def main():
     start_sim = time.time()
     nphotons = 0
 
-    for i, ev in enumerate(event_iterator):
-        photons = ev.photon_start
-        assert len(photons.positions) > 0, 'GEANT4 generated event with no photons!'
-
-        nphotons += len(photons.positions)
-
-        gpu_worker.load_photons(photons)
-
-        gpu_worker.propagate()
-        gpu_worker.run_daq()
+    for i, ev in enumerate(simulation.simulate(options.nevents, vertex_generator,
+                                               keep_photon_start=options.save_photon_start,
+                                               keep_photon_stop=options.save_photon_stop)):
+        assert ev.nphoton > 0, 'GEANT4 generated event with no photons!'
+        nphotons += ev.nphoton
 
-        ev.hits = gpu_worker.get_hits()        
-        if not options.save_photon_start:
-            ev.photon_start = None
-        if options.save_photon_stop:
-            ev.photon_stop = gpu_worker.get_photons()
         writer.write_event(ev)
 
         if i % 10 == 0: