diff options
Diffstat (limited to 'generator/vertex.py')
| -rw-r--r-- | generator/vertex.py | 51 |
1 files changed, 21 insertions, 30 deletions
diff --git a/generator/vertex.py b/generator/vertex.py index 450c0e4..5521d6b 100644 --- a/generator/vertex.py +++ b/generator/vertex.py @@ -2,7 +2,7 @@ import numpy as np from itertools import izip, count from chroma.pi0 import pi0_decay -from chroma.event import Event, Subtrack +from chroma import event from chroma.sample import uniform_sphere from chroma.itertoolset import repeat_func @@ -40,43 +40,34 @@ def flat(e_lo, e_hi): # vertex generators -def particle_gun(particle_name, position, direction, total_energy, start_id=0): - for i, ev_particle, ev_position, ev_direction, ev_total_energy in \ - izip(count(start_id), - particle_name, position, direction, total_energy): - ev = Event(event_id=i, particle_name=ev_particle, - gen_position=ev_position, gen_direction=ev_direction, gen_total_energy=ev_total_energy) - yield ev +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): + dir /= np.linalg.norm(dir) + vertex = event.Vertex(particle_name, pos, dir, None, ke) + ev_vertex = event.Event(i, vertex, [vertex]) + yield ev_vertex -def pi0_gun(pi0_position, pi0_direction, pi0_total_energy, start_id=0): - for i, ev_position, ev_direction, ev_total_energy in \ - izip(count(start_id), pi0_position, pi0_direction, pi0_total_energy): +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): + dir /= np.linalg.norm(dir) + primary_vertex = event.Vertex('pi0', pos, dir, None, ke) - ev = Event(event_id=i, particle_name='pi0', - gen_position=ev_position, gen_direction=ev_direction, gen_total_energy=ev_total_energy) - - ev_direction = ev_direction / np.linalg.norm(ev_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_decay(energy=ev_total_energy, - direction=ev_direction, - theta=theta_rest, - phi=phi_rest) + (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) + + ev_vertex = event.Event(i, primary_vertex, [gamma1_vertex, gamma2_vertex]) - ev.subtracks = [Subtrack(particle_name='gamma', position=ev_position, direction=gamma1_dir, - start_time=0.0, total_energy=gamma1_e), - Subtrack(particle_name='gamma', position=ev_position, direction=gamma2_dir, - start_time=0.0, total_energy=gamma2_e)] - yield ev + yield ev_vertex -def constant_particle_gun(particle_name, position, direction, total_energy, - start_id=0): +def constant_particle_gun(particle_name, pos, dir, ke, start_id=0): '''Convenience wrapper around particle gun that assumes all arguments are constants, rather than generators.''' - return particle_gun(constant(particle_name), constant(position), - constant(direction), constant(total_energy), - start_id=start_id) + return particle_gun(constant(particle_name), constant(pos), constant(dir), constant(ke), start_id=start_id) |