merge

1 files changed, 34 insertions, 32 deletions

diff --git a/fileio/root.py b/fileio/root.py
index 004484f..5bebda0 100644
--- a/fileio/root.py
+++ b/fileio/root.py
@@ -3,7 +3,8 @@ import os.path
 
 ROOT.gROOT.ProcessLine('.L '+os.path.join(os.path.dirname(__file__), 'root.C+g'))
 
-from ROOT import Event
+import ROOT
+import event
 
 class RootWriter(object):
     def __init__(self, filename):
@@ -14,42 +15,43 @@ class RootWriter(object):
         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, 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_pos.SetXYZ(*pyev.gen_position)
+        self.ev.mc.gen_dir.SetXYZ(*pyev.gen_direction)
+        self.ev.mc.gen_total_energy = pyev.gen_total_energy
 
-    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
+        if pyev.photon_start is not None:
+            photons = pyev.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:
+                              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, 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_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 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'])
+        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].pos.SetXYZ(*subtrack.position)
+                self.ev.mc.subtrack[i].dir.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):