merge

1 files changed, 98 insertions, 88 deletions

diff --git a/fileio/root.py b/fileio/root.py
index 4c9d9bb..507c249 100644
--- a/fileio/root.py
+++ b/fileio/root.py
@@ -14,61 +14,66 @@ def tvector3_to_ndarray(vec):
 def make_photon_with_arrays(size):
     '''Returns a new chroma.event.Photons object for `size` number of
     photons with empty arrays set for all the photon attributes.'''
-    return event.Photons(positions=np.empty((size,3), dtype=np.float32),
-                         directions=np.empty((size,3), dtype=np.float32),
-                         polarizations=np.empty((size,3), dtype=np.float32),
+    return event.Photons(pos=np.empty((size,3), dtype=np.float32),
+                         dir=np.empty((size,3), dtype=np.float32),
+                         pol=np.empty((size,3), dtype=np.float32),
                          wavelengths=np.empty(size, dtype=np.float32),
-                         times=np.empty(size, dtype=np.float32),
-                         histories=np.empty(size, dtype=np.uint32),
+                         t=np.empty(size, dtype=np.float32),
+                         flags=np.empty(size, dtype=np.uint32),
                          last_hit_triangles=np.empty(size, dtype=np.int32))
 
+def root_vertex_to_python_vertex(vertex):
+    "Returns a chroma.event.Vertex object from a root Vertex object."
+    return event.Vertex(str(vertex.particle_name),
+                        tvector3_to_ndarray(vertex.pos),
+                        tvector3_to_ndarray(vertex.dir),
+                        tvector3_to_ndarray(vertex.pol),
+                        vertex.ke,
+                        vertex.t0)
 
 def root_event_to_python_event(ev):
     '''Returns a new chroma.event.Event object created from the
     contents of the ROOT event `ev`.'''
-    pyev = event.Event(ev.event_id)
-
-    # MC
-    pyev.particle_name = str(ev.mc.particle)
-    pyev.gen_position = tvector3_to_ndarray(ev.mc.gen_position)
-    pyev.gen_direction = tvector3_to_ndarray(ev.mc.gen_direction)
-    pyev.gen_total_energy = ev.mc.gen_total_energy
-
-    pyev.nphoton = ev.mc.nphoton
-
-    for subtrack in ev.mc.subtrack:
-        pysubtrack = event.Subtrack(str(subtrack.particle),
-                                    tvector3_to_ndarray(subtrack.position),
-                                    tvector3_to_ndarray(subtrack.direction),
-                                    subtrack.start_time,
-                                    subtrack.total_energy)
-        pyev.subtracks.append(pysubtrack)
-
-    # photon start
-    if ev.mc.photon_start.size() > 0:
-        photons = make_photon_with_arrays(ev.mc.photon_start.size())
-        ROOT.get_photons(ev.mc.photon_start, photons.positions.ravel(), photons.directions.ravel(),
-                         photons.polarizations.ravel(), photons.wavelengths, photons.times,
-                         photons.histories, photons.last_hit_triangles)
-        pyev.photon_start = photons
-
-    # photon stop
-    if ev.mc.photon_stop.size() > 0:
-        photons = make_photon_with_arrays(ev.mc.photon_stop.size())
-        ROOT.get_photons(ev.mc.photon_stop, photons.positions.ravel(), photons.directions.ravel(),
-                         photons.polarizations.ravel(), photons.wavelengths, photons.times,
-                         photons.histories, photons.last_hit_triangles)
-        pyev.photon_stop = photons
-
-    # hits
-    max_channel_id = ev.max_channel_id
-    hit = np.empty(shape=max_channel_id+1, dtype=np.int32)
-    t = np.empty(shape=max_channel_id+1, dtype=np.float32)
-    q = np.empty(shape=max_channel_id+1, dtype=np.float32)
-    histories = np.empty(shape=max_channel_id+1, dtype=np.uint32)
-
-    ev.get_channels(max_channel_id+1, hit, t, q, histories)
-    pyev.channels = event.Channels(hit.astype(bool), t, q, histories)
+    pyev = event.Event(ev.id)
+    pyev.primary_vertex = root_vertex_to_python_vertex(ev.primary_vertex)
+
+    for vertex in ev.vertices:
+        pyev.vertices.append(root_vertex_to_python_vertex(vertex))
+
+    # photon begin
+    if ev.photons_beg.size() > 0:
+        photons = make_photon_with_arrays(ev.photons_beg.size())
+        ROOT.get_photons(ev.photons_beg,
+                         photons.pos.ravel(),
+                         photons.dir.ravel(),
+                         photons.pol.ravel(),
+                         photons.wavelengths,
+                         photons.t,
+                         photons.last_hit_triangles,
+                         photons.flags)
+        pyev.photons_beg = photons
+
+    # photon end
+    if ev.photons_end.size() > 0:
+        photons = make_photon_with_arrays(ev.photons_end.size())
+        ROOT.get_photons(ev.photons_end,
+                         photons.pos.ravel(),
+                         photons.dir.ravel(),
+                         photons.pol.ravel(),
+                         photons.wavelengths,
+                         photons.t,
+                         photons.last_hit_triangles,
+                         photons.flags)
+        pyev.photons_end = photons
+
+    # channels
+    hit = np.empty(ev.nchannels, dtype=np.int32)
+    t = np.empty(ev.nchannels, dtype=np.float32)
+    q = np.empty(ev.nchannels, dtype=np.float32)
+    flags = np.empty(ev.nchannels, dtype=np.uint32)
+
+    ROOT.get_channels(ev, hit, t, q, flags)
+    pyev.channels = event.Channels(hit.astype(bool), t, q, flags)
     return pyev
     
 class RootReader(object):
@@ -140,48 +145,53 @@ class RootWriter(object):
         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.mc.photon_start,
-                              len(photons.positions), 
-                              np.ravel(photons.positions),
-                              np.ravel(photons.directions),
-                              np.ravel(photons.polarizations),
-                              photons.wavelengths, photons.times,
-                              photons.histories, photons.last_hit_triangles)
-        if pyev.photon_stop is not None:
-            photons = pyev.photon_stop
-            ROOT.fill_photons(self.ev.mc.photon_stop,
-                              len(photons.positions), 
-                              np.ravel(photons.positions),
-                              np.ravel(photons.directions),
-                              np.ravel(photons.polarizations),
-                              photons.wavelengths, photons.times,
-                              photons.histories, photons.last_hit_triangles)
-
-        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.channels.t), pyev.channels.t, pyev.channels.q, pyev.channels.histories)
+        "Write an event.Event object to the ROOT tree as a ROOT.Event object."
+        self.ev.id = pyev.id
+
+        self.ev.primary_vertex.particle_name = pyev.primary_vertex.particle_name
+        self.ev.primary_vertex.pos.SetXYZ(*pyev.primary_vertex.pos)
+        self.ev.primary_vertex.dir.SetXYZ(*pyev.primary_vertex.dir)
+        if pyev.primary_vertex.pol is not None:
+            self.ev.primary_vertex.pol.SetXYZ(*pyev.primary_vertex.pol)
+        self.ev.primary_vertex.ke = pyev.primary_vertex.ke
+
+        if pyev.photons_beg is not None:
+            photons = pyev.photons_beg
+            ROOT.fill_photons(self.ev.photons_beg,
+                              len(photons.pos),
+                              photons.pos.ravel(),
+                              photons.dir.ravel(),
+                              photons.pol.ravel(),
+                              photons.wavelengths, photons.t,
+                              photons.last_hit_triangles, photons.flags)
+
+        if pyev.photons_end is not None:
+            photons = pyev.photons_end
+            ROOT.fill_photons(self.ev.photons_end,
+                              len(photons.pos),
+                              photons.pos.ravel(),
+                              photons.dir.ravel(),
+                              photons.pol.ravel(),
+                              photons.wavelengths, photons.t,
+                              photons.last_hit_triangles, photons.flags)
+
+        self.ev.vertices.resize(0)
+        if pyev.vertices is not None:
+            self.ev.vertices.resize(len(pyev.vertices))
+            for i, vertex in enumerate(pyev.vertices):
+                self.ev.vertices[i].particle_name = vertex.particle_name
+                self.ev.vertices[i].pos.SetXYZ(*vertex.pos)
+                self.ev.vertices[i].dir.SetXYZ(*vertex.dir)
+                if vertex.pol is not None:
+                    self.ev.vertices[i].pol.SetXYZ(*vertex.pol)
+                self.ev.vertices[i].ke = vertex.ke
+                self.ev.vertices[i].t0 = vertex.t0
+
+        if pyev.channels is not None:
+            ROOT.fill_channels(self.ev, np.count_nonzero(pyev.channels.hit), np.arange(len(pyev.channels.t))[pyev.channels.hit].astype(np.int32), pyev.channels.t, pyev.channels.q, pyev.channels.flags, len(pyev.channels.hit))
+
         self.T.Fill()
 
     def close(self):
         self.T.Write()
         self.file.Close()
-