From 7d9b50e9e64c9d8d9a25942e2ffaca52142c6c2b Mon Sep 17 00:00:00 2001 From: Stan Seibert Date: Tue, 16 Aug 2011 13:52:00 -0400 Subject: Epic restructuring of code to switch to a generator-based style of event creation. Now we have vertex generators (that produce initial particles), photon generators (that create photons to propagate), and a standard data structure using Python class containers and numpy arrays to hand around the code. Also cleaned up some naming of things before they become conventions. --- generator/G4chroma.cc | 230 ++++++++++++++++++++++++++++++++++++++++++++++++++ generator/G4chroma.hh | 49 +++++++++++ generator/__init__.py | 5 ++ generator/g4gen.py | 147 ++++++++++++++++++++++++++++++++ generator/photon.py | 74 ++++++++++++++++ generator/vertex.py | 71 ++++++++++++++++ 6 files changed, 576 insertions(+) create mode 100644 generator/G4chroma.cc create mode 100644 generator/G4chroma.hh create mode 100644 generator/__init__.py create mode 100644 generator/g4gen.py create mode 100644 generator/photon.py create mode 100644 generator/vertex.py (limited to 'generator') diff --git a/generator/G4chroma.cc b/generator/G4chroma.cc new file mode 100644 index 0000000..1aba133 --- /dev/null +++ b/generator/G4chroma.cc @@ -0,0 +1,230 @@ +#include "G4chroma.hh" +#include +#include + +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(track)->SetTrackStatus(fStopAndKill); + } +} + +#include +#include + +using namespace boost::python; + +pyublas::numpy_vector PTA_GetX(const PhotonTrackingAction *pta) +{ + pyublas::numpy_vector r(pta->GetNumPhotons()); + pta->GetX(&r[0]); + return r; +} + +pyublas::numpy_vector PTA_GetY(const PhotonTrackingAction *pta) +{ + pyublas::numpy_vector r(pta->GetNumPhotons()); + pta->GetY(&r[0]); + return r; +} + +pyublas::numpy_vector PTA_GetZ(const PhotonTrackingAction *pta) +{ + pyublas::numpy_vector r(pta->GetNumPhotons()); + pta->GetZ(&r[0]); + return r; +} + +pyublas::numpy_vector PTA_GetDirX(const PhotonTrackingAction *pta) +{ + pyublas::numpy_vector r(pta->GetNumPhotons()); + pta->GetDirX(&r[0]); + return r; +} + +pyublas::numpy_vector PTA_GetDirY(const PhotonTrackingAction *pta) +{ + pyublas::numpy_vector r(pta->GetNumPhotons()); + pta->GetDirY(&r[0]); + return r; +} + +pyublas::numpy_vector PTA_GetDirZ(const PhotonTrackingAction *pta) +{ + pyublas::numpy_vector r(pta->GetNumPhotons()); + pta->GetDirZ(&r[0]); + return r; +} + +pyublas::numpy_vector PTA_GetPolX(const PhotonTrackingAction *pta) +{ + pyublas::numpy_vector r(pta->GetNumPhotons()); + pta->GetPolX(&r[0]); + return r; +} + +pyublas::numpy_vector PTA_GetPolY(const PhotonTrackingAction *pta) +{ + pyublas::numpy_vector r(pta->GetNumPhotons()); + pta->GetPolY(&r[0]); + return r; +} + +pyublas::numpy_vector PTA_GetPolZ(const PhotonTrackingAction *pta) +{ + pyublas::numpy_vector r(pta->GetNumPhotons()); + pta->GetPolZ(&r[0]); + return r; +} + +pyublas::numpy_vector PTA_GetWave(const PhotonTrackingAction *pta) +{ + pyublas::numpy_vector r(pta->GetNumPhotons()); + pta->GetWavelength(&r[0]); + return r; +} + +pyublas::numpy_vector PTA_GetT0(const PhotonTrackingAction *pta) +{ + pyublas::numpy_vector r(pta->GetNumPhotons()); + pta->GetT0(&r[0]); + return r; +} + + +void export_Chroma() +{ + class_, boost::noncopyable > ("ChromaPhysicsList", "EM+Optics physics list") + .def(init<>()) + ; + + class_, + 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/generator/G4chroma.hh b/generator/G4chroma.hh new file mode 100644 index 0000000..4f085aa --- /dev/null +++ b/generator/G4chroma.hh @@ -0,0 +1,49 @@ +#ifndef __G4chroma_hh__ +#define __G4chroma_hh__ + +#include +class ChromaPhysicsList: public G4VModularPhysicsList +{ +public: + ChromaPhysicsList(); + virtual ~ChromaPhysicsList(); + virtual void SetCuts(); +}; + +#include +#include +#include + +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 pos; + std::vector dir; + std::vector pol; + std::vector wavelength; + std::vector t0; +}; + +#endif diff --git a/generator/__init__.py b/generator/__init__.py new file mode 100644 index 0000000..c3179d0 --- /dev/null +++ b/generator/__init__.py @@ -0,0 +1,5 @@ +from vertex import * +from event import * +from photon import * + + diff --git a/generator/g4gen.py b/generator/g4gen.py new file mode 100644 index 0000000..d4e79a1 --- /dev/null +++ b/generator/g4gen.py @@ -0,0 +1,147 @@ +from Geant4 import * +import g4py.ezgeom +import g4py.NISTmaterials +import g4py.ParticleGun +import pyublas +import numpy as np +import event + +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: 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) + + 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() + + 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) + times = self.tracking_action.GetT0().astype(np.float32) + + return event.Photons(positions=pos, directions=dir, polarizations=pol, times=times, wavelengths=wavelengths) + + def generate_photons(self, ev): + '''Use GEANT4 to generate photons produced by the given particle. + + 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_pos, + direction=ev.gen_dir, + 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()) + + 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) + # prime things + gen.generate_photons(event.Event('e-', (0,0,0), (1,0,0), 1.0)) + + 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' diff --git a/generator/photon.py b/generator/photon.py new file mode 100644 index 0000000..dcfc292 --- /dev/null +++ b/generator/photon.py @@ -0,0 +1,74 @@ +import g4gen +import multiprocessing +import threading +import numpy as np +import itertools + +class G4GeneratorProcess(multiprocessing.Process): + def __init__(self, idnum, material, input_queue, output_queue, seed=None): + multiprocessing.Process.__init__(self) + + self.idnum = idnum + self.material = material + self.input_queue = input_queue + self.output_queue = output_queue + self.seed = seed + self.daemon = True + + def run(self): + gen = g4gen.G4Generator(self.material, seed=self.seed) + while True: + ev = self.input_queue.get() + ev.photon_start = gen.generate_photons(ev) + self.output_queue.put(ev) + +class G4VertexGeneratorProcess(multiprocessing.Process): + def __init__(self, vertex_iterator, queue): + multiprocessing.Process.__init__(self) + self.vertex_iterator = vertex_iterator + self.queue = queue + + def run(self): + for ev in self.vertex_iterator: + self.queue.put(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 + ''' + step = num // partitions + for i in xrange(partitions): + if i < partitions - 1: + yield step + else: + yield step + (num % partitions) + +def queue_iterator(nelements, queue): + for i in xrange(nelements): + yield queue.get() + +class G4ParallelGenerator(object): + def __init__(self, nprocesses, material, base_seed=None): + self.material = material + self.vertex_queue = multiprocessing.Queue() + self.photon_queue = multiprocessing.Queue() + self.processes = [ G4GeneratorProcess(i, material, self.vertex_queue, self.photon_queue, seed=base_seed + i) + for i in xrange(nprocesses) ] + for p in self.processes: + p.start() + + def generate_events(self, nevents, vertex_iterator): + # Doing this to avoid a deadlock caused by putting to one queue while getting from another + for ev in itertools.islice(vertex_iterator, nevents): + self.vertex_queue.put(ev) + + return queue_iterator(nevents, self.photon_queue) + + + diff --git a/generator/vertex.py b/generator/vertex.py new file mode 100644 index 0000000..1fa61cb --- /dev/null +++ b/generator/vertex.py @@ -0,0 +1,71 @@ +import itertools +import numpy as np +from math import sin, cos, sqrt +import copy + +import pi0 +from event import Event, Subtrack + +# generator parts for use with gun() + +def constant(obj): + while True: + yield obj + +def isotropic(): + while True: + cos_theta = np.random.uniform(-1.0, 1.0) + sin_theta = sqrt(1.0 - cos_theta**2) + phi = np.random.uniform(0.0, 2*np.pi) + direction = np.array([sin_theta * cos(phi), + sin_theta * sin(phi), + cos_theta]) + yield direction + +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, position, direction, total_energy, start_id=0): + for i, ev_particle, ev_position, ev_direction, ev_total_energy in \ + itertools.izip(itertools.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 pi0_gun(pi0_position, pi0_direction, pi0_total_energy, start_id=0): + for i, ev_position, ev_direction, ev_total_energy in \ + itertools.izip(itertools.count(start_id), pi0_position, pi0_direction, pi0_total_energy): + + 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.pi0_decay(energy=ev_total_energy, + direction=ev_direction, + theta=theta_rest, + phi=phi_rest) + + 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 -- cgit