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.

3 files changed, 67 insertions, 12 deletions

diff --git a/fileio/root.py b/fileio/root.py
index 0b80745..004484f 100644
--- a/fileio/root.py
+++ b/fileio/root.py
@@ -20,7 +20,7 @@ class RootWriter(object):
         self.ev.mc.gen_dir.SetXYZ(*direction)
         self.ev.mc.gen_total_e = total_e
 
-    def write_event(self, event_id, hits, photon_start=None, photon_stop=None):
+    def write_event(self, event_id, hits, photon_start=None, photon_stop=None, subtracks=None):
         self.ev.event_id = event_id
         if photon_start is not None:
             photons = photon_start
@@ -40,6 +40,15 @@ class RootWriter(object):
                               photons['wavelength'], photons['t0'],
                               photons['histories'], photons['last_hit_triangles'])
 
+        self.ev.mc.subtrack.resize(0)
+        if subtracks is not None:
+            self.ev.mc.subtrack.resize(len(subtracks))
+            for i, subtrack in enumerate(subtracks):
+                self.ev.mc.subtrack[i].name = subtrack['name']
+                self.ev.mc.subtrack[i].pos.SetXYZ(*subtrack['pos'])
+                self.ev.mc.subtrack[i].dir.SetXYZ(*subtrack['dir'])
+                self.ev.mc.subtrack[i].total_energy = subtrack['total_e']
+
         ROOT.fill_hits(self.ev, len(hits['t']), hits['t'], hits['q'], hits['history'])
         self.T.Fill()
 
diff --git a/g4gen.py b/g4gen.py
index 718df6e..6766d42 100644
--- a/g4gen.py
+++ b/g4gen.py
@@ -3,7 +3,8 @@ import g4py.ezgeom
 import g4py.NISTmaterials
 import g4py.ParticleGun
 import pyublas
-import numpy
+import numpy as np
+import pi0
 
 try:
     import G4chroma
@@ -70,6 +71,40 @@ class G4Generator(object):
         g4material.SetMaterialPropertiesTable(prop_table)
         return g4material
 
+    def generate_pi0(self, total_energy, position, direction):
+        '''Use GEANT4 to generate photons produced by a pi0 decay.
+
+        pi0 event will be generated by running two gammas with the appropriate
+        energy and angular distribution.
+
+               total_energy: Total energy of pi0 (incl rest mass) in MeV
+               position: 3-tuple of position of particle in global coordinates
+               direction: 3-tuple direction vector.  
+                          Does not have to be normalized.
+        '''
+        direction = direction / np.linalg.norm(direction)
+        
+        cos_theta_rest = np.random.random_sample() * 2 - 1
+        theta_rest = np.arccos(cos_theta_rest)
+        phi_rest = np.random.random_sample() * 2 * np.pi
+
+        (gamma1_e, gamma1_dir), (gamma2_e, gamma2_dir) = pi0.pi0_decay(energy=total_energy,
+                                                                       direction=direction,
+                                                                       theta=theta_rest,
+                                                                       phi=phi_rest)
+
+        gamma1 = self.generate_photons('gamma', gamma1_e, position, gamma1_dir)
+        gamma2 = self.generate_photons('gamma', gamma2_e, position, gamma2_dir)
+
+        photons = {'subtracks' : [ {'name':'gamma', 'pos': position, 't0' : 0.0,
+                                    'dir' : gamma1_dir, 'total_e' : gamma1_e },
+                                   {'name':'gamma', 'pos': position, 't0' : 0.0,
+                                    'dir' : gamma2_dir, 'total_e' : gamma2_e } ]
+                    }
+        for key in gamma1.keys():
+            photons[key] = np.concatenate((gamma1[key], gamma2[key]))
+                                                                       
+        return photons
 
     def generate_photons(self, particle_name, total_energy, position, direction):
         '''Use GEANT4 to generate photons produced by the given particle.
@@ -89,23 +124,23 @@ class G4Generator(object):
 
         gRunManager.BeamOn(1)
         n = self.tracking_action.GetNumPhotons()        
-        pos = numpy.zeros(shape=(n,3), dtype=numpy.float32)
+        pos = np.zeros(shape=(n,3), dtype=np.float32)
         pos[:,0] = self.tracking_action.GetX()
         pos[:,1] = self.tracking_action.GetY()
         pos[:,2] = self.tracking_action.GetZ()
 
-        dir = numpy.zeros(shape=(n,3), dtype=numpy.float32)
+        dir = np.zeros(shape=(n,3), dtype=np.float32)
         dir[:,0] = self.tracking_action.GetDirX()
         dir[:,1] = self.tracking_action.GetDirY()
         dir[:,2] = self.tracking_action.GetDirZ()
 
-        pol = numpy.zeros(shape=(n,3), dtype=numpy.float32)
+        pol = np.zeros(shape=(n,3), dtype=np.float32)
         pol[:,0] = self.tracking_action.GetPolX()
         pol[:,1] = self.tracking_action.GetPolY()
         pol[:,2] = self.tracking_action.GetPolZ()
         
-        wavelength = self.tracking_action.GetWavelength().astype(numpy.float32)
-        t0 = self.tracking_action.GetT0().astype(numpy.float32)
+        wavelength = self.tracking_action.GetWavelength().astype(np.float32)
+        t0 = self.tracking_action.GetT0().astype(np.float32)
         return { 'pos' : pos, 
                  'dir' : dir, 
                  'pol' : pol,
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,