diff options
Diffstat (limited to 'sim.py')
| -rwxr-xr-x | sim.py | 43 |
1 files changed, 41 insertions, 2 deletions
@@ -52,6 +52,29 @@ class GeneratorThread(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), hits) + T.Fill() + def main(): parser = optparse.OptionParser('%prog') parser.add_option('-b', type='int', dest='nbits', default=10) @@ -64,6 +87,12 @@ def main(): parser.add_option('--energy', type='float', dest='energy', default=100.0) parser.add_option('--pos', type='string', dest='pos', default='(0,0,0)') parser.add_option('--dir', type='string', dest='dir', default='(1,0,0)') + parser.add_option('--save-photon-start', action='store_true', + dest='save_photon_start', default=False, + help='Save initial photon vertices to disk') + parser.add_option('--save-photon-stop', action='store_true', + dest='save_photon_stop', default=False, + help='Save final photon vertices to disk') options, args = parser.parse_args() if len(args) != 1: @@ -125,8 +154,18 @@ def main(): gpu_worker.propagate() gpu_worker.run_daq() hit_times = gpu_worker.get_hits() - root.fill_event(T, ev, i-1, position, len(hit_times), hit_times) - + + if options.save_photon_start: + photon_start = photons + else: + photon_start = None + if options.save_photon_stop: + photon_stop = gpu_worker.get_photons() + else: + photon_stop = None + + write_event(T, ev, i, hit_times, + photon_start=photon_start, photon_stop=photon_stop) if i % 10 == 0: print >>sys.stderr, "\rEvent:", i, |