G4Generator and sim.py can now generate boosted pi0 decays.

The ROOT data structure has been extended to allow storage
of the individual gamma rays as subtracks.

1 files changed, 16 insertions, 5 deletions

diff --git a/sim.py b/sim.py
index e123e5e..b8382e6 100755
--- a/sim.py
+++ b/sim.py
@@ -51,10 +51,15 @@ class GeneratorProcess(multiprocessing.Process):
         generator = g4gen.G4Generator(self.material, seed=self.seed)
 
         for i in xrange(self.nevents):
-            photons = generator.generate_photons(particle_name=self.particle, 
-                                                 total_energy=self.energy,
+            if self.particle == 'pi0':
+                photons = generator.generate_pi0(total_energy=self.energy,
                                                  position=self.position,
                                                  direction=self.direction)
+            else:
+                photons = generator.generate_photons(particle_name=self.particle, 
+                                                     total_energy=self.energy,
+                                                     position=self.position,
+                                                     direction=self.direction)
             self.queue.put(photons)
 
 
@@ -150,7 +155,8 @@ def main():
 
         nphotons += len(photons['pos'])
 
-        gpu_worker.load_photons(**photons)
+        gpu_worker.load_photons(pos=photons['pos'], dir=photons['dir'], pol=photons['pol'],
+                                t0=photons['t0'], wavelength=photons['wavelength'])
         gpu_worker.propagate()
         gpu_worker.run_daq()
         hits = gpu_worker.get_hits()
@@ -163,8 +169,13 @@ def main():
             photon_stop = gpu_worker.get_photons()
         else:
             photon_stop = None
-            
-        writer.write_event(i, hits, photon_start=photon_start, photon_stop=photon_stop)
+        
+        if 'subtracks' in photons:
+            subtracks = photons['subtracks']
+        else:
+            subtracks = None
+        writer.write_event(i, hits, photon_start=photon_start, photon_stop=photon_stop,
+                           subtracks=subtracks)
         if i % 10 == 0:
             print >>sys.stderr, "\rEvent:", i,