1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
|
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()
|
