diff options
Diffstat (limited to 'generator/vertex.py')
-rw-r--r-- | generator/vertex.py | 74 |
1 files changed, 0 insertions, 74 deletions
diff --git a/generator/vertex.py b/generator/vertex.py deleted file mode 100644 index 0ec4b31..0000000 --- a/generator/vertex.py +++ /dev/null @@ -1,74 +0,0 @@ -import numpy as np -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 repeatfunc - -# 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 repeatfunc(np.random.random_sample): - yield h.bincenters[np.searchsorted(cdf, x)] - -def constant(obj): - while True: - yield obj - -def isotropic(): - while True: - yield uniform_sphere() - -def line_segment(point1, point2): - while True: - frac = np.random.uniform(0.0, 1.0) - yield frac * point1 + (1.0 - frac) * point2 - -def fill_shell(center, radius): - for direction in isotropic(): - r = radius * np.random.uniform(0.0, 1.0)**(1.0/3.0) - yield r * direction - -def flat(e_lo, e_hi): - while True: - yield np.random.uniform(e_lo, e_hi) - -# vertex generators - -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, 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, ke, t0=t0) - - 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(ke+134.9766, dir, theta_rest, phi_rest) - - 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]) - - yield ev_vertex - - -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), constant(t0), start_id=start_id) |