diff options
| -rw-r--r-- | generator/g4gen.py | 1 | ||||
| -rw-r--r-- | generator/vertex.py | 21 |
2 files changed, 12 insertions, 10 deletions
diff --git a/generator/g4gen.py b/generator/g4gen.py index 8dca086..7266456 100644 --- a/generator/g4gen.py +++ b/generator/g4gen.py @@ -120,6 +120,7 @@ class G4Generator(object): self.particle_gun.SetParticleEnergy(total_energy) self.particle_gun.SetParticlePosition(G4ThreeVector(*vertex.pos)*m) self.particle_gun.SetParticleMomentumDirection(G4ThreeVector(*vertex.dir).unit()) + self.particle_gun.SetParticleTime(vertex.t0*s) self.tracking_action.Clear() gRunManager.BeamOn(1) diff --git a/generator/vertex.py b/generator/vertex.py index 5521d6b..b775318 100644 --- a/generator/vertex.py +++ b/generator/vertex.py @@ -40,17 +40,18 @@ def flat(e_lo, e_hi): # vertex generators -def particle_gun(particle_name_iter, pos_iter, dir_iter, ke_iter, start_id=0): - for i, particle_name, pos, dir, ke in izip(count(start_id), particle_name_iter, pos_iter, dir_iter, ke_iter): +def particle_gun(particle_name_iter, pos_iter, dir_iter, ke_iter, + t0_iter=constant(0.0), start_id=0): + for i, particle_name, pos, dir, ke, t0 in izip(count(start_id), particle_name_iter, pos_iter, dir_iter, ke_iter, t0_iter): dir /= np.linalg.norm(dir) - vertex = event.Vertex(particle_name, pos, dir, None, ke) + vertex = event.Vertex(particle_name, pos, dir, None, ke, t0) ev_vertex = event.Event(i, vertex, [vertex]) yield ev_vertex -def pi0_gun(pos_iter, dir_iter, ke_iter, start_id=0): - for i, pos, dir, ke in izip(count(start_id), pos_iter, dir_iter, ke_iter): +def pi0_gun(pos_iter, dir_iter, ke_iter, t0_iter=constant(0.0), start_id=0): + for i, pos, dir, ke, t0 in izip(count(start_id), pos_iter, dir_iter, ke_iter, t0_iter): dir /= np.linalg.norm(dir) - primary_vertex = event.Vertex('pi0', pos, dir, None, ke) + primary_vertex = event.Vertex('pi0', pos, dir, None, ke, t0) cos_theta_rest = np.random.random_sample() * 2 - 1 theta_rest = np.arccos(cos_theta_rest) @@ -59,15 +60,15 @@ def pi0_gun(pos_iter, dir_iter, ke_iter, start_id=0): (gamma1_e, gamma1_dir), (gamma2_e, gamma2_dir) = \ pi0_decay(ke+134.9766, dir, theta_rest, phi_rest) - gamma1_vertex = event.Vertex('gamma', pos, gamma1_dir, None, gamma1_e) - gamma2_vertex = event.Vertex('gamma', pos, gamma2_dir, None, gamma2_e) + gamma1_vertex = event.Vertex('gamma', pos, gamma1_dir, None, gamma1_e, t0) + gamma2_vertex = event.Vertex('gamma', pos, gamma2_dir, None, gamma2_e, t0) ev_vertex = event.Event(i, primary_vertex, [gamma1_vertex, gamma2_vertex]) yield ev_vertex -def constant_particle_gun(particle_name, pos, dir, ke, start_id=0): +def constant_particle_gun(particle_name, pos, dir, ke, t0=0.0, start_id=0): '''Convenience wrapper around particle gun that assumes all arguments are constants, rather than generators.''' - return particle_gun(constant(particle_name), constant(pos), constant(dir), constant(ke), start_id=start_id) + return particle_gun(constant(particle_name), constant(pos), constant(dir), constant(ke), constant(t0), start_id=start_id) |