import ROOT import os.path ROOT.gROOT.ProcessLine('.L '+os.path.join(os.path.dirname(__file__), 'root.C+g')) from ROOT import Event 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, subtracks=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']) self.ev.mc.subtrack.resize(0) if subtracks is not None: self.ev.mc.subtrack.resize(len(subtracks)) for i, subtrack in enumerate(subtracks): self.ev.mc.subtrack[i].name = subtrack['name'] self.ev.mc.subtrack[i].pos.SetXYZ(*subtrack['pos']) self.ev.mc.subtrack[i].dir.SetXYZ(*subtrack['dir']) self.ev.mc.subtrack[i].total_energy = subtrack['total_e'] 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()