diff options
Diffstat (limited to 'generator')
| -rw-r--r-- | generator/g4gen.py | 96 | ||||
| -rw-r--r-- | generator/photon.py | 52 | ||||
| -rw-r--r-- | generator/vertex.py | 51 |
3 files changed, 86 insertions, 113 deletions
diff --git a/generator/g4gen.py b/generator/g4gen.py index 9650b68..8dca086 100644 --- a/generator/g4gen.py +++ b/generator/g4gen.py @@ -4,7 +4,7 @@ import g4py.NISTmaterials import g4py.ParticleGun import pyublas import numpy as np -import chroma.event as event +from chroma.event import Photons, Vertex try: import G4chroma @@ -20,16 +20,18 @@ except: class G4Generator(object): def __init__(self, material, seed=None): - '''Create generator to produce photons inside the specified material. + """Create generator to produce photons inside the specified material. material: chroma.geometry.Material object with density, composition dict and refractive_index. composition dictionary should be { element_symbol : fraction_by_weight, ... }. - seed: Random number generator seed for HepRandom. If None, - generator is not seeded. - ''' + + seed: int, *optional* + Random number generator seed for HepRandom. If None, generator + is not seeded. + """ if seed is not None: HepRandom.setTheSeed(seed) @@ -51,11 +53,8 @@ class G4Generator(object): gRunManager.SetUserAction(self.tracking_action) gRunManager.Initialize() - #preinitialize the process by running a simple event - self.generate_photons(event.Event(event_id=0, particle_name='e-', - gen_position=(0,0,0), - gen_direction=(1,0,0), - gen_total_energy=1.0)) + # preinitialize the process by running a simple event + self.generate_photons([Vertex('e-', (0,0,0), (1,0,0), 0, 1.0)]) def create_g4material(self, material): g4material = G4Material('world_material', material.density * g / cm3, @@ -96,57 +95,38 @@ class G4Generator(object): pol[:,2] = self.tracking_action.GetPolZ() wavelengths = self.tracking_action.GetWavelength().astype(np.float32) - times = self.tracking_action.GetT0().astype(np.float32) - return event.Photons(positions=pos, directions=dir, polarizations=pol, times=times, wavelengths=wavelengths) + t0 = self.tracking_action.GetT0().astype(np.float32) + + return Photons(pos, dir, pol, wavelengths, t0) - def generate_photons(self, ev): - '''Use GEANT4 to generate photons produced by the given particle. + def generate_photons(self, vertices): + """Use GEANT4 to generate photons produced by propagating `vertices`. - ev: a generator.event.Event object with the particle - properties set. If it contains subtracks, those - will be used to create the photon vertices rather - than the main particle. - - Returns an instance of event.Photons containing the - generated photon vertices for the primary particle or - all the subtracks, if present. - ''' - photons = [] - if ev.subtracks: - subtracks = ev.subtracks - else: - # Create temporary subtrack for single primary particle - subtracks = [event.Subtrack(particle_name=ev.particle_name, - position=ev.gen_position, - direction=ev.gen_direction, - start_time=0.0, - total_energy=ev.gen_total_energy)] - - for subtrack in subtracks: - self.particle_gun.SetParticleByName(subtrack.particle_name) - self.particle_gun.SetParticleEnergy(subtrack.total_energy * MeV) - self.particle_gun.SetParticlePosition(G4ThreeVector(*subtrack.position)*m) - self.particle_gun.SetParticleMomentumDirection(G4ThreeVector(*subtrack.direction).unit()) + Args: + vertices: list of event.Vertex objects + List of initial particle vertices. + + Returns: + photons: event.Photons + Photon vertices generated by the propagation of `vertices`. + """ + photons = None + + for vertex in vertices: + self.particle_gun.SetParticleByName(vertex.particle_name) + mass = G4ParticleTable.GetParticleTable().FindParticle(vertex.particle_name).GetPDGMass() + total_energy = vertex.ke*MeV + mass + self.particle_gun.SetParticleEnergy(total_energy) + self.particle_gun.SetParticlePosition(G4ThreeVector(*vertex.pos)*m) + self.particle_gun.SetParticleMomentumDirection(G4ThreeVector(*vertex.dir).unit()) self.tracking_action.Clear() gRunManager.BeamOn(1) - photons.append(self._extract_photons_from_tracking_action()) - - # Merge all photon lists into one big list - return event.concatenate_photons(photons) - -if __name__ == '__main__': - import time - import optics - gen = G4Generator(optics.water) - - start = time.time() - n = 0 - for i in xrange(100): - photons = gen.generate_photons(event.Event('mu-', (0,0,0), (1,0,0), 1.0)) - n += len(photons.times) - print photons.positions[0].min(), photons.positions[0].max() - stop = time.time() - print stop - start, 'sec' - print n / (stop-start), 'photons/sec' + + if photons is None: + photons = self._extract_photons_from_tracking_action() + else: + photons += self._extract_photons_from_tracking_action() + + return photons diff --git a/generator/photon.py b/generator/photon.py index e656428..6c656b2 100644 --- a/generator/photon.py +++ b/generator/photon.py @@ -25,24 +25,29 @@ class G4GeneratorProcess(multiprocessing.Process): photon_socket = context.socket(zmq.PUSH) photon_socket.connect(self.photon_socket_address) - # Signal with the photon socket that we are online and ready for messages + # Signal with the photon socket that we are online + # and ready for messages. photon_socket.send('READY') while True: ev = vertex_socket.recv_pyobj() - ev.photon_start = gen.generate_photons(ev) + ev.photons_beg = gen.generate_photons(ev.vertices) + #print 'type(ev.photons_beg) is %s' % type(ev.photons_beg) photon_socket.send_pyobj(ev) def partition(num, partitions): - '''Generator that returns num//partitions, with the last item including the remainder. - - Useful for partitioning a number into mostly equal parts while preserving the sum. - - >>> list(partition(800, 3)) - [266, 266, 268] - >>> sum(list(partition(800, 3))) - 800 - ''' + """Generator that returns num//partitions, with the last item including + the remainder. + + Useful for partitioning a number into mostly equal parts while preserving + the sum. + + Examples: + >>> list(partition(800, 3)) + [266, 266, 268] + >>> sum(list(partition(800, 3))) + 800 + """ step = num // partitions for i in xrange(partitions): if i < partitions - 1: @@ -51,8 +56,8 @@ def partition(num, partitions): yield step + (num % partitions) def vertex_sender(vertex_iterator, vertex_socket): - for ev in vertex_iterator: - vertex_socket.send_pyobj(ev) + for vertex in vertex_iterator: + vertex_socket.send_pyobj(vertex) def socket_iterator(nelements, socket): for i in xrange(nelements): @@ -66,8 +71,7 @@ class G4ParallelGenerator(object): base_address = 'ipc:///tmp/chroma_'+str(uuid.uuid4()) self.vertex_address = base_address + '.vertex' self.photon_address = base_address + '.photon' - self.processes = [ G4GeneratorProcess(i, material, self.vertex_address, self.photon_address, seed=base_seed + i) - for i in xrange(nprocesses) ] + self.processes = [ G4GeneratorProcess(i, material, self.vertex_address, self.photon_address, seed=base_seed + i) for i in xrange(nprocesses) ] for p in self.processes: p.start() @@ -78,18 +82,16 @@ class G4ParallelGenerator(object): self.photon_socket = self.zmq_context.socket(zmq.PULL) self.photon_socket.bind(self.photon_address) - # Verify everyone is running and connected to avoid sending all the events to one client + # Verify everyone is running and connected to avoid + # sending all the events to one client. for i in xrange(nprocesses): msg = self.photon_socket.recv() assert msg == 'READY' - def generate_events(self, nevents, vertex_iterator): - # Doing this to avoid a deadlock caused by putting to one queue while getting from another - limited_iterator = itertools.islice(vertex_iterator, nevents) - sender_thread = threading.Thread(target=vertex_sender, args=(limited_iterator, - self.vertex_socket)) + def generate_events(self, vertex_iterator): + # Doing this to avoid a deadlock caused by putting to one queue + # while getting from another. + vertex_list = list(vertex_iterator) + sender_thread = threading.Thread(target=vertex_sender, args=(vertex_list, self.vertex_socket)) sender_thread.start() - return socket_iterator(nevents, self.photon_socket) - - - + return socket_iterator(len(vertex_list), self.photon_socket) diff --git a/generator/vertex.py b/generator/vertex.py index 450c0e4..5521d6b 100644 --- a/generator/vertex.py +++ b/generator/vertex.py @@ -2,7 +2,7 @@ import numpy as np from itertools import izip, count from chroma.pi0 import pi0_decay -from chroma.event import Event, Subtrack +from chroma import event from chroma.sample import uniform_sphere from chroma.itertoolset import repeat_func @@ -40,43 +40,34 @@ def flat(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 \ - izip(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 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) + ev_vertex = event.Event(i, vertex, [vertex]) + yield ev_vertex -def pi0_gun(pi0_position, pi0_direction, pi0_total_energy, start_id=0): - for i, ev_position, ev_direction, ev_total_energy in \ - izip(count(start_id), pi0_position, pi0_direction, pi0_total_energy): +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) - 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_decay(energy=ev_total_energy, - direction=ev_direction, - theta=theta_rest, - phi=phi_rest) + (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) + + ev_vertex = event.Event(i, primary_vertex, [gamma1_vertex, gamma2_vertex]) - 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 + yield ev_vertex -def constant_particle_gun(particle_name, position, direction, total_energy, - start_id=0): +def constant_particle_gun(particle_name, pos, dir, ke, 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) + return particle_gun(constant(particle_name), constant(pos), constant(dir), constant(ke), start_id=start_id) |