#!/usr/bin/env python import sys import time import os import numpy as np import itertools import threading from chroma import detectors from chroma import optics from chroma import generator from chroma import gpu from chroma.fileio import root from chroma import tools from chroma import event from chroma.itertoolset import peek, repeat_copy, repeating_iterator 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, seed=None, cuda_device=None, geant4_processes=4, bvh_bits=11, use_cache=True, nthreads_per_block=64, max_blocks=1024): self.detector = detector self.nthreads_per_block = nthreads_per_block self.max_blocks = max_blocks if seed is None: self.seed = pick_seed() else: self.seed = seed print 'RNG seed: %i' % self.seed # We have three generators to seed: numpy.random, GEANT4, and CURAND. # The latter two are done below. np.random.seed(self.seed) if geant4_processes > 0: self.photon_generator = generator.photon.G4ParallelGenerator(geant4_processes, detector.detector_material, base_seed=self.seed) else: self.photon_generator = None if not hasattr(detector, 'mesh'): # need to build geometry print 'Creating BVH with %i bits...' % bvh_bits detector.build(bits=bvh_bits, use_cache=use_cache) self.gpu = gpu.GPU(cuda_device) self.gpu_geometry = gpu.GPUGeometry(detector) self.gpu_daq = gpu.GPUDaq(self.gpu_geometry, max(self.detector.pmtids)) self.gpu_pdf = gpu.GPUPDF() print 'Initializing random numbers generators...' self.rng_states = gpu.get_rng_states(self.nthreads_per_block*self.max_blocks, seed=self.seed) self.pdf_config = None def simulate(self, iterable, keep_photons_beg=False, keep_photons_end=False, run_daq=True, max_steps=10): try: first_element, iterable = peek(iterable) except TypeError: first_element, iterable = iterable, [iterable] if isinstance(first_element, event.Event): iterable = self.photon_generator.generate_events(iterable) elif isinstance(first_element, event.Photons): iterable = (event.Event(photons_beg=x) for x in iterable) elif isinstance(first_element, event.Vertex): iterable = (event.Event(vertices=[vertex]) for vertex in iterable) iterable = self.photon_generator.generate_events(iterable) for ev in iterable: gpu_photons = gpu.GPUPhotons(ev.photons_beg) gpu_photons.propagate(self.gpu_geometry, self.rng_states, nthreads_per_block=self.nthreads_per_block, max_blocks=self.max_blocks, max_steps=max_steps) ev.nphotons = len(ev.photons_beg.pos) if not keep_photons_beg: ev.photons_beg = None if keep_photons_end: ev.photons_end = gpu_photons.get() if run_daq: gpu_channels = self.gpu_daq.acquire(gpu_photons, self.rng_states, nthreads_per_block=self.nthreads_per_block, max_blocks=self.max_blocks) ev.channels = gpu_channels.get() yield ev def create_pdf(self, iterable, tbins, trange, qbins, qrange, nreps=1): """Returns tuple: 1D array of channel hit counts, 3D array of (channel, time, charge) pdfs.""" first_element, iterable = peek(iterable) if isinstance(first_element, event.Event): iterable = self.photon_generator.generate_events(iterable) pdf_config = (tbins, trange, qbins, qrange) if pdf_config != self.pdf_config: self.pdf_config = pdf_config self.gpu_pdf.setup_pdf(max(self.detector.pmtids), tbins, trange, qbins, qrange) else: self.gpu_pdf.clear_pdf() if nreps > 1: iterable = repeating_iterator(iterable, nreps) for ev in iterable: gpu_photons = gpu.GPUPhotons(ev.photons_beg) gpu_photons.propagate(self.gpu_geometry, self.rng_states, nthreads_per_block=self.nthreads_per_block, max_blocks=self.max_blocks) gpu_channels = self.gpu_daq.acquire(gpu_photons, self.rng_states, nthreads_per_block=self.nthreads_per_block, max_blocks=self.max_blocks) self.gpu_pdf.add_hits_to_pdf(gpu_channels) return self.gpu_pdf.get_pdfs() def eval_pdf(self, event_channels, iterable, min_twidth, trange, min_qwidth, qrange, min_bin_content=20, nreps=1, time_only=True): """Returns tuple: 1D array of channel hit counts, 1D array of PDF probability densities.""" self.gpu_pdf.setup_pdf_eval(event_channels.hit, event_channels.t, event_channels.q, min_twidth, trange, min_qwidth, qrange, min_bin_content=min_bin_content, time_only=True) first_element, iterable = peek(iterable) if isinstance(first_element, event.Event): iterable = self.photon_generator.generate_events(iterable) if nreps > 1: iterable = repeating_iterator(iterable, nreps) for ev in iterable: gpu_photons = gpu.GPUPhotons(ev.photons_beg) gpu_photons.propagate(self.gpu_geometry, self.rng_states, nthreads_per_block=self.nthreads_per_block, max_blocks=self.max_blocks) gpu_channels = self.gpu_daq.acquire(gpu_photons, self.rng_states, nthreads_per_block=self.nthreads_per_block, max_blocks=self.max_blocks) self.gpu_pdf.accumulate_pdf_eval(gpu_channels) return self.gpu_pdf.get_pdf_eval() @tools.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', help='particle name', default='e-') parser.add_option('--ke', type='float', dest='ke', help='particle kinetic energy in MeV', default=100.0) parser.add_option('--pos', type='string', dest='pos', help='particle vertex origin.', default='(0,0,0)') parser.add_option('--dir', type='string', dest='dir', help='particle vertex direction.', default='(1,0,0)') parser.add_option('--save-photon-beg', action='store_true', dest='save_photons_beg', default=False, help='Save initial photon vertices to disk') parser.add_option('--save-photon-end', action='store_true', dest='save_photons_end', 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!') pos = np.array(eval(options.pos), dtype=float) dir = np.array(eval(options.dir), dtype=float) print 'Loading detector %s...' % options.detector sys.stdout.flush() detector = detectors.find(options.detector) print 'Creating particle vertex generator...' sys.stdout.flush() if options.particle == 'pi0': ev_vertex_iter = itertools.islice(generator.vertex.pi0_gun(itertools.repeat(pos), itertools.repeat(dir), itertools.repeat(options.ke)), options.nevents) else: vertex = event.Vertex(options.particle, pos, dir, None, options.ke) ev_vertex_iter = (event.Event(i, vertex, [vertex]) for i, vertex in zip(range(options.nevents), repeat_copy(vertex))) # Initializing simulation simulation = Simulation(detector=detector, 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 ev_iter = simulation.simulate(ev_vertex_iter, keep_photons_beg=options.save_photons_beg, keep_photons_end=options.save_photons_end) print 'Starting simulation...' nphotons = 0 t0 = time.time() for i, ev in enumerate(ev_iter): print "\rEvent: %i" % (i+1), sys.stdout.flush() assert ev.nphotons > 0, 'Geant4 generated event with no photons!' nphotons += ev.nphotons writer.write_event(ev) print elapsed = time.time() - t0 writer.close() print '%f elapsed, %1.1f events/sec, %1.0f photons/sec.' % \ (elapsed, options.nevents/elapsed, nphotons/elapsed) if __name__ == '__main__': tools.enable_debug_on_crash() main()