diff options
Diffstat (limited to 'chroma/generator')
| -rw-r--r-- | chroma/generator/G4chroma.cc | 230 | ||||
| -rw-r--r-- | chroma/generator/G4chroma.hh | 49 | ||||
| -rw-r--r-- | chroma/generator/__init__.py | 4 | ||||
| -rw-r--r-- | chroma/generator/g4gen.py | 136 | ||||
| -rw-r--r-- | chroma/generator/photon.py | 97 | ||||
| -rw-r--r-- | chroma/generator/vertex.py | 74 |
6 files changed, 590 insertions, 0 deletions
diff --git a/chroma/generator/G4chroma.cc b/chroma/generator/G4chroma.cc new file mode 100644 index 0000000..1aba133 --- /dev/null +++ b/chroma/generator/G4chroma.cc @@ -0,0 +1,230 @@ +#include "G4chroma.hh" +#include <geant4/G4OpticalPhysics.hh> +#include <geant4/G4EmPenelopePhysics.hh> + +ChromaPhysicsList::ChromaPhysicsList(): G4VModularPhysicsList() +{ + // default cut value (1.0mm) + defaultCutValue = 1.0*mm; + + // General Physics + RegisterPhysics( new G4EmPenelopePhysics(0) ); + // Optical Physics + G4OpticalPhysics* opticalPhysics = new G4OpticalPhysics(); + RegisterPhysics( opticalPhysics ); + opticalPhysics->SetTrackSecondariesFirst(false); +} + +ChromaPhysicsList::~ChromaPhysicsList() +{ +} + +void ChromaPhysicsList::SetCuts(){ + // " G4VUserPhysicsList::SetCutsWithDefault" method sets + // the default cut value for all particle types + SetCutsWithDefault(); +} + + +PhotonTrackingAction::PhotonTrackingAction() +{ +} + +PhotonTrackingAction::~PhotonTrackingAction() +{ +} + +int PhotonTrackingAction::GetNumPhotons() const +{ + return pos.size(); +} + +void PhotonTrackingAction::Clear() +{ + pos.clear(); + dir.clear(); + pol.clear(); + wavelength.clear(); + t0.clear(); +} + +void PhotonTrackingAction::GetX(double *x) const +{ + for (unsigned i=0; i < pos.size(); i++) x[i] = pos[i].x(); +} + +void PhotonTrackingAction::GetY(double *y) const +{ + for (unsigned i=0; i < pos.size(); i++) y[i] = pos[i].y(); +} + +void PhotonTrackingAction::GetZ(double *z) const +{ + for (unsigned i=0; i < pos.size(); i++) z[i] = pos[i].z(); +} + +void PhotonTrackingAction::GetDirX(double *dir_x) const +{ + for (unsigned i=0; i < dir.size(); i++) dir_x[i] = dir[i].x(); +} + +void PhotonTrackingAction::GetDirY(double *dir_y) const +{ + for (unsigned i=0; i < dir.size(); i++) dir_y[i] = dir[i].y(); +} + +void PhotonTrackingAction::GetDirZ(double *dir_z) const +{ + for (unsigned i=0; i < dir.size(); i++) dir_z[i] = dir[i].z(); +} + +void PhotonTrackingAction::GetPolX(double *pol_x) const +{ + for (unsigned i=0; i < pol.size(); i++) pol_x[i] = pol[i].x(); +} + +void PhotonTrackingAction::GetPolY(double *pol_y) const +{ + for (unsigned i=0; i < pol.size(); i++) pol_y[i] = pol[i].y(); +} + +void PhotonTrackingAction::GetPolZ(double *pol_z) const +{ + for (unsigned i=0; i < pol.size(); i++) pol_z[i] = pol[i].z(); +} + +void PhotonTrackingAction::GetWavelength(double *wl) const +{ + for (unsigned i=0; i < wavelength.size(); i++) wl[i] = wavelength[i]; +} + +void PhotonTrackingAction::GetT0(double *t) const +{ + for (unsigned i=0; i < t0.size(); i++) t[i] = t0[i]; +} + +void PhotonTrackingAction::PreUserTrackingAction(const G4Track *track) +{ + G4ParticleDefinition *particle = track->GetDefinition(); + if (particle->GetParticleName() == "opticalphoton") { + pos.push_back(track->GetPosition()/m); + dir.push_back(track->GetMomentumDirection()); + pol.push_back(track->GetPolarization()); + wavelength.push_back( (h_Planck * c_light / track->GetKineticEnergy()) / nanometer ); + t0.push_back(track->GetGlobalTime() / s); + const_cast<G4Track *>(track)->SetTrackStatus(fStopAndKill); + } +} + +#include <boost/python.hpp> +#include <pyublas/numpy.hpp> + +using namespace boost::python; + +pyublas::numpy_vector<double> PTA_GetX(const PhotonTrackingAction *pta) +{ + pyublas::numpy_vector<double> r(pta->GetNumPhotons()); + pta->GetX(&r[0]); + return r; +} + +pyublas::numpy_vector<double> PTA_GetY(const PhotonTrackingAction *pta) +{ + pyublas::numpy_vector<double> r(pta->GetNumPhotons()); + pta->GetY(&r[0]); + return r; +} + +pyublas::numpy_vector<double> PTA_GetZ(const PhotonTrackingAction *pta) +{ + pyublas::numpy_vector<double> r(pta->GetNumPhotons()); + pta->GetZ(&r[0]); + return r; +} + +pyublas::numpy_vector<double> PTA_GetDirX(const PhotonTrackingAction *pta) +{ + pyublas::numpy_vector<double> r(pta->GetNumPhotons()); + pta->GetDirX(&r[0]); + return r; +} + +pyublas::numpy_vector<double> PTA_GetDirY(const PhotonTrackingAction *pta) +{ + pyublas::numpy_vector<double> r(pta->GetNumPhotons()); + pta->GetDirY(&r[0]); + return r; +} + +pyublas::numpy_vector<double> PTA_GetDirZ(const PhotonTrackingAction *pta) +{ + pyublas::numpy_vector<double> r(pta->GetNumPhotons()); + pta->GetDirZ(&r[0]); + return r; +} + +pyublas::numpy_vector<double> PTA_GetPolX(const PhotonTrackingAction *pta) +{ + pyublas::numpy_vector<double> r(pta->GetNumPhotons()); + pta->GetPolX(&r[0]); + return r; +} + +pyublas::numpy_vector<double> PTA_GetPolY(const PhotonTrackingAction *pta) +{ + pyublas::numpy_vector<double> r(pta->GetNumPhotons()); + pta->GetPolY(&r[0]); + return r; +} + +pyublas::numpy_vector<double> PTA_GetPolZ(const PhotonTrackingAction *pta) +{ + pyublas::numpy_vector<double> r(pta->GetNumPhotons()); + pta->GetPolZ(&r[0]); + return r; +} + +pyublas::numpy_vector<double> PTA_GetWave(const PhotonTrackingAction *pta) +{ + pyublas::numpy_vector<double> r(pta->GetNumPhotons()); + pta->GetWavelength(&r[0]); + return r; +} + +pyublas::numpy_vector<double> PTA_GetT0(const PhotonTrackingAction *pta) +{ + pyublas::numpy_vector<double> r(pta->GetNumPhotons()); + pta->GetT0(&r[0]); + return r; +} + + +void export_Chroma() +{ + class_<ChromaPhysicsList, ChromaPhysicsList*, bases<G4VModularPhysicsList>, boost::noncopyable > ("ChromaPhysicsList", "EM+Optics physics list") + .def(init<>()) + ; + + class_<PhotonTrackingAction, PhotonTrackingAction*, bases<G4UserTrackingAction>, + boost::noncopyable > ("PhotonTrackingAction", "Tracking action that saves photons") + .def(init<>()) + .def("GetNumPhotons", &PhotonTrackingAction::GetNumPhotons) + .def("Clear", &PhotonTrackingAction::Clear) + .def("GetX", PTA_GetX) + .def("GetY", PTA_GetY) + .def("GetZ", PTA_GetZ) + .def("GetDirX", PTA_GetDirX) + .def("GetDirY", PTA_GetDirY) + .def("GetDirZ", PTA_GetDirZ) + .def("GetPolX", PTA_GetPolX) + .def("GetPolY", PTA_GetPolY) + .def("GetPolZ", PTA_GetPolZ) + .def("GetWavelength", PTA_GetWave) + .def("GetT0", PTA_GetT0) + ; +} + +BOOST_PYTHON_MODULE(G4chroma) +{ + export_Chroma(); +} diff --git a/chroma/generator/G4chroma.hh b/chroma/generator/G4chroma.hh new file mode 100644 index 0000000..4f085aa --- /dev/null +++ b/chroma/generator/G4chroma.hh @@ -0,0 +1,49 @@ +#ifndef __G4chroma_hh__ +#define __G4chroma_hh__ + +#include <geant4/G4VModularPhysicsList.hh> +class ChromaPhysicsList: public G4VModularPhysicsList +{ +public: + ChromaPhysicsList(); + virtual ~ChromaPhysicsList(); + virtual void SetCuts(); +}; + +#include <geant4/G4UserTrackingAction.hh> +#include <vector> +#include <geant4/G4ThreeVector.hh> + +class PhotonTrackingAction : public G4UserTrackingAction +{ +public: + PhotonTrackingAction(); + virtual ~PhotonTrackingAction(); + + int GetNumPhotons() const; + void Clear(); + + void GetX(double *x) const; + void GetY(double *y) const; + void GetZ(double *z) const; + void GetDirX(double *dir_x) const; + void GetDirY(double *dir_y) const; + void GetDirZ(double *dir_z) const; + void GetPolX(double *pol_x) const; + void GetPolY(double *pol_y) const; + void GetPolZ(double *pol_z) const; + + void GetWavelength(double *wl) const; + void GetT0(double *t) const; + + virtual void PreUserTrackingAction(const G4Track *); + +protected: + std::vector<G4ThreeVector> pos; + std::vector<G4ThreeVector> dir; + std::vector<G4ThreeVector> pol; + std::vector<double> wavelength; + std::vector<double> t0; +}; + +#endif 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..91ddc18 --- /dev/null +++ b/chroma/generator/g4gen.py @@ -0,0 +1,136 @@ +from Geant4 import * +import g4py.ezgeom +import g4py.NISTmaterials +import g4py.ParticleGun +import pyublas +import numpy as np +from chroma.event import Photons, Vertex + +try: + import G4chroma +except: + # Try building the module + import subprocess + import sys, os + module_dir = os.path.split(os.path.realpath(__file__))[0] + print >>sys.stderr, 'Compiling G4chroma.so...' + retcode = subprocess.call('g++ -o \'%s/G4chroma.so\' -shared \'%s/G4chroma.cc\' -fPIC `geant4-config --cflags --libs` `python-config --cflags --libs --ldflags` -lboost_python' % (module_dir, module_dir), shell=True) + assert retcode == 0 + import G4chroma + +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) + gRunManager.Initialize() + + # 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..6c656b2 --- /dev/null +++ b/chroma/generator/photon.py @@ -0,0 +1,97 @@ +import g4gen +import multiprocessing +import numpy as np +import itertools +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) |