diff options
Diffstat (limited to 'generator')
| -rw-r--r-- | generator/g4gen.py | 3 | ||||
| -rw-r--r-- | generator/vertex.py | 12 |
2 files changed, 9 insertions, 6 deletions
diff --git a/generator/g4gen.py b/generator/g4gen.py index 7266456..91ddc18 100644 --- a/generator/g4gen.py +++ b/generator/g4gen.py @@ -122,6 +122,9 @@ class G4Generator(object): self.particle_gun.SetParticleMomentumDirection(G4ThreeVector(*vertex.dir).unit()) self.particle_gun.SetParticleTime(vertex.t0*s) + if vertex.pol is not None: + self.particle_gun.SetParticlePolarization(G4ThreeVector(*vertex.pol).unit()) + self.tracking_action.Clear() gRunManager.BeamOn(1) diff --git a/generator/vertex.py b/generator/vertex.py index b775318..0ec4b31 100644 --- a/generator/vertex.py +++ b/generator/vertex.py @@ -4,7 +4,7 @@ from itertools import izip, count from chroma.pi0 import pi0_decay from chroma import event from chroma.sample import uniform_sphere -from chroma.itertoolset import repeat_func +from chroma.itertoolset import repeatfunc # generator parts for use with gun() @@ -13,7 +13,7 @@ def from_histogram(h): pdf = h.hist/h.hist.sum() cdf = np.cumsum(pdf) - for x in repeat_func(np.random.random_sample): + for x in repeatfunc(np.random.random_sample): yield h.bincenters[np.searchsorted(cdf, x)] def constant(obj): @@ -44,14 +44,14 @@ 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, t0) + vertex = event.Vertex(particle_name, pos, dir, ke, t0=t0) ev_vertex = event.Event(i, vertex, [vertex]) yield ev_vertex 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, t0) + primary_vertex = event.Vertex('pi0', pos, dir, ke, t0=t0) cos_theta_rest = np.random.random_sample() * 2 - 1 theta_rest = np.arccos(cos_theta_rest) @@ -60,8 +60,8 @@ def pi0_gun(pos_iter, dir_iter, ke_iter, t0_iter=constant(0.0), 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, t0) - gamma2_vertex = event.Vertex('gamma', pos, gamma2_dir, None, gamma2_e, t0) + gamma1_vertex = event.Vertex('gamma', pos, gamma1_dir, gamma1_e, t0=t0) + gamma2_vertex = event.Vertex('gamma', pos, gamma2_dir, gamma2_e, t0=t0) ev_vertex = event.Event(i, primary_vertex, [gamma1_vertex, gamma2_vertex]) |