diff options
| -rw-r--r-- | chroma/generator/g4gen.py | 53 |
1 files changed, 32 insertions, 21 deletions
diff --git a/chroma/generator/g4gen.py b/chroma/generator/g4gen.py index 0b8169d..4b0a139 100644 --- a/chroma/generator/g4gen.py +++ b/chroma/generator/g4gen.py @@ -53,7 +53,7 @@ class G4Generator(object): gRunManager.Initialize() g4unmute() # preinitialize the process by running a simple event - self.generate_photons([Vertex('e-', (0,0,0), (1,0,0), 0, 1.0)]) + self.generate_photons([Vertex('e-', (0,0,0), (1,0,0), 0, 1.0)], mute=True) def create_g4material(self, material): g4material = G4Material('world_material', material.density * g / cm3, @@ -107,37 +107,48 @@ class G4Generator(object): return Photons(pos, dir, pol, wavelengths, t0) - def generate_photons(self, vertices): + def generate_photons(self, vertices, mute=False): """Use GEANT4 to generate photons produced by propagating `vertices`. Args: vertices: list of event.Vertex objects List of initial particle vertices. + mute: bool + Disable GEANT4 output to console during generation. (GEANT4 can + be quite chatty.) + Returns: photons: event.Photons Photon vertices generated by the propagation of `vertices`. """ + if mute: + g4mute() + photons = None - for vertex in vertices: - self.particle_gun.SetParticleByName(vertex.particle_name) - mass = G4ParticleTable.GetParticleTable().FindParticle(vertex.particle_name).GetPDGMass() - total_energy = vertex.ke*MeV + mass - self.particle_gun.SetParticleEnergy(total_energy) - self.particle_gun.SetParticlePosition(G4ThreeVector(*vertex.pos)*mm) - self.particle_gun.SetParticleMomentumDirection(G4ThreeVector(*vertex.dir).unit()) - self.particle_gun.SetParticleTime(vertex.t0*ns) - - if vertex.pol is not None: - self.particle_gun.SetParticlePolarization(G4ThreeVector(*vertex.pol).unit()) - - self.tracking_action.Clear() - gRunManager.BeamOn(1) - - if photons is None: - photons = self._extract_photons_from_tracking_action() - else: - photons += self._extract_photons_from_tracking_action() + try: + for vertex in vertices: + self.particle_gun.SetParticleByName(vertex.particle_name) + mass = G4ParticleTable.GetParticleTable().FindParticle(vertex.particle_name).GetPDGMass() + total_energy = vertex.ke*MeV + mass + self.particle_gun.SetParticleEnergy(total_energy) + self.particle_gun.SetParticlePosition(G4ThreeVector(*vertex.pos)*mm) + self.particle_gun.SetParticleMomentumDirection(G4ThreeVector(*vertex.dir).unit()) + self.particle_gun.SetParticleTime(vertex.t0*ns) + + if vertex.pol is not None: + self.particle_gun.SetParticlePolarization(G4ThreeVector(*vertex.pol).unit()) + + self.tracking_action.Clear() + gRunManager.BeamOn(1) + + if photons is None: + photons = self._extract_photons_from_tracking_action() + else: + photons += self._extract_photons_from_tracking_action() + finally: + if mute: + g4unmute() return photons |