Add t0 parameters to the vertex generators that set the start time

of the event.

1 files changed, 11 insertions, 10 deletions

diff --git a/generator/vertex.py b/generator/vertex.py
index 5521d6b..b775318 100644
--- a/generator/vertex.py
+++ b/generator/vertex.py
@@ -40,17 +40,18 @@ def flat(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):
+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, None, ke)
+        vertex = event.Vertex(particle_name, pos, dir, None, ke, t0)
         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):
+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, None, ke)
+        primary_vertex = event.Vertex('pi0', pos, dir, None, ke, t0)
 
         cos_theta_rest = np.random.random_sample() * 2 - 1
         theta_rest = np.arccos(cos_theta_rest)
@@ -59,15 +60,15 @@ def pi0_gun(pos_iter, dir_iter, ke_iter, start_id=0):
         (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)
+        gamma1_vertex = event.Vertex('gamma', pos, gamma1_dir, None, gamma1_e, t0)
+        gamma2_vertex = event.Vertex('gamma', pos, gamma2_dir, None, gamma2_e, t0)
 
         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):
+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), start_id=start_id)
+    return particle_gun(constant(particle_name), constant(pos), constant(dir), constant(ke), constant(t0), start_id=start_id)