diff options
Diffstat (limited to 'chroma/generator')
| -rw-r--r-- | chroma/generator/__init__.py | 4 | ||||
| -rw-r--r-- | chroma/generator/g4gen.py | 134 | ||||
| -rw-r--r-- | chroma/generator/photon.py | 96 | ||||
| -rw-r--r-- | chroma/generator/vertex.py | 74 |
4 files changed, 308 insertions, 0 deletions
diff --git a/chroma/generator/__init__.py b/chroma/generator/__init__.py new file mode 100644 index 0000000..5dd5eca --- /dev/null +++ b/chroma/generator/__init__.py @@ -0,0 +1,4 @@ +from vertex import * +from photon import * + + diff --git a/chroma/generator/g4gen.py b/chroma/generator/g4gen.py new file mode 100644 index 0000000..0ed5433 --- /dev/null +++ b/chroma/generator/g4gen.py @@ -0,0 +1,134 @@ +from chroma.generator.mute import * + +import pyublas +import numpy as np +from chroma.event import Photons, Vertex + +g4mute() +from Geant4 import * +g4unmute() +import g4py.ezgeom +import g4py.NISTmaterials +import g4py.ParticleGun +from chroma.generator import _g4chroma + +#cppmute() +#cppunmute() + +class G4Generator(object): + def __init__(self, material, seed=None): + """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: int, *optional* + Random number generator seed for HepRandom. If None, generator + is not seeded. + """ + if seed is not None: + HepRandom.setTheSeed(seed) + + g4py.NISTmaterials.Construct() + g4py.ezgeom.Construct() + self.physics_list = _g4chroma.ChromaPhysicsList() + gRunManager.SetUserInitialization(self.physics_list) + self.particle_gun = g4py.ParticleGun.Construct() + + self.world_material = self.create_g4material(material) + g4py.ezgeom.SetWorldMaterial(self.world_material) + g4py.ezgeom.ResizeWorld(100*m, 100*m, 100*m) + + self.world = g4py.ezgeom.G4EzVolume('world') + self.world.CreateBoxVolume(self.world_material, 100*m, 100*m, 100*m) + self.world.PlaceIt(G4ThreeVector(0,0,0)) + + self.tracking_action = _g4chroma.PhotonTrackingAction() + gRunManager.SetUserAction(self.tracking_action) + g4mute() + gRunManager.Initialize() + g4unmute() + # 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, + len(material.composition)) + + # Add elements + for element_name, element_frac_by_weight in material.composition.items(): + g4material.AddElement(G4Element.GetElement(element_name, True), + element_frac_by_weight) + + # Set index of refraction + prop_table = G4MaterialPropertiesTable() + # Reverse entries so they are in ascending energy order rather + # than wavelength + energy = list((2*pi*hbarc / (material.refractive_index[::-1,0] * nanometer)).astype(float)) + values = list(material.refractive_index[::-1, 1].astype(float)) + prop_table.AddProperty('RINDEX', energy, values) + + # Load properties + g4material.SetMaterialPropertiesTable(prop_table) + return g4material + + def _extract_photons_from_tracking_action(self): + n = self.tracking_action.GetNumPhotons() + pos = np.zeros(shape=(n,3), dtype=np.float32) + pos[:,0] = self.tracking_action.GetX() + pos[:,1] = self.tracking_action.GetY() + pos[:,2] = self.tracking_action.GetZ() + + dir = np.zeros(shape=(n,3), dtype=np.float32) + dir[:,0] = self.tracking_action.GetDirX() + dir[:,1] = self.tracking_action.GetDirY() + dir[:,2] = self.tracking_action.GetDirZ() + + pol = np.zeros(shape=(n,3), dtype=np.float32) + pol[:,0] = self.tracking_action.GetPolX() + pol[:,1] = self.tracking_action.GetPolY() + pol[:,2] = self.tracking_action.GetPolZ() + + wavelengths = self.tracking_action.GetWavelength().astype(np.float32) + + t0 = self.tracking_action.GetT0().astype(np.float32) + + return Photons(pos, dir, pol, wavelengths, t0) + + def generate_photons(self, vertices): + """Use GEANT4 to generate photons produced by propagating `vertices`. + + 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.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) + + if photons is None: + photons = self._extract_photons_from_tracking_action() + else: + photons += self._extract_photons_from_tracking_action() + + return photons diff --git a/chroma/generator/photon.py b/chroma/generator/photon.py new file mode 100644 index 0000000..39e8cf4 --- /dev/null +++ b/chroma/generator/photon.py @@ -0,0 +1,96 @@ +import g4gen +import multiprocessing +import numpy as np +import threading +import zmq +import uuid + +class G4GeneratorProcess(multiprocessing.Process): + def __init__(self, idnum, material, vertex_socket_address, photon_socket_address, seed=None): + multiprocessing.Process.__init__(self) + + self.idnum = idnum + self.material = material + self.vertex_socket_address = vertex_socket_address + self.photon_socket_address = photon_socket_address + self.seed = seed + self.daemon = True + + def run(self): + gen = g4gen.G4Generator(self.material, seed=self.seed) + context = zmq.Context() + vertex_socket = context.socket(zmq.PULL) + vertex_socket.connect(self.vertex_socket_address) + 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. + photon_socket.send('READY') + + while True: + ev = vertex_socket.recv_pyobj() + 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. + + 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: + yield step + else: + yield step + (num % partitions) + +def vertex_sender(vertex_iterator, vertex_socket): + for vertex in vertex_iterator: + vertex_socket.send_pyobj(vertex) + +def socket_iterator(nelements, socket): + for i in xrange(nelements): + yield socket.recv_pyobj() + +class G4ParallelGenerator(object): + def __init__(self, nprocesses, material, base_seed=None): + self.material = material + if base_seed is None: + base_seed = np.random.randint(100000000) + 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) ] + + for p in self.processes: + p.start() + + self.zmq_context = zmq.Context() + self.vertex_socket = self.zmq_context.socket(zmq.PUSH) + self.vertex_socket.bind(self.vertex_address) + 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. + for i in xrange(nprocesses): + msg = self.photon_socket.recv() + assert msg == 'READY' + + 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(len(vertex_list), self.photon_socket) diff --git a/chroma/generator/vertex.py b/chroma/generator/vertex.py new file mode 100644 index 0000000..0ec4b31 --- /dev/null +++ b/chroma/generator/vertex.py @@ -0,0 +1,74 @@ +import numpy as np +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 repeatfunc + +# generator parts for use with gun() + +def from_histogram(h): + "Yield values drawn randomly from the histogram `h` interpreted as a pdf." + pdf = h.hist/h.hist.sum() + cdf = np.cumsum(pdf) + + for x in repeatfunc(np.random.random_sample): + yield h.bincenters[np.searchsorted(cdf, x)] + +def constant(obj): + while True: + yield obj + +def isotropic(): + while True: + yield uniform_sphere() + +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_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, 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, ke, t0=t0) + + 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(ke+134.9766, dir, theta_rest, phi_rest) + + 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]) + + yield ev_vertex + + +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), constant(t0), start_id=start_id) |