merge

6 files changed, 576 insertions, 0 deletions

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 <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/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 <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/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