diff options
| -rw-r--r-- | event.py | 1 | ||||
| -rw-r--r-- | fileio/root.C | 49 | ||||
| -rw-r--r-- | fileio/root.py | 9 | ||||
| -rw-r--r-- | generator/g4gen.py | 4 | ||||
| -rwxr-xr-x | sim.py | 123 |
5 files changed, 112 insertions, 74 deletions
@@ -46,6 +46,7 @@ class Event(object): self.gen_total_energy = gen_total_energy self.subtracks = [] + self.nphoton = 0 self.photon_start = None self.photon_stop = None self.channels = None diff --git a/fileio/root.C b/fileio/root.C index 8f7d0a9..7052a71 100644 --- a/fileio/root.C +++ b/fileio/root.C @@ -4,10 +4,10 @@ #include <string> struct Photon { - double t; - TVector3 pos; - TVector3 dir; - TVector3 pol; + double time; + TVector3 position; + TVector3 direction; + TVector3 polarization; double wavelength; // nm unsigned int history; int last_hit_triangle; @@ -15,19 +15,20 @@ struct Photon { struct Track { std::string particle; - double t; - TVector3 pos; - TVector3 dir; + double time; + TVector3 position; + TVector3 direction; double start_time; double total_energy; }; struct MC { std::string particle; - TVector3 gen_pos; - TVector3 gen_dir; - double gen_total_e; + TVector3 gen_position; + TVector3 gen_direction; + double gen_total_energy; + int nphoton; std::vector<Track> subtrack; std::vector<Photon> photon_start; @@ -36,10 +37,10 @@ struct MC { }; struct Channel { - Channel() : channel_id(-1), t(-9999.0), q(-9999.0) { }; + Channel() : channel_id(-1), time(-9999.0), charge(-9999.0) { }; int channel_id; - double t; - double q; + double time; + double charge; unsigned int mc_history; }; @@ -65,8 +66,8 @@ struct Event { if (channel_id < nentries) { hit[channel_id] = 1; - time[channel_id] = channel[i].t; - charge[channel_id] = channel[i].q; + time[channel_id] = channel[i].time; + charge[channel_id] = channel[i].charge; } } } @@ -83,10 +84,10 @@ void fill_photons(Event *ev, bool start, for (unsigned int i=0; i < nphotons; i++) { Photon &photon = photons[i]; - photon.t = t0[i]; - photon.pos.SetXYZ(pos[3*i], pos[3*i + 1], pos[3*i + 2]); - photon.dir.SetXYZ(dir[3*i], dir[3*i + 1], dir[3*i + 2]); - photon.pol.SetXYZ(pol[3*i], pol[3*i + 1], pol[3*i + 2]); + photon.time = t0[i]; + photon.position.SetXYZ(pos[3*i], pos[3*i + 1], pos[3*i + 2]); + photon.direction.SetXYZ(dir[3*i], dir[3*i + 1], dir[3*i + 2]); + photon.polarization.SetXYZ(pol[3*i], pol[3*i + 1], pol[3*i + 2]); photon.wavelength = wavelength[i]; if (histories) photon.history = histories[i]; @@ -101,18 +102,18 @@ void fill_photons(Event *ev, bool start, } -void fill_hits(Event *ev, unsigned int nchannels, float *t, - float *q, unsigned int *history) +void fill_hits(Event *ev, unsigned int nchannels, float *time, + float *charge, unsigned int *history) { ev->channel.resize(0); ev->nhit = 0; Channel ch; for (unsigned int i=0; i < nchannels; i++) { - if (t[i] < 1e8) { + if (time[i] < 1e8) { ev->nhit++; ch.channel_id = i; - ch.t = t[i]; - ch.q = q[i]; + ch.time = time[i]; + ch.charge = charge[i]; ch.mc_history = history[i]; ev->channel.push_back(ch); } diff --git a/fileio/root.py b/fileio/root.py index 5bebda0..77c23ff 100644 --- a/fileio/root.py +++ b/fileio/root.py @@ -20,9 +20,10 @@ class RootWriter(object): self.ev.event_id = pyev.event_id self.ev.mc.particle = pyev.particle_name - self.ev.mc.gen_pos.SetXYZ(*pyev.gen_position) - self.ev.mc.gen_dir.SetXYZ(*pyev.gen_direction) + self.ev.mc.gen_position.SetXYZ(*pyev.gen_position) + self.ev.mc.gen_direction.SetXYZ(*pyev.gen_direction) self.ev.mc.gen_total_energy = pyev.gen_total_energy + self.ev.mc.nphoton = pyev.nphoton if pyev.photon_start is not None: photons = pyev.photon_start @@ -46,8 +47,8 @@ class RootWriter(object): self.ev.mc.subtrack.resize(len(pyev.subtracks)) for i, subtrack in enumerate(pyev.subtracks): self.ev.mc.subtrack[i].name = subtrack.particle_name - self.ev.mc.subtrack[i].pos.SetXYZ(*subtrack.position) - self.ev.mc.subtrack[i].dir.SetXYZ(*subtrack.direction) + self.ev.mc.subtrack[i].position.SetXYZ(*subtrack.position) + self.ev.mc.subtrack[i].direction.SetXYZ(*subtrack.direction) self.ev.mc.subtrack[i].start_time = subtrack.start_time self.ev.mc.subtrack[i].total_energy = subtrack.total_energy diff --git a/generator/g4gen.py b/generator/g4gen.py index d4e79a1..4416c45 100644 --- a/generator/g4gen.py +++ b/generator/g4gen.py @@ -111,8 +111,8 @@ class G4Generator(object): else: # Create temporary subtrack for single primary particle subtracks = [event.Subtrack(particle_name=ev.particle_name, - position=ev.gen_pos, - direction=ev.gen_dir, + position=ev.gen_position, + direction=ev.gen_direction, start_time=0.0, total_energy=ev.gen_total_energy)] @@ -9,8 +9,10 @@ import detectors import optics import generator from generator import constant +import itertools import gpu from fileio import root + from tools import profile_if_possible, enable_debug_on_crash def pick_seed(): @@ -18,12 +20,77 @@ def pick_seed(): 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): + 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 for detector "%s" with %d bits...' % (self.detector, bvh_bits) + detector.build(bits=bvh_bits) + + 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)) + + def simulate(self, nevents, vertex_generator, keep_photon_start=False, keep_photon_stop=False, + run_daq=True): + return self.simulate_photons(nevents, self.photon_generator.generate_events(nevents, vertex_generator), + 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() + + yield ev + + def propagate_photons(self, photons, max_steps=10): + self.gpu_worker.load_photons(photons) + self.gpu_worker.propagate() + return self.gpu_worker.get_photons() + @profile_if_possible def main(): parser = optparse.OptionParser('%prog') parser.add_option('-b', type='int', dest='nbits', default=10) parser.add_option('-j', type='int', dest='device', default=None) - parser.add_option('-n', type='int', dest='nblocks', default=64) 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, @@ -55,13 +122,6 @@ def main(): position = np.array(eval(options.pos), dtype=float) direction = np.array(eval(options.dir), dtype=float) detector = detectors.find(options.detector) - if options.seed is None: - options.seed = pick_seed() - - print >>sys.stderr, 'RNG seed:', options.seed - # We have three generators to seed: numpy.random, GEANT4, and CURAND. - # The latter two are done below - np.random.seed(options.seed) print >>sys.stderr, 'Creating generator...' if options.particle == 'pi0': @@ -73,27 +133,12 @@ def main(): position=constant(position), direction=constant(direction), total_energy=constant(options.energy)) - detector_material = optics.water_wcsim - photon_generator = generator.photon.G4ParallelGenerator(options.ngenerators, detector_material, - base_seed=options.seed) - print >>sys.stderr, 'WARNING: ASSUMING DETECTOR IS WCSIM WATER!!' - - # Do this now so we can get ahead of the photon propagation - print >>sys.stderr, 'Starting GEANT4 generators...' - event_iterator = photon_generator.generate_events(options.nevents, vertex_generator) - - print >>sys.stderr, 'Creating BVH for detector "%s" with %d bits...' % (options.detector, options.nbits) - detector.build(bits=options.nbits) - - print >>sys.stderr, 'Initializing GPU...' - gpu_worker = gpu.GPU(options.device) - - print >>sys.stderr, 'Loading detector onto GPU...' - gpu_worker.load_geometry(detector) - - print >>sys.stderr, 'Initializing random numbers generators...' - gpu_worker.setup_propagate(seed=options.seed) - gpu_worker.setup_daq(max(detector.pmtids)) + + # 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) @@ -102,22 +147,12 @@ def main(): start_sim = time.time() nphotons = 0 - for i, ev in enumerate(event_iterator): - photons = ev.photon_start - assert len(photons.positions) > 0, 'GEANT4 generated event with no photons!' - - nphotons += len(photons.positions) - - gpu_worker.load_photons(photons) - - gpu_worker.propagate() - gpu_worker.run_daq() + for i, ev in enumerate(simulation.simulate(options.nevents, vertex_generator, + keep_photon_start=options.save_photon_start, + keep_photon_stop=options.save_photon_stop)): + assert ev.nphoton > 0, 'GEANT4 generated event with no photons!' + nphotons += ev.nphoton - ev.hits = gpu_worker.get_hits() - if not options.save_photon_start: - ev.photon_start = None - if options.save_photon_stop: - ev.photon_stop = gpu_worker.get_photons() writer.write_event(ev) if i % 10 == 0: |