merge

2 files changed, 191 insertions, 204 deletions

diff --git a/fileio/root.C b/fileio/root.C
index 07e50a1..2b39b1b 100644
--- a/fileio/root.C
+++ b/fileio/root.C
@@ -3,166 +3,143 @@
 #include <TTree.h>
 #include <string>
 
+struct Vertex {
+  std::string particle_name;
+  TVector3 pos;
+  TVector3 dir;
+  TVector3 pol;
+  double ke;
+  double t0;
+
+  ClassDef(Vertex, 1);
+};
+
 struct Photon {
-  double time;
-  TVector3 position;
-  TVector3 direction;
-  TVector3 polarization;
+  double t;
+  TVector3 pos;
+  TVector3 dir;
+  TVector3 pol;
   double wavelength; // nm
-  unsigned int history;
+  unsigned int flag;
   int last_hit_triangle;
 
   ClassDef(Photon, 1);
 };
 
-struct Track {
-  std::string particle;
-  TVector3 position;
-  TVector3 direction;
-  double start_time;
-  double total_energy;
+struct Channel {
+  Channel() : id(-1), t(-1e9), q(-1e9) { };
+  int id;
+  double t;
+  double q;
+  unsigned int flag;
 
-  ClassDef(Track, 1);
+  ClassDef(Channel, 1);
 };
 
-struct MC {
-  std::string particle;
-  TVector3 gen_position;
-  TVector3 gen_direction;
-  double gen_total_energy;
-
-  int nphoton;
-  
-  std::vector<Track> subtrack;
-  std::vector<Photon> photon_start;
-  std::vector<Photon> photon_stop;
+struct Event {
+  int id;
+  unsigned int nhit;
+  unsigned int nchannels;
 
-  ClassDef(MC, 1);
-};
+  Vertex primary_vertex;
 
-struct Channel {
-  Channel() : channel_id(-1), time(-9999.0), charge(-9999.0) { };
-  int channel_id;
-  double time;
-  double charge;
-  unsigned int mc_history;
+  std::vector<Vertex> vertices;
+  std::vector<Photon> photons_beg;
+  std::vector<Photon> photons_end;
+  std::vector<Channel> channels;
 
-  ClassDef(Channel, 1);
+  ClassDef(Event, 1);
 };
-  
-struct Event {
-  int event_id;
-  MC mc;
-  int nhit;
-  int max_channel_id;
-  std::vector<Channel> channel;
 
-  ClassDef(Event, 1);
+void fill_channels(Event *ev, unsigned int nhit, unsigned int *ids, float *t,
+		   float *q, unsigned int *flags, unsigned int nchannels)
+{
+  ev->nhit = 0;
+  ev->nchannels = nchannels;
+  ev->channels.resize(0);
 
-  // Populate arrays of length nentries with hit, time, and charge
-  // information, indexed by channel ID
-  void get_channels(unsigned int nentries, int *hit, float *time,
-		    float *charge, unsigned int *mc_history=0)
-  {
-    for (unsigned int i=0; i < nentries; i++) {
-      hit[i] = 0;
-      time[i] = -1e9f;
-      charge[i] = -1e9f;
-      if (mc_history)
-	mc_history[i] = 0;
-    }
+  Channel ch;
+  unsigned int id;
+  for (unsigned int i=0; i < nhit; i++) {
+      ev->nhit++;
+      id = ids[i];
+      ch.id = id;
+      ch.t = t[id];
+      ch.q = q[id];
+      ch.flag = flags[id];
+      ev->channels.push_back(ch);
+  }
+}
+
+void get_channels(Event *ev, int *hit, float *t, float *q, unsigned int *flags)
+{
+  for (unsigned int i=0; i < ev->nchannels; i++) {
+    hit[i] = 0;
+    t[i] = -1e9f;
+    q[i] = -1e9f;
+    flags[i] = 0;
+  }
 
-    for (unsigned int i=0; i < channel.size(); i++) {
-      unsigned int channel_id = channel[i].channel_id;
+  unsigned int id;
+  for (unsigned int i=0; i < ev->channels.size(); i++) {
+    id = ev->channels[i].id;
 
-      if (channel_id < nentries) {
-	hit[channel_id] = 1;
-	time[channel_id] = channel[i].time;
-	charge[channel_id] = channel[i].charge;
-	if (mc_history)
-	  mc_history[channel_id] = channel[i].mc_history;
-      }
+    if (id < ev->nchannels) {
+      hit[id] = 1;
+      t[id] = ev->channels[i].t;
+      q[id] = ev->channels[i].q;
+      flags[id] = ev->channels[i].flag;
     }
   }
+}
 
-};
-
-void get_photons(const std::vector<Photon> &photons, float *positions,
-		 float *directions, float *polarizations, float *wavelengths,
-		 float *times, unsigned int *histories, int *last_hit_triangles)
+void get_photons(const std::vector<Photon> &photons, float *pos, float *dir,
+		 float *pol, float *wavelengths, float *t,
+		 int *last_hit_triangles, unsigned int *flags)
 {
   for (unsigned int i=0; i < photons.size(); i++) {
     const Photon &photon = photons[i];
-    positions[3*i] = photon.position.X();
-    positions[3*i+1] = photon.position.Y();
-    positions[3*i+2] = photon.position.Z();
+    pos[3*i] = photon.pos.X();
+    pos[3*i+1] = photon.pos.Y();
+    pos[3*i+2] = photon.pos.Z();
 
-    directions[3*i] = photon.direction.X();
-    directions[3*i+1] = photon.direction.Y();
-    directions[3*i+2] = photon.direction.Z();
+    dir[3*i] = photon.dir.X();
+    dir[3*i+1] = photon.dir.Y();
+    dir[3*i+2] = photon.dir.Z();
     
-    polarizations[3*i] = photon.polarization.X();
-    polarizations[3*i+1] = photon.polarization.Y();
-    polarizations[3*i+2] = photon.polarization.Z();
+    pol[3*i] = photon.pol.X();
+    pol[3*i+1] = photon.pol.Y();
+    pol[3*i+2] = photon.pol.Z();
 
     wavelengths[i] = photon.wavelength;
-    times[i] = photon.time;
-    histories[i] = photon.history;
+    t[i] = photon.t;
+    flags[i] = photon.flag;
     last_hit_triangles[i] = photon.last_hit_triangle;
   }
 }
 		 
-
 void fill_photons(std::vector<Photon> &photons,
 		  unsigned int nphotons, float *pos, float *dir,
-		  float *pol, float *wavelength, float *t0,
-		  unsigned int *histories=0, int *last_hit_triangle=0)
+		  float *pol, float *wavelengths, float *t,
+		  int *last_hit_triangles, unsigned int *flags)
 {
   photons.resize(nphotons);
   
   for (unsigned int i=0; i < nphotons; i++) {
     Photon &photon = photons[i];
-    photon.time = t0[i];
-    photon.position.SetXYZ(pos[3*i], pos[3*i + 1], pos[3*i + 2]);
-    photon.direction.SetXYZ(dir[3*i], dir[3*i + 1], dir[3*i + 2]);
-    photon.polarization.SetXYZ(pol[3*i], pol[3*i + 1], pol[3*i + 2]);
-    photon.wavelength = wavelength[i];
-    if (histories)
-      photon.history = histories[i];
-    else
-      photon.history = 0;
-
-    if (last_hit_triangle)
-      photon.last_hit_triangle = last_hit_triangle[i];
-    else
-      photon.last_hit_triangle = -1;
-  }
-}
-		  
-
-void fill_hits(Event *ev, unsigned int nchannels, float *time,
-	       float *charge, unsigned int *history)
-{
-  ev->channel.resize(0);
-  ev->nhit = 0;
-  ev->max_channel_id = nchannels - 1;
+    photon.t = t[i];
+    photon.pos.SetXYZ(pos[3*i], pos[3*i + 1], pos[3*i + 2]);
+    photon.dir.SetXYZ(dir[3*i], dir[3*i + 1], dir[3*i + 2]);
+    photon.pol.SetXYZ(pol[3*i], pol[3*i + 1], pol[3*i + 2]);
+    photon.wavelength = wavelengths[i];
+    photon.last_hit_triangle = last_hit_triangles[i];
+    photon.flag = flags[i];
 
-  Channel ch;
-  for (unsigned int i=0; i < nchannels; i++) {
-    if (time[i] < 1e8) {
-      ev->nhit++;
-      ch.channel_id = i;
-      ch.time = time[i];
-      ch.charge = charge[i];
-      ch.mc_history = history[i];
-      ev->channel.push_back(ch);
-    }
   }
 }
 
-
 #ifdef __MAKECINT__
-#pragma link C++ class vector<Track>;
+#pragma link C++ class vector<Vertex>;
 #pragma link C++ class vector<Photon>;
 #pragma link C++ class vector<Channel>;
 #endif
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()
-