diff options
Diffstat (limited to 'fileio/root.py')
| -rw-r--r-- | fileio/root.py | 54 |
1 files changed, 42 insertions, 12 deletions
diff --git a/fileio/root.py b/fileio/root.py index e43f5d4..0b80745 100644 --- a/fileio/root.py +++ b/fileio/root.py @@ -5,15 +5,45 @@ ROOT.gROOT.ProcessLine('.L '+os.path.join(os.path.dirname(__file__), 'root.C+g') from ROOT import Event -fill_photons = ROOT.fill_photons -fill_hits = ROOT.fill_hits - -def make_tree(name, desc=''): - '''Create a ROOT tree for holding event information. - - Returns tuple of Event object for filling and TTree. - ''' - tree = ROOT.TTree(name, desc) - ev = ROOT.Event() - tree.Branch('ev', ev) - return ev, tree +class RootWriter(object): + def __init__(self, filename): + self.filename = filename + self.file = ROOT.TFile(filename, 'RECREATE') + + self.T = ROOT.TTree('T', 'Chroma events') + self.ev = ROOT.Event() + self.T.Branch('ev', self.ev) + + def set_generated_particle(self, name, position, direction, total_e): + self.ev.mc.particle = name + self.ev.mc.gen_pos.SetXYZ(*position) + self.ev.mc.gen_dir.SetXYZ(*direction) + self.ev.mc.gen_total_e = total_e + + def write_event(self, event_id, hits, photon_start=None, photon_stop=None): + self.ev.event_id = event_id + if photon_start is not None: + photons = photon_start + ROOT.fill_photons(self.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(self.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(self.ev, len(hits['t']), hits['t'], hits['q'], hits['history']) + self.T.Fill() + + def close(self): + self.T.Write() + self.file.Close() + |