diff options
| -rw-r--r-- | __init__.py | 8 | ||||
| -rwxr-xr-x | chroma-sim | 87 | ||||
| -rw-r--r-- | event.py | 2 | ||||
| -rw-r--r-- | fileio/root.py | 24 | ||||
| -rw-r--r-- | generator/g4gen.py | 4 | ||||
| -rw-r--r-- | generator/vertex.py | 12 | ||||
| -rw-r--r-- | gpu.py | 25 | ||||
| -rw-r--r-- | itertoolset.py | 16 | ||||
| -rwxr-xr-x | sim.py | 132 | ||||
| -rw-r--r-- | solids/__init__.py | 2 | ||||
| -rw-r--r-- | tests/test_fileio.py | 3 | ||||
| -rw-r--r-- | tests/test_pdf.py | 8 | ||||
| -rw-r--r-- | tests/test_propagation.py | 3 | ||||
| -rw-r--r-- | tests/test_rayleigh.py | 5 |
14 files changed, 157 insertions, 174 deletions
diff --git a/__init__.py b/__init__.py index 07a096c..3359bbc 100644 --- a/__init__.py +++ b/__init__.py @@ -1,15 +1,15 @@ from camera import Camera, EventViewer, view, build import geometry import event -import fileio -#import generator +from fileio import root +import generator import gpu import itertoolset -#import likelihood +import likelihood import make import optics import project import sample -#import sim +from sim import Simulation from stl import mesh_from_stl import transform diff --git a/chroma-sim b/chroma-sim new file mode 100755 index 0000000..fa106a2 --- /dev/null +++ b/chroma-sim @@ -0,0 +1,87 @@ +#!/usr/bin/env python +#--*-python-*- + +if __name__ == '__main__': + import optparse + import sys + import imp + import os + import inspect + import numpy as np + + from chroma import event + from chroma import itertoolset + from chroma import Simulation + from chroma import root + + parser = optparse.OptionParser('%prog <detector>') + parser.add_option('-o', dest='output_filename', + help='output filename', default='out.root') + parser.add_option('-j', type='int', dest='device', + help='CUDA device number', default=None) + parser.add_option('-s', type='int', dest='seed', + help='random number generator seed') + parser.add_option('-g', type='int', dest='ngenerators', + help='number of GEANT4 processes', default=4) + parser.add_option('-n', '--nevents', type='int', dest='nevents', + default=100) + parser.add_option('-p', '--particle', dest='particle', + help='particle name', default='e-') + parser.add_option('-k', '--ke', type='float', dest='ke', + help='kinetic energy (MeV)', default=100.0) + parser.add_option('--pos', dest='pos', + help='particle vertex origin', default='0,0,0') + parser.add_option('--dir', dest='dir', + help='particle vertex direction', default='1,0,0') + parser.add_option('--save-photons-beg', action='store_true', + dest='save_photons_beg', + help='save initial photons to disk', default=False) + parser.add_option('--save-photons-end', action='store_true', + dest='save_photons_end', + help='save final photons to disk', default=False) + + options, args = parser.parse_args() + + if len(args) < 1: + sys.exit(parser.format_help()) + + name, attr = args[0].split('.') + + try: + file, path, description = imp.find_module(name) + except ImportError: + raise + + module = imp.load_module(name, file, path, description) + + detector = getattr(module, attr) + + if inspect.isfunction(detector): + detector = detector() + + pos = np.array([float(s) for s in options.pos.split(',')], dtype=float) + dir = np.array([float(s) for s in options.dir.split(',')], dtype=float) + + vertex = event.Vertex(options.particle, pos, dir, options.ke) + + sim = Simulation(detector, seed=options.seed, cuda_device=options.device, + geant4_processes=options.ngenerators) + + print 'RNG seed: %i' % sim.seed + + writer = root.RootWriter(options.output_filename) + + for i, ev in \ + enumerate(sim.simulate(itertoolset.repeatcopy(vertex, + options.nevents), + keep_photons_beg=options.save_photons_beg, + keep_photons_end=options.save_photons_end)): + print '\rEvent: %i' % (i+1), + sys.stdout.flush() + + assert ev.nphotons > 0, 'Geant4 generated event with no photons!' + + writer.write_event(ev) + print + + writer.close() @@ -1,7 +1,7 @@ import numpy as np class Vertex(object): - def __init__(self, particle_name, pos, dir, pol, ke, t0=0.0): + def __init__(self, particle_name, pos, dir, ke, t0=0.0, pol=None): self.particle_name = particle_name self.pos = pos self.dir = dir diff --git a/fileio/root.py b/fileio/root.py index 507c249..af86deb 100644 --- a/fileio/root.py +++ b/fileio/root.py @@ -25,11 +25,11 @@ def make_photon_with_arrays(size): def root_vertex_to_python_vertex(vertex): "Returns a chroma.event.Vertex object from a root Vertex object." return event.Vertex(str(vertex.particle_name), - tvector3_to_ndarray(vertex.pos), - tvector3_to_ndarray(vertex.dir), - tvector3_to_ndarray(vertex.pol), - vertex.ke, - vertex.t0) + pos=tvector3_to_ndarray(vertex.pos), + dir=tvector3_to_ndarray(vertex.dir), + ke=vertex.ke, + t0=vertex.t0, + pol=tvector3_to_ndarray(vertex.pol)) def root_event_to_python_event(ev): '''Returns a new chroma.event.Event object created from the @@ -148,12 +148,14 @@ class RootWriter(object): "Write an event.Event object to the ROOT tree as a ROOT.Event object." self.ev.id = pyev.id - self.ev.primary_vertex.particle_name = pyev.primary_vertex.particle_name - self.ev.primary_vertex.pos.SetXYZ(*pyev.primary_vertex.pos) - self.ev.primary_vertex.dir.SetXYZ(*pyev.primary_vertex.dir) - if pyev.primary_vertex.pol is not None: - self.ev.primary_vertex.pol.SetXYZ(*pyev.primary_vertex.pol) - self.ev.primary_vertex.ke = pyev.primary_vertex.ke + if pyev.primary_vertex is not None: + self.ev.primary_vertex.particle_name = \ + pyev.primary_vertex.particle_name + self.ev.primary_vertex.pos.SetXYZ(*pyev.primary_vertex.pos) + self.ev.primary_vertex.dir.SetXYZ(*pyev.primary_vertex.dir) + if pyev.primary_vertex.pol is not None: + self.ev.primary_vertex.pol.SetXYZ(*pyev.primary_vertex.pol) + self.ev.primary_vertex.ke = pyev.primary_vertex.ke if pyev.photons_beg is not None: photons = pyev.photons_beg diff --git a/generator/g4gen.py b/generator/g4gen.py index 8dca086..3836897 100644 --- a/generator/g4gen.py +++ b/generator/g4gen.py @@ -120,7 +120,9 @@ class G4Generator(object): self.particle_gun.SetParticleEnergy(total_energy) self.particle_gun.SetParticlePosition(G4ThreeVector(*vertex.pos)*m) self.particle_gun.SetParticleMomentumDirection(G4ThreeVector(*vertex.dir).unit()) - + if vertex.pol is not None: + self.particle_gun.SetParticlePolarization(G4ThreeVector(*vertex.pol).unit()) + self.tracking_action.Clear() gRunManager.BeamOn(1) diff --git a/generator/vertex.py b/generator/vertex.py index 5521d6b..b4e2d23 100644 --- a/generator/vertex.py +++ b/generator/vertex.py @@ -4,7 +4,7 @@ from itertools import izip, count from chroma.pi0 import pi0_decay from chroma import event from chroma.sample import uniform_sphere -from chroma.itertoolset import repeat_func +from chroma.itertoolset import repeatfunc # generator parts for use with gun() @@ -13,7 +13,7 @@ def from_histogram(h): pdf = h.hist/h.hist.sum() cdf = np.cumsum(pdf) - for x in repeat_func(np.random.random_sample): + for x in repeatfunc(np.random.random_sample): yield h.bincenters[np.searchsorted(cdf, x)] def constant(obj): @@ -43,14 +43,14 @@ def flat(e_lo, e_hi): def particle_gun(particle_name_iter, pos_iter, dir_iter, ke_iter, start_id=0): for i, particle_name, pos, dir, ke in izip(count(start_id), particle_name_iter, pos_iter, dir_iter, ke_iter): dir /= np.linalg.norm(dir) - vertex = event.Vertex(particle_name, pos, dir, None, ke) + vertex = event.Vertex(particle_name, pos, dir, ke) ev_vertex = event.Event(i, vertex, [vertex]) yield ev_vertex def pi0_gun(pos_iter, dir_iter, ke_iter, start_id=0): for i, pos, dir, ke in izip(count(start_id), pos_iter, dir_iter, ke_iter): dir /= np.linalg.norm(dir) - primary_vertex = event.Vertex('pi0', pos, dir, None, ke) + primary_vertex = event.Vertex('pi0', pos, dir, ke) cos_theta_rest = np.random.random_sample() * 2 - 1 theta_rest = np.arccos(cos_theta_rest) @@ -59,8 +59,8 @@ def pi0_gun(pos_iter, dir_iter, ke_iter, start_id=0): (gamma1_e, gamma1_dir), (gamma2_e, gamma2_dir) = \ pi0_decay(ke+134.9766, dir, theta_rest, phi_rest) - gamma1_vertex = event.Vertex('gamma', pos, gamma1_dir, None, gamma1_e) - gamma2_vertex = event.Vertex('gamma', pos, gamma2_dir, None, gamma2_e) + gamma1_vertex = event.Vertex('gamma', pos, gamma1_dir, gamma1_e) + gamma2_vertex = event.Vertex('gamma', pos, gamma2_dir, gamma2_e) ev_vertex = event.Event(i, primary_vertex, [gamma1_vertex, gamma2_vertex]) @@ -797,28 +797,3 @@ def create_cuda_context(device_id=None): context.set_cache_config(cuda.func_cache.PREFER_L1) return context - -class GPU(object): - def __init__(self, device_id=None): - """Initialize a GPU context on the specified device. - If device_id is None, the default device is used.""" - - if device_id is None: - self.context = pycuda.tools.make_default_context() - else: - device = cuda.Device(device_id) - self.context = device.make_context() - - print 'device %s' % self.context.get_device().name() - - self.print_mem_info() - - self.context.set_cache_config(cuda.func_cache.PREFER_L1) - - def print_mem_info(self): - free, total = cuda.mem_get_info() - - print '%.1f %% of device memory is free.' % ((free/float(total))*100) - - def __del__(self): - self.context.pop() diff --git a/itertoolset.py b/itertoolset.py index 498b940..4d293e8 100644 --- a/itertoolset.py +++ b/itertoolset.py @@ -15,16 +15,16 @@ def peek(iterable): first_element = next(it) return first_element, chain([first_element], it) -def repeat_func(func, times=None, args=()): - "equivalent to (func(*args) for i in xrange(times))." +def repeatfunc(func, times=None, *args): + """Repeat calls to func with specified arguments. + + Example: repeatfunc(random.random) + """ if times is None: - while True: - yield func(*args) - else: - for i in xrange(times): - yield func(*args) + return starmap(func, repeat(args)) + return starmap(func, repeat(args, times)) -def repeat_copy(object, times=None): +def repeatcopy(object, times=None): """Returns deep copies of `object` over and over again. Runs indefinitely unless the `times` argument is specified.""" if times is None: @@ -1,19 +1,12 @@ #!/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 +from chroma import itertoolset def pick_seed(): """Returns a seed for a random number generator selected using @@ -21,7 +14,8 @@ def pick_seed(): 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): + 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 @@ -32,7 +26,6 @@ class Simulation(object): 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) @@ -42,24 +35,24 @@ class Simulation(object): else: self.photon_generator = None + self.context = gpu.create_cuda_context(cuda_device) + 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) + 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() - 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): + def simulate(self, iterable, keep_photons_beg=False, + keep_photons_end=False, run_daq=True, max_steps=10): try: - first_element, iterable = peek(iterable) + first_element, iterable = itertoolset.peek(iterable) except TypeError: first_element, iterable = iterable, [iterable] @@ -68,13 +61,16 @@ class Simulation(object): 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 = (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) + 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) @@ -93,7 +89,7 @@ class Simulation(object): 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) + first_element, iterable = itertoolset.peek(iterable) if isinstance(first_element, event.Event): iterable = self.photon_generator.generate_events(iterable) @@ -102,12 +98,12 @@ class Simulation(object): if pdf_config != self.pdf_config: self.pdf_config = pdf_config self.gpu_pdf.setup_pdf(max(self.detector.pmtids), tbins, trange, - qbins, qrange) + qbins, qrange) else: self.gpu_pdf.clear_pdf() if nreps > 1: - iterable = repeating_iterator(iterable, nreps) + iterable = itertoolset.repeating_iterator(iterable, nreps) for ev in iterable: gpu_photons = gpu.GPUPhotons(ev.photons_beg) @@ -132,7 +128,7 @@ class Simulation(object): min_bin_content=min_bin_content, time_only=True) - first_element, iterable = peek(iterable) + first_element, iterable = itertoolset.peek(iterable) if isinstance(first_element, event.Event): iterable = self.photon_generator.generate_events(iterable) @@ -149,91 +145,5 @@ class Simulation(object): 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() + def __del__(self): + self.context.pop() diff --git a/solids/__init__.py b/solids/__init__.py index 5ec11cf..c4d227f 100644 --- a/solids/__init__.py +++ b/solids/__init__.py @@ -29,7 +29,7 @@ def build_12inch_pmt_with_lc_hd(outer_material=water, nsteps=128): rmax=lc_12inch_rmax, npoints=100) def build_8inch_pmt(outer_material=water, nsteps=24): - return build_pmt(dirname(__file__) + '/sno_pmt_reduced.txt', 0.003, + return build_pmt(dirname(__file__) + '/sno_pmt.txt', 0.003, outer_material, nsteps) def build_8inch_pmt_with_lc(outer_material=water, nsteps=24): diff --git a/tests/test_fileio.py b/tests/test_fileio.py index 2911976..3869a9f 100644 --- a/tests/test_fileio.py +++ b/tests/test_fileio.py @@ -5,7 +5,8 @@ import numpy as np class TestFileIO(unittest.TestCase): def test_file_write_and_read(self): - ev = event.Event(1, event.Vertex('e-', (0,0,1), (1,0,0), (0,1,0), 15)) + ev = event.Event(1, event.Vertex('e-', pos=(0,0,1), dir=(1,0,0), + ke=15.0, pol=(0,1,0))) photons_beg = root.make_photon_with_arrays(1) photons_beg.pos[0] = (1,2,3) diff --git a/tests/test_pdf.py b/tests/test_pdf.py index c7c9394..571cbd4 100644 --- a/tests/test_pdf.py +++ b/tests/test_pdf.py @@ -11,7 +11,7 @@ from chroma.sim import Simulation class TestPDF(unittest.TestCase): def setUp(self): - self.detector = chroma.detectors.find("microlbne") + self.detector = chroma.detectors.microlbne() self.detector.build() self.vertex_gen = constant_particle_gun('e-', (0,0,0), (1,0,0), 10) @@ -19,7 +19,9 @@ class TestPDF(unittest.TestCase): '''Create a hit count and (q,t) PDF for 10 MeV events in MicroLBNE''' g4generator = G4ParallelGenerator(1, water_wcsim) - gpu_instance = gpu.GPU(0) + + context = gpu.create_cuda_context() + gpu_geometry = gpu.GPUGeometry(self.detector) nthreads_per_block, max_blocks = 64, 1024 @@ -48,6 +50,8 @@ class TestPDF(unittest.TestCase): for i, nhits in enumerate(hitcount): self.assertEqual(nhits, pdf[i].sum()) + context.pop() + def testSimPDF(self): sim = Simulation(self.detector) diff --git a/tests/test_propagation.py b/tests/test_propagation.py index 8f9d084..43ecb9b 100644 --- a/tests/test_propagation.py +++ b/tests/test_propagation.py @@ -20,7 +20,8 @@ class TestPropagation(unittest.TestCase): cube = Geometry(vacuum) cube.add_solid(Solid(box(100,100,100), vacuum, vacuum)) cube.pmtids = [0] - sim = Simulation(cube, bvh_bits=4, geant4_processes=0, use_cache=False) + cube.build(use_cache=False) + sim = Simulation(cube, geant4_processes=0) # Create initial photons nphotons = 10000 diff --git a/tests/test_rayleigh.py b/tests/test_rayleigh.py index 7e5fc92..02ccb41 100644 --- a/tests/test_rayleigh.py +++ b/tests/test_rayleigh.py @@ -16,8 +16,9 @@ class TestRayleigh(unittest.TestCase): self.cube = Geometry(water_wcsim) self.cube.add_solid(Solid(box(100,100,100), water_wcsim, water_wcsim)) self.cube.pmtids = [0] - self.sim = Simulation(self.cube, bvh_bits=4, geant4_processes=0, - use_cache=False) + self.cube.build(use_cache=False) + self.sim = Simulation(self.cube, geant4_processes=0) + nphotons = 100000 pos = np.tile([0,0,0], (nphotons,1)).astype(np.float32) dir = np.tile([0,0,1], (nphotons,1)).astype(np.float32) |