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.event import Event, Subtrack
from chroma.sample import uniform_sphere
# generator parts for use with gun()
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, 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 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):
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)
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
def constant_particle_gun(particle_name, position, direction, total_energy,
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)
|
