diff options
-rw-r--r-- | G4chroma.cc | 230 | ||||
-rw-r--r-- | G4chroma.hh | 49 | ||||
-rw-r--r-- | g4gen.py | 125 | ||||
-rw-r--r-- | geometry.py | 2 | ||||
-rw-r--r-- | optics.py | 2 |
5 files changed, 408 insertions, 0 deletions
diff --git a/G4chroma.cc b/G4chroma.cc new file mode 100644 index 0000000..580aa63 --- /dev/null +++ b/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 ); + // 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/G4chroma.hh b/G4chroma.hh new file mode 100644 index 0000000..4f085aa --- /dev/null +++ b/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/g4gen.py b/g4gen.py new file mode 100644 index 0000000..409570f --- /dev/null +++ b/g4gen.py @@ -0,0 +1,125 @@ +from Geant4 import * +import g4py.ezgeom +import g4py.NISTmaterials +import g4py.ParticleGun +import pyublas +import numpy + +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): + '''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, ... }. + ''' + 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 generate_photons(self, particle_name, total_energy, position, direction): + '''Use GEANT4 to generate photons produced by the given particle. + + particle_name: GEANT4 name of particle. 'e-', 'mu-', etc + total_energy: Total energy of particle (incl rest mass) in MeV + position: 3-tuple of position of particle in global coordinates + direction: 3-tuple direction vector. + Does not have to be normalized. + ''' + self.particle_gun.SetParticleByName(particle_name) + self.particle_gun.SetParticleEnergy(total_energy * MeV) + self.particle_gun.SetParticlePosition(G4ThreeVector(*position)) + self.particle_gun.SetParticleMomentumDirection(G4ThreeVector(*direction).unit()) + + self.tracking_action.Clear() + + gRunManager.BeamOn(1) + n = self.tracking_action.GetNumPhotons() + pos = numpy.zeros(shape=(n,3), dtype=numpy.float32) + pos[:,0] = self.tracking_action.GetX() + pos[:,1] = self.tracking_action.GetY() + pos[:,2] = self.tracking_action.GetZ() + + dir = numpy.zeros(shape=(n,3), dtype=numpy.float32) + dir[:,0] = self.tracking_action.GetDirX() + dir[:,1] = self.tracking_action.GetDirY() + dir[:,2] = self.tracking_action.GetDirZ() + + pol = numpy.zeros(shape=(n,3), dtype=numpy.float32) + pol[:,0] = self.tracking_action.GetPolX() + pol[:,1] = self.tracking_action.GetPolY() + pol[:,2] = self.tracking_action.GetPolZ() + + wavelength = self.tracking_action.GetWavelength().astype(numpy.float32) + t0 = self.tracking_action.GetT0().astype(numpy.float32) + return { 'pos' : pos, + 'dir' : dir, + 'pol' : pol, + 'wavelength' : wavelength, + 't0' : t0 } + +if __name__ == '__main__': + import time + import optics + gen = G4Generator(optics.water) + # prime things + gen.generate_photons('e-', 1, (0,0,0), (1,0,0)) + + start = time.time() + n = 0 + for i in xrange(100): + photons = gen.generate_photons('mu-', 700, (0,0,0), (1,0,0)) + n += len(photons['t0']) + print photons['pos'][0].min(), photons['pos'][0].max() + stop = time.time() + print stop - start, 'sec' + print n / (stop-start), 'photons/sec' diff --git a/geometry.py b/geometry.py index 5fdfa1f..ac71429 100644 --- a/geometry.py +++ b/geometry.py @@ -109,6 +109,8 @@ class Material(object): self.refractive_index = None self.absorption_length = None self.scattering_length = None + self.density = 0.0 # g/cm^3 + self.composition = {} # by mass def set(self, name, value, wavelengths=standard_wavelengths): if np.iterable(value): @@ -49,6 +49,8 @@ r7081hqe_photocathode.set('reflect_diffuse', 1.0 - r7081hqe_photocathode.detect[ # water data comes from 'lightwater_sno' material in the SNO+ optics database water = Material('water') +water.density = 1.0 # g/cm^3 +water.composition = { 'H' : 0.1119, 'O' : 0.8881 } # fraction by mass water.absorption_length = \ np.array([[ 200. , 57.51539993], [ 220. , 64.22219849], |