diff options
Diffstat (limited to 'generator/g4gen.py')
| -rw-r--r-- | generator/g4gen.py | 136 |
1 files changed, 0 insertions, 136 deletions
diff --git a/generator/g4gen.py b/generator/g4gen.py deleted file mode 100644 index 91ddc18..0000000 --- a/generator/g4gen.py +++ /dev/null @@ -1,136 +0,0 @@ -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 |