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
59
60
61
|
import ROOT
import os.path
ROOT.gROOT.ProcessLine('.L '+os.path.join(os.path.dirname(__file__), 'root.C+g'))
import ROOT
import chroma.event as 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 write_event(self, pyev):
'''Write an event.Event object to the ROOT tree as a ROOT.Event object.'''
self.ev.event_id = pyev.event_id
self.ev.mc.particle = pyev.particle_name
self.ev.mc.gen_position.SetXYZ(*pyev.gen_position)
self.ev.mc.gen_direction.SetXYZ(*pyev.gen_direction)
self.ev.mc.gen_total_energy = pyev.gen_total_energy
self.ev.mc.nphoton = pyev.nphoton
if pyev.photon_start is not None:
photons = pyev.photon_start
ROOT.fill_photons(self.ev, True,
len(photons.positions),
np.ravel(photons.positions),
np.ravel(photons.directions),
np.ravel(photons.polarizations),
photons.wavelengths, photons.times)
if pyev.photon_stop is not None:
photons = photon_stop
ROOT.fill_photons(self.ev, True,
len(photons.positions),
np.ravel(photons.positions),
np.ravel(photons.directions),
np.ravel(photons.polarizations),
photons.wavelengths, photons.times)
self.ev.mc.subtrack.resize(0)
if pyev.subtracks is not None:
self.ev.mc.subtrack.resize(len(pyev.subtracks))
for i, subtrack in enumerate(pyev.subtracks):
self.ev.mc.subtrack[i].name = subtrack.particle_name
self.ev.mc.subtrack[i].position.SetXYZ(*subtrack.position)
self.ev.mc.subtrack[i].direction.SetXYZ(*subtrack.direction)
self.ev.mc.subtrack[i].start_time = subtrack.start_time
self.ev.mc.subtrack[i].total_energy = subtrack.total_energy
ROOT.fill_hits(self.ev, len(pyev.hits.t), pyev.hits.t, pyev.hits.q, pyev.hits.histories)
self.T.Fill()
def close(self):
self.T.Write()
self.file.Close()
|
