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, 41 insertions, 6 deletions

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,