#!/usr/bin/env python import sys import time import os import numpy as np import detectors import optics import generator from generator import constant import itertools import threading import gpu from fileio import root from chroma.itertoolset import repeating_iterator from tools import profile_if_possible, enable_debug_on_crash 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) class Simulation(object): def __init__(self, detector, detector_material, seed=None, cuda_device=None, geant4_processes=4, bvh_bits=11, use_cache=True): self.detector = detector if seed is None: self.seed = pick_seed() else: self.seed = seed print >>sys.stderr, 'RNG seed:', self.seed # We have three generators to seed: numpy.random, GEANT4, and CURAND. # The latter two are done below np.random.seed(self.seed) self.detector_material = detector_material if geant4_processes > 0: self.photon_generator = generator.photon.G4ParallelGenerator(geant4_processes, detector_material, base_seed=self.seed) else: self.photon_generator = None print >>sys.stderr, 'Creating BVH with %d bits...' % (bvh_bits) detector.build(bits=bvh_bits, use_cache=use_cache) print >>sys.stderr, 'Initializing GPU...' self.gpu_worker = gpu.GPU(cuda_device) print >>sys.stderr, 'Loading detector onto GPU...' self.gpu_worker.load_geometry(detector) print >>sys.stderr, 'Initializing random numbers generators...' self.gpu_worker.setup_propagate(seed=self.seed) self.gpu_worker.setup_daq(max(self.detector.pmtids)) self.pdf_config = None def simulate(self, nevents, vertex_generator, keep_photon_start=False, keep_photon_stop=False, run_daq=True, nreps=1): photon_gen = repeating_iterator(self.photon_generator.generate_events(nevents, vertex_generator), nreps) return self.simulate_photons(nevents*nreps, photon_gen, keep_photon_start=keep_photon_start, keep_photon_stop=keep_photon_stop, run_daq=run_daq) def simulate_photons(self, nevents, photon_generator, keep_photon_start=False, keep_photon_stop=False, run_daq=True): for ev in itertools.islice(photon_generator, nevents): self.gpu_worker.load_photons(ev.photon_start) self.gpu_worker.propagate() self.gpu_worker.run_daq() ev.nphoton = len(ev.photon_start.positions) if not keep_photon_start: ev.photon_start = None if keep_photon_stop: ev.photon_stop = self.gpu_worker.get_photons() if run_daq: ev.hits = self.gpu_worker.get_hits() ev.channels = ev.hits yield ev def propagate_photons(self, photons, max_steps=10): self.gpu_worker.load_photons(photons) self.gpu_worker.propagate(max_steps=max_steps) return self.gpu_worker.get_photons() def create_pdf(self, nevents, vertex_generator, tbins, trange, qbins, qrange, nreps=1): photon_gen = repeating_iterator(self.photon_generator.generate_events(nevents, vertex_generator), nreps) return self.create_pdf_from_photons(nevents*nreps, photon_gen, tbins, trange, qbins, qrange) def create_pdf_from_photons(self, nevents, photon_generator, tbins, trange, qbins, qrange): '''Returns tuple: 1D array of channel hit counts, 3D array of (channel, time, charge) pdfs''' pdf_config = (tbins, trange, qbins, qrange) if pdf_config != self.pdf_config: self.pdf_config = pdf_config self.gpu_worker.setup_pdf(max(self.detector.pmtids), tbins, trange, qbins, qrange) else: self.gpu_worker.clear_pdf() for ev in itertools.islice(photon_generator, nevents): self.gpu_worker.load_photons(ev.photon_start) self.gpu_worker.propagate() self.gpu_worker.run_daq() self.gpu_worker.add_hits_to_pdf() return self.gpu_worker.get_pdfs() @profile_if_possible def main(): import optparse parser = optparse.OptionParser('%prog filename') parser.add_option('-b', type='int', dest='nbits', default=11) parser.add_option('-j', type='int', dest='device', default=None) 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: sys.exit(parser.format_help()) else: output_filename = args[0] if options.nevents <= 0: sys.exit('--nevents must be greater than 0!') position = np.array(eval(options.pos), dtype=float) direction = np.array(eval(options.dir), dtype=float) print >>sys.stderr, 'Loading detector %s...' % options.detector detector = detectors.find(options.detector) print >>sys.stderr, 'Creating generator...' if options.particle == 'pi0': vertex_generator = generator.vertex.pi0_gun(pi0_position=constant(position), pi0_direction=constant(direction), pi0_total_energy=constant(options.energy)) else: vertex_generator = generator.vertex.particle_gun(particle_name=constant(options.particle), position=constant(position), direction=constant(direction), total_energy=constant(options.energy)) # Initializing simulation print >>sys.stderr, 'WARNING: ASSUMING DETECTOR IS WCSIM WATER!!' simulation = Simulation(detector=detector, detector_material=optics.water_wcsim, seed=options.seed, cuda_device=options.device, geant4_processes=options.ngenerators, bvh_bits=options.nbits) # Create output file writer = root.RootWriter(output_filename) # Preheat generator event_iterator = simulation.simulate(options.nevents, vertex_generator, keep_photon_start=options.save_photon_start, keep_photon_stop=options.save_photon_stop) print >>sys.stderr, 'Starting simulation...' start_sim = time.time() nphotons = 0 for i, ev in enumerate(event_iterator): assert ev.nphoton > 0, 'GEANT4 generated event with no photons!' nphotons += ev.nphoton writer.write_event(ev) 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__': enable_debug_on_crash() main()