merge

2 files changed, 9 insertions, 6 deletions

diff --git a/generator/g4gen.py b/generator/g4gen.py
index 7266456..91ddc18 100644
--- a/generator/g4gen.py
+++ b/generator/g4gen.py
@@ -122,6 +122,9 @@ class G4Generator(object):
             self.particle_gun.SetParticleMomentumDirection(G4ThreeVector(*vertex.dir).unit())
             self.particle_gun.SetParticleTime(vertex.t0*s)
             
+            if vertex.pol is not None:
+                self.particle_gun.SetParticlePolarization(G4ThreeVector(*vertex.pol).unit())
+
             self.tracking_action.Clear()
             gRunManager.BeamOn(1)
 
diff --git a/generator/vertex.py b/generator/vertex.py
index b775318..0ec4b31 100644
--- a/generator/vertex.py
+++ b/generator/vertex.py
@@ -4,7 +4,7 @@ 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
+from chroma.itertoolset import repeatfunc
 
 # generator parts for use with gun()
 
@@ -13,7 +13,7 @@ def from_histogram(h):
     pdf = h.hist/h.hist.sum()
     cdf = np.cumsum(pdf)
 
-    for x in repeat_func(np.random.random_sample):
+    for x in repeatfunc(np.random.random_sample):
         yield h.bincenters[np.searchsorted(cdf, x)]
 
 def constant(obj):
@@ -44,14 +44,14 @@ 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, t0)
+        vertex = event.Vertex(particle_name, pos, dir, ke, t0=t0)
         ev_vertex = event.Event(i, vertex, [vertex])
         yield ev_vertex
 
 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, t0)
+        primary_vertex = event.Vertex('pi0', pos, dir, ke, t0=t0)
 
         cos_theta_rest = np.random.random_sample() * 2 - 1
         theta_rest = np.arccos(cos_theta_rest)
@@ -60,8 +60,8 @@ def pi0_gun(pos_iter, dir_iter, ke_iter, t0_iter=constant(0.0), 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, t0)
-        gamma2_vertex = event.Vertex('gamma', pos, gamma2_dir, None, gamma2_e, t0)
+        gamma1_vertex = event.Vertex('gamma', pos, gamma1_dir, gamma1_e, t0=t0)
+        gamma2_vertex = event.Vertex('gamma', pos, gamma2_dir, gamma2_e, t0=t0)
 
         ev_vertex = event.Event(i, primary_vertex, [gamma1_vertex, gamma2_vertex])