diff options
| -rwxr-xr-x | sim.py | 132 |
1 files changed, 132 insertions, 0 deletions
@@ -0,0 +1,132 @@ +#!/usr/bin/env python +import sys +import optparse +import time + +import detectors +import optics +import gpu +import g4gen +from io import root +import numpy as np +import math +import ROOT + +def info(type, value, tb): + if hasattr(sys, 'ps1') or not sys.stderr.isatty(): + # we are in interactive mode or we don't have a tty-like + # device, so we call the default hook + sys.__excepthook__(type, value, tb) + else: + import traceback, pdb + # we are NOT in interactive mode, print the exception... + traceback.print_exception(type, value, tb) + print + # ...then start the debugger in post-mortem mode. + pdb.pm() + +sys.excepthook = info + +if __name__ == '__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) + parser.add_option('-n', type='int', dest='nblocks', default=64) + + parser.add_option('--detector', type='string', dest='detector', default='microlbne') + parser.add_option('--nevents', type='int', dest='nevents', default=100) + 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='pos', default='(1,0,0)') + + options, args = parser.parse_args() + if len(args) != 1: + print 'Must specify output filename!' + sys.exit(1) + else: + output_filename = args[0] + + if options.nevents <= 0: + print '--nevents must be greater than 0!' + sys.exit(1) + + position = np.array(eval(options.pos), dtype=float) + direction = np.array(eval(options.pos), dtype=float) + detector = detectors.find(options.detector) + + print >>sys.stderr, 'Creating detector...' + detector.build(bits=options.nbits) + + print >>sys.stderr, 'Initializing generator...' + detector_material = optics.water + generator = g4gen.G4Generator(detector_material) + print >>sys.stderr, 'WARNING: ASSUMING DETECTOR IS WATER!!' + + print >>sys.stderr, 'Initializing GPU...' + gpu_worker = gpu.GPU(options.device) + + print >>sys.stderr, 'Loading detector onto GPU...' + gpu_worker.load_geometry(detector) + + print >>sys.stderr, 'Initializing random numbers generators...' + gpu_worker.setup_propagate() + gpu_worker.setup_daq(max(detector.pmtids)) + + # 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 += 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() + + # write out this event + root.fill_event(T, ev, i-1, position, len(hit_times), hit_times) + + 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 + + T.Write() + 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)) + + gpu_worker.shutdown() |