diff options
Diffstat (limited to 'generator/vertex.py')
| -rw-r--r-- | generator/vertex.py | 75 |
1 files changed, 34 insertions, 41 deletions
diff --git a/generator/vertex.py b/generator/vertex.py index 139d675..5521d6b 100644 --- a/generator/vertex.py +++ b/generator/vertex.py @@ -1,26 +1,28 @@ -import itertools import numpy as np -from math import sin, cos, sqrt -import copy +from itertools import izip, count -import chroma.pi0 as pi0 -from chroma.event import Event, Subtrack +from chroma.pi0 import pi0_decay +from chroma import event +from chroma.sample import uniform_sphere +from chroma.itertoolset import repeat_func # generator parts for use with gun() +def from_histogram(h): + "Yield values drawn randomly from the histogram `h` interpreted as a pdf." + pdf = h.hist/h.hist.sum() + cdf = np.cumsum(pdf) + + for x in repeat_func(np.random.random_sample): + yield h.bincenters[np.searchsorted(cdf, x)] + def constant(obj): while True: yield obj def isotropic(): while True: - cos_theta = np.random.uniform(-1.0, 1.0) - sin_theta = sqrt(1.0 - cos_theta**2) - phi = np.random.uniform(0.0, 2*np.pi) - direction = np.array([sin_theta * cos(phi), - sin_theta * sin(phi), - cos_theta]) - yield direction + yield uniform_sphere() def line_segment(point1, point2): while True: @@ -38,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 \ - itertools.izip(itertools.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 pi0_gun(pi0_position, pi0_direction, pi0_total_energy, start_id=0): - for i, ev_position, ev_direction, ev_total_energy in \ - itertools.izip(itertools.count(start_id), pi0_position, pi0_direction, pi0_total_energy): +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 - ev = Event(event_id=i, particle_name='pi0', - gen_position=ev_position, gen_direction=ev_direction, gen_total_energy=ev_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_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.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) |