1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
|
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 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:
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, 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(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)
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, 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])
yield ev_vertex
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(pos), constant(dir), constant(ke), start_id=start_id)
|
