diff options
Diffstat (limited to 'sim.py')
| -rwxr-xr-x | sim.py | 37 |
1 files changed, 5 insertions, 32 deletions
@@ -58,28 +58,6 @@ class GeneratorProcess(multiprocessing.Process): self.queue.put(photons) -def write_event(T, ev, event_id, hits, photon_start=None, photon_stop=None): - ev.event_id = event_id - if photon_start is not None: - photons = photon_start - root.fill_photons(ev, True, - len(photons['pos']), - np.ravel(photons['pos']), - np.ravel(photons['dir']), - np.ravel(photons['pol']), - photons['wavelength'], photons['t0']) - if photon_stop is not None: - photons = photon_stop - root.fill_photons(ev, False, - len(photons['pos']), - np.ravel(photons['pos']), - np.ravel(photons['dir']), - np.ravel(photons['pol']), - photons['wavelength'], photons['t0'], - photons['histories'], photons['last_hit_triangles']) - - root.fill_hits(ev, len(hits['t']), hits['t'], hits['q'], hits['history']) - T.Fill() # Allow profile decorator to exist, but do nothing if not running under kernprof try: @@ -156,14 +134,11 @@ def main(): gpu_worker.setup_daq(max(detector.pmtids)) # Create output file - f = ROOT.TFile(output_filename, 'RECREATE') - ev, T = root.make_tree('T') + writer = root.RootWriter(output_filename) # Set generator info - ev.mc.particle = options.particle - ev.mc.gen_pos.SetXYZ(*position) - ev.mc.gen_dir.SetXYZ(*direction) - ev.mc.gen_total_e = options.energy + writer.set_generated_particle(name=options.particle, position=position, + direction=direction, total_e=options.energy) print >>sys.stderr, 'Starting simulation...' start_sim = time.time() @@ -189,16 +164,14 @@ def main(): else: photon_stop = None - write_event(T, ev, i, hits, - photon_start=photon_start, photon_stop=photon_stop) + writer.write_event(i, hits, photon_start=photon_start, photon_stop=photon_stop) if i % 10 == 0: print >>sys.stderr, "\rEvent:", i, end_sim = time.time() print >>sys.stderr, "\rEvent:", options.nevents - 1 - T.Write() - f.Close() + writer.close() print >>sys.stderr, 'Done. %1.1f events/sec, %1.0f photons/sec.' % (options.nevents/(end_sim - start_sim), nphotons/(end_sim - start_sim)) |