diff options
| -rw-r--r-- | io/root.C | 58 | ||||
| -rw-r--r-- | io/root.py | 3 | ||||
| -rwxr-xr-x | sim.py | 133 |
3 files changed, 126 insertions, 68 deletions
@@ -9,23 +9,19 @@ struct Photon { TVector3 dir; TVector3 pol; double wavelength; // nm - int state; - int last_triangle; - int last_solid; + int history; + int last_hit_triangle; }; struct MC { std::string particle; TVector3 gen_pos; TVector3 gen_dir; - double gen_ke; + double gen_total_e; - std::vector<Photon> photon; + std::vector<Photon> photon_start; + std::vector<Photon> photon_stop; - void clear() { - particle = "none"; - photon.clear(); - } }; struct Channel { @@ -41,23 +37,42 @@ struct Event { int nhit; std::vector<Channel> channel; - void clear() { - mc.clear(); - channel.clear(); - } }; -void fill_event(TTree *T, - Event *ev, int event_id, double *gen_pos, int nchannels, - float *channel_times) +void fill_photons(Event *ev, bool start, + unsigned int nphotons, float *pos, float *dir, + float *pol, float *wavelength, float *t0, + int *histories=0, int *last_hit_triangle=0) +{ + std::vector<Photon> &photons = start ? ev->mc.photon_start : ev->mc.photon_stop; + photons.resize(nphotons); + + for (unsigned int i=0; i < nphotons; i++) { + Photon &photon = photons[i]; + photon.t = t0[i]; + photon.pos.SetXYZ(pos[3*i], pos[3*i + 1], pos[3*i + 2]); + photon.dir.SetXYZ(dir[3*i], dir[3*i + 1], dir[3*i + 2]); + photon.pol.SetXYZ(pol[3*i], pol[3*i + 1], pol[3*i + 2]); + photon.wavelength = wavelength[i]; + if (histories) + photon.history = histories[i]; + else + photon.history = 0; + + if (last_hit_triangle) + photon.last_hit_triangle = last_hit_triangle[i]; + else + photon.last_hit_triangle = -1; + } +} + + +void fill_hits(Event *ev, unsigned int nchannels, float *channel_times) { - ev->clear(); - ev->event_id = event_id; - MC &mc = ev->mc; - mc.gen_pos.SetXYZ(gen_pos[0], gen_pos[1], gen_pos[2]); + ev->channel.resize(0); ev->nhit = 0; Channel ch; - for (int i=0; i < nchannels; i++) { + for (unsigned int i=0; i < nchannels; i++) { if (channel_times[i] < 1e8) { ev->nhit++; ch.channel_id = i; @@ -66,7 +81,6 @@ void fill_event(TTree *T, ev->channel.push_back(ch); } } - T->Fill(); } @@ -5,7 +5,8 @@ ROOT.gROOT.ProcessLine('.L '+os.path.join(os.path.dirname(__file__), 'root.C+g') from ROOT import Event -fill_event = ROOT.fill_event +fill_photons = ROOT.fill_photons +fill_hits = ROOT.fill_hits def make_tree(name, desc=''): '''Create a ROOT tree for holding event information. @@ -2,6 +2,7 @@ import sys import optparse import time +import multiprocessing import detectors import optics @@ -27,8 +28,54 @@ def info(type, value, tb): sys.excepthook = info -if __name__ == '__main__': - +class GeneratorProcess(multiprocessing.Process): + def __init__(self, particle, energy, position, direction, nevents, material): + multiprocessing.Process.__init__(self) + + self.particle = particle + self.energy = energy + self.position = position + self.direction = direction + self.nevents = nevents + self.material = material + self.queue = multiprocessing.Queue() + + def run(self): + print >>sys.stderr, 'Starting generator thread...' + generator = g4gen.G4Generator(self.material) + + for i in xrange(self.nevents): + photons = generator.generate_photons(particle_name=self.particle, + total_energy=self.energy, + position=self.position, + direction=self.direction) + self.queue.put(photons) + + +def write_event(T, ev, event_id, hits, photon_start=None, photon_stop=None): + ev.event_id = event_id + if photon_start is not None: + photons = photon_start + root.fill_photons(ev, True, + len(photons['pos']), + np.ravel(photons['pos']), + np.ravel(photons['dir']), + np.ravel(photons['pol']), + photons['wavelength'], photons['t0']) + if photon_stop is not None: + photons = photon_stop + root.fill_photons(ev, False, + len(photons['pos']), + np.ravel(photons['pos']), + np.ravel(photons['dir']), + np.ravel(photons['pol']), + photons['wavelength'], photons['t0'], + photons['histories'], photons['last_hit_triangles']) + + root.fill_hits(ev, len(hits), hits) + T.Fill() + +def main(): parser = optparse.OptionParser('%prog') parser.add_option('-b', type='int', dest='nbits', default=10) parser.add_option('-j', type='int', dest='device', default=None) @@ -39,7 +86,13 @@ if __name__ == '__main__': parser.add_option('--particle', type='string', dest='particle', default='e-') parser.add_option('--energy', type='float', dest='energy', default=100.0) parser.add_option('--pos', type='string', dest='pos', default='(0,0,0)') - parser.add_option('--dir', type='string', dest='dir', default='(0,1,0)') + parser.add_option('--dir', type='string', dest='dir', default='(1,0,0)') + parser.add_option('--save-photon-start', action='store_true', + dest='save_photon_start', default=False, + help='Save initial photon vertices to disk') + parser.add_option('--save-photon-stop', action='store_true', + dest='save_photon_stop', default=False, + help='Save final photon vertices to disk') options, args = parser.parse_args() if len(args) != 1: @@ -59,11 +112,20 @@ if __name__ == '__main__': print >>sys.stderr, 'Creating detector...' detector.build(bits=options.nbits) - print >>sys.stderr, 'Initializing generator...' + print >>sys.stderr, 'Creating generator...' detector_material = optics.water - generator = g4gen.G4Generator(detector_material) + generator_thread = GeneratorProcess(particle=options.particle, + energy=options.energy, + position=position, + direction=direction, + nevents=options.nevents, + material=detector_material) print >>sys.stderr, 'WARNING: ASSUMING DETECTOR IS WATER!!' + # Do this now so we can get ahead of the photon propagation + print >>sys.stderr, 'Starting GEANT4 generator...' + generator_thread.start() + print >>sys.stderr, 'Initializing GPU...' gpu_worker = gpu.GPU(options.device) @@ -74,61 +136,39 @@ if __name__ == '__main__': gpu_worker.setup_propagate() gpu_worker.setup_daq(max(detector.pmtids)) - print >>sys.stderr, 'Loading tables in GEANT4...' - # Do easy event to force tables to load, throw away photons - generator.generate_photons(particle_name='e-', - total_energy=1.5, - position=(0,0,0), - direction=(1,0,0)) - - # Create output file f = ROOT.TFile(output_filename, 'RECREATE') ev, T = root.make_tree('T') print >>sys.stderr, 'Starting simulation...' start_sim = time.time() - #### Do the generation and writing in this offset order to ensure - #### that propagation on the GPU overlaps with CPU work to create - #### and save photons. - #### WARNING: THIS MAKES THE CODE LOOK A LITTLE CRAZY. I AM SORRY - - # Do first event - photons = generator.generate_photons(particle_name=options.particle, - total_energy=options.energy, - position=position, - direction=direction) - nphotons = len(photons['pos']) - gpu_worker.load_photons(**photons) - gpu_worker.propagate() - gpu_worker.run_daq() - - for i in xrange(1, options.nevents-1): - # photons for next event while previous event propagates on GPU - photons = generator.generate_photons(particle_name=options.particle, - total_energy=options.energy, - position=position, - direction=direction) + nphotons = 0 + for i in xrange(options.nevents): + photons = generator_thread.queue.get() + + assert len(photons['pos']) > 0, 'GEANT4 generated event with no photons!' + nphotons += len(photons['pos']) - - # this will stop and wait for event to finish - hit_times = gpu_worker.get_hits() - # turn around next event gpu_worker.load_photons(**photons) gpu_worker.propagate() gpu_worker.run_daq() + hit_times = gpu_worker.get_hits() - # write out this event - root.fill_event(T, ev, i-1, position, len(hit_times), hit_times) - + if options.save_photon_start: + photon_start = photons + else: + photon_start = None + if options.save_photon_stop: + photon_stop = gpu_worker.get_photons() + else: + photon_stop = None + + write_event(T, ev, i, hit_times, + photon_start=photon_start, photon_stop=photon_stop) if i % 10 == 0: print >>sys.stderr, "\rEvent:", i, - # Get results for last event and write it out - hit_times = gpu_worker.get_hits() - root.fill_event(T, ev, options.nevents - 1, position, len(hit_times), hit_times) - end_sim = time.time() print >>sys.stderr, "\rEvent:", options.nevents - 1 @@ -136,3 +176,6 @@ if __name__ == '__main__': f.Close() print >>sys.stderr, 'Done. %1.1f events/sec, %1.0f photons/sec.' % (options.nevents/(end_sim - start_sim), nphotons/(end_sim - start_sim)) + +if __name__ == '__main__': + main() |