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
74
75
76
77
78
79
80
|
import itertools
import numpy as np
from math import sin, cos, sqrt
import copy
import chroma.pi0 as pi0
from chroma.event import Event, Subtrack
# generator parts for use with gun()
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
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 \
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):
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.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)
|
