#!/usr/bin/env python import sys import optparse import time import os import multiprocessing import detectors import optics import gpu import g4gen from fileio import root import numpy as np import math import ROOT def pick_seed(): '''Returns a seed for a random number generator selected using a mixture of the current time and the current process ID.''' return int(time.time()) ^ (os.getpid() << 16) 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() class GeneratorProcess(multiprocessing.Process): def __init__(self, particle, energy, position, direction, nevents, material, queue, seed=None): multiprocessing.Process.__init__(self) self.particle = particle self.energy = energy self.position = position self.direction = direction self.nevents = nevents self.material = material self.seed = seed self.queue = queue self.daemon = True def run(self): print >>sys.stderr, 'Starting generator thread...' generator = g4gen.G4Generator(self.material, seed=self.seed) for i in xrange(self.nevents): if self.particle == 'pi0': photons = generator.generate_pi0(total_energy=self.energy, position=self.position, direction=self.direction) else: photons = generator.generate_photons(particle_name=self.particle, total_energy=self.energy, position=self.position, direction=self.direction) self.queue.put(photons) 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) # Allow profile decorator to exist, but do nothing if not running under kernprof try: profile = profile except NameError: profile = lambda x: x @profile 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) parser.add_option('-n', type='int', dest='nblocks', default=64) parser.add_option('-s', type='int', dest='seed', default=None, help='Set random number generator seed') parser.add_option('-g', type='int', dest='ngenerators', default=4, help='Number of GEANT4 generator processes') 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='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: 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.dir), dtype=float) detector = detectors.find(options.detector) if options.seed is None: options.seed = pick_seed() print >>sys.stderr, 'RNG seed:', options.seed print >>sys.stderr, 'Creating generator...' detector_material = optics.water_wcsim queue = multiprocessing.Queue() generators = [GeneratorProcess(particle=options.particle, energy=options.energy, position=position, direction=direction, nevents=nevents, material=detector_material, seed=options.seed + seed_offset, queue=queue) for seed_offset, nevents in enumerate(partition(options.nevents, options.ngenerators))] print >>sys.stderr, 'WARNING: ASSUMING DETECTOR IS WCSIM WATER!!' # Do this now so we can get ahead of the photon propagation print >>sys.stderr, 'Starting GEANT4 generators...' for generator in generators: generator.start() print >>sys.stderr, 'Creating BVH for detector "%s" with %d bits...' % (options.detector, options.nbits) detector.build(bits=options.nbits) 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(seed=options.seed) gpu_worker.setup_daq(max(detector.pmtids)) # Create output file writer = root.RootWriter(output_filename) # Set generator info writer.set_generated_particle(name=options.particle, position=position, direction=direction, total_e=options.energy) print >>sys.stderr, 'Starting simulation...' start_sim = time.time() nphotons = 0 for i in xrange(options.nevents): photons = queue.get() assert len(photons['pos']) > 0, 'GEANT4 generated event with no photons!' nphotons += len(photons['pos']) gpu_worker.load_photons(pos=photons['pos'], dir=photons['dir'], pol=photons['pol'], t0=photons['t0'], wavelength=photons['wavelength']) gpu_worker.propagate() gpu_worker.run_daq() hits = gpu_worker.get_hits() 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 if 'subtracks' in photons: subtracks = photons['subtracks'] else: subtracks = None writer.write_event(i, hits, photon_start=photon_start, photon_stop=photon_stop, subtracks=subtracks) if i % 10 == 0: print >>sys.stderr, "\rEvent:", i, end_sim = time.time() print >>sys.stderr, "\rEvent:", options.nevents - 1 writer.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__': sys.excepthook = info main()