diff options
Diffstat (limited to 'g4gen.py')
| -rw-r--r-- | g4gen.py | 47 |
1 files changed, 41 insertions, 6 deletions
@@ -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, |