From 6f0703602270d03f4025221f13fee21aa842a863 Mon Sep 17 00:00:00 2001 From: Anthony LaTorre Date: Mon, 12 Sep 2011 16:41:17 -0400 Subject: add chroma-sim script. get rid of GPU class; contexts should now be created with create_cuda_context(). --- generator/g4gen.py | 4 +++- generator/vertex.py | 12 ++++++------ 2 files changed, 9 insertions(+), 7 deletions(-) (limited to 'generator') diff --git a/generator/g4gen.py b/generator/g4gen.py index 8dca086..3836897 100644 --- a/generator/g4gen.py +++ b/generator/g4gen.py @@ -120,7 +120,9 @@ class G4Generator(object): self.particle_gun.SetParticleEnergy(total_energy) self.particle_gun.SetParticlePosition(G4ThreeVector(*vertex.pos)*m) self.particle_gun.SetParticleMomentumDirection(G4ThreeVector(*vertex.dir).unit()) - + 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 5521d6b..b4e2d23 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): @@ -43,14 +43,14 @@ def flat(e_lo, e_hi): 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): dir /= np.linalg.norm(dir) - vertex = event.Vertex(particle_name, pos, dir, None, ke) + vertex = event.Vertex(particle_name, pos, dir, ke) 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): dir /= np.linalg.norm(dir) - primary_vertex = event.Vertex('pi0', pos, dir, None, ke) + primary_vertex = event.Vertex('pi0', pos, dir, ke) cos_theta_rest = np.random.random_sample() * 2 - 1 theta_rest = np.arccos(cos_theta_rest) @@ -59,8 +59,8 @@ 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, gamma1_e) + gamma2_vertex = event.Vertex('gamma', pos, gamma2_dir, gamma2_e) ev_vertex = event.Event(i, primary_vertex, [gamma1_vertex, gamma2_vertex]) -- cgit