path: root/sim.py

#!/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)

        # geometry is loaded onto gpu by default
        self.gpu_geometry = gpu.GPUGeometry(self.gpu, 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.propagate(self.gpu, gpu_photons, 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)

        for ev in iterable:
            gpu_photons = gpu.GPUPhotons(ev.photons_beg)
            gpu.propagate(self.gpu, gpu_photons, 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)

        for ev in iterable:
            gpu_photons = gpu.GPUPhotons(ev.photons_beg)
            gpu.propagate(self.gpu, gpu_photons, 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()