diff options
Diffstat (limited to 'sim.py')
| -rwxr-xr-x | sim.py | 149 |
1 files changed, 0 insertions, 149 deletions
@@ -1,149 +0,0 @@ -#!/usr/bin/env python -import time -import os -import numpy as np - -from chroma import generator -from chroma import gpu -from chroma import event -from chroma import itertoolset - -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, 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 - - # 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 - - self.context = gpu.create_cuda_context(cuda_device) - - if not hasattr(detector, 'mesh'): - # need to build geometry - detector.build() - - self.gpu_geometry = gpu.GPUGeometry(detector) - self.gpu_daq = gpu.GPUDaq(self.gpu_geometry, max(self.detector.pmtids)) - self.gpu_pdf = gpu.GPUPDF() - - 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 = itertoolset.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(primary_vertex=vertex, 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 = itertoolset.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 = itertoolset.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, ndaq=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 = itertoolset.peek(iterable) - - if isinstance(first_element, event.Event): - iterable = self.photon_generator.generate_events(iterable) - - for ev in iterable: - gpu_photons = gpu.GPUPhotons(ev.photons_beg, ncopies=nreps) - gpu_photons.propagate(self.gpu_geometry, self.rng_states, - nthreads_per_block=self.nthreads_per_block, - max_blocks=self.max_blocks) - for gpu_photon_slice in gpu_photons.iterate_copies(): - for idaq in xrange(ndaq): - gpu_channels = self.gpu_daq.acquire(gpu_photon_slice, 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() - - def __del__(self): - self.context.pop() |