Refactor ROOT file writing into fileio.root.RootWriter class

2 files changed, 47 insertions, 44 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()
+    
diff --git a/sim.py b/sim.py
index 48bf313..e123e5e 100755
--- a/sim.py
+++ b/sim.py
@@ -58,28 +58,6 @@ class GeneratorProcess(multiprocessing.Process):
             self.queue.put(photons)
 
 
-def write_event(T, ev, event_id, hits, photon_start=None, photon_stop=None):
-    ev.event_id = event_id
-    if photon_start is not None:
-        photons = photon_start
-        root.fill_photons(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(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(ev, len(hits['t']), hits['t'], hits['q'], hits['history'])
-    T.Fill()
 
 # Allow profile decorator to exist, but do nothing if not running under kernprof
 try:
@@ -156,14 +134,11 @@ def main():
     gpu_worker.setup_daq(max(detector.pmtids))
 
     # Create output file
-    f = ROOT.TFile(output_filename, 'RECREATE')
-    ev, T = root.make_tree('T')
+    writer = root.RootWriter(output_filename)
 
     # Set generator info
-    ev.mc.particle = options.particle
-    ev.mc.gen_pos.SetXYZ(*position)
-    ev.mc.gen_dir.SetXYZ(*direction)
-    ev.mc.gen_total_e = options.energy
+    writer.set_generated_particle(name=options.particle, position=position,
+                                  direction=direction, total_e=options.energy)
 
     print >>sys.stderr, 'Starting simulation...'
     start_sim = time.time()
@@ -189,16 +164,14 @@ def main():
         else:
             photon_stop = None
             
-        write_event(T, ev, i, hits,
-                    photon_start=photon_start, photon_stop=photon_stop)
+        writer.write_event(i, hits, photon_start=photon_start, photon_stop=photon_stop)
         if i % 10 == 0:
             print >>sys.stderr, "\rEvent:", i,
 
     end_sim = time.time()
     print >>sys.stderr, "\rEvent:", options.nevents - 1
 
-    T.Write()
-    f.Close()
+    writer.close()
 
     print >>sys.stderr, 'Done. %1.1f events/sec, %1.0f photons/sec.' % (options.nevents/(end_sim - start_sim), nphotons/(end_sim - start_sim))