1 files changed, 71 insertions, 0 deletions
diff --git a/generator/vertex.py b/generator/vertex.py
new file mode 100644
index 0000000..1fa61cb
--- /dev/null
+++ b/generator/vertex.py
@@ -0,0 +1,71 @@
+import itertools
+import numpy as np
+from math import sin, cos, sqrt
+import copy
+
+import pi0
+from 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