Modify data structure to hold photon start and stop vertices separately.

Add --save-photon-start and --save-photon-stop options to sim.py to
save these vertices.  Otherwise, only hit information is recorded.

3 files changed, 79 insertions, 25 deletions

diff --git a/io/root.C b/io/root.C
index dfcb3c0..43f4e73 100644
--- a/io/root.C
+++ b/io/root.C
@@ -9,23 +9,19 @@ struct Photon {
   TVector3 dir;
   TVector3 pol;
   double wavelength; // nm
-  int state;
-  int last_triangle;
-  int last_solid;
+  int history;
+  int last_hit_triangle;
 };
 
 struct MC {
   std::string particle;
   TVector3 gen_pos;
   TVector3 gen_dir;
-  double gen_ke;
+  double gen_total_e;
 
-  std::vector<Photon> photon;
+  std::vector<Photon> photon_start;
+  std::vector<Photon> photon_stop;
 
-  void clear() {
-    particle = "none";
-    photon.clear();
-  }
 };
 
 struct Channel {
@@ -41,23 +37,42 @@ struct Event {
   int nhit;
   std::vector<Channel> channel;
 
-  void clear() {
-    mc.clear();
-    channel.clear();
-  }
 };
 
-void fill_event(TTree *T,
-		Event *ev, int event_id, double *gen_pos, int nchannels, 
-		float *channel_times)
+void fill_photons(Event *ev, bool start, 
+		  unsigned int nphotons, float *pos, float *dir,
+		  float *pol, float *wavelength, float *t0,
+		  int *histories=0, int *last_hit_triangle=0)
+{
+  std::vector<Photon> &photons = start ? ev->mc.photon_start : ev->mc.photon_stop;
+  photons.resize(nphotons);
+  
+  for (unsigned int i=0; i < nphotons; i++) {
+    Photon &photon = photons[i];
+    photon.t = t0[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 = 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 *channel_times)
 {
-  ev->clear();
-  ev->event_id = event_id;
-  MC &mc = ev->mc;
-  mc.gen_pos.SetXYZ(gen_pos[0], gen_pos[1], gen_pos[2]);
+  ev->channel.resize(0);
   ev->nhit = 0;
   Channel ch;
-  for (int i=0; i < nchannels; i++) {
+  for (unsigned int i=0; i < nchannels; i++) {
     if (channel_times[i] < 1e8) {
       ev->nhit++;
       ch.channel_id = i;
@@ -66,7 +81,6 @@ void fill_event(TTree *T,
       ev->channel.push_back(ch);
     }
   }
-  T->Fill();
 }
 
 
diff --git a/io/root.py b/io/root.py
index d970d1d..e43f5d4 100644
--- a/io/root.py
+++ b/io/root.py
@@ -5,7 +5,8 @@ ROOT.gROOT.ProcessLine('.L '+os.path.join(os.path.dirname(__file__), 'root.C+g')
 
 from ROOT import Event
 
-fill_event = ROOT.fill_event
+fill_photons = ROOT.fill_photons
+fill_hits = ROOT.fill_hits
 
 def make_tree(name, desc=''):
     '''Create a ROOT tree for holding event information.
diff --git a/sim.py b/sim.py
index 44fb1dd..0a1f9e6 100755
--- a/sim.py
+++ b/sim.py
@@ -52,6 +52,29 @@ class GeneratorThread(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), hits)
+    T.Fill()
+
 def main():
     parser = optparse.OptionParser('%prog')
     parser.add_option('-b', type='int', dest='nbits', default=10)
@@ -64,6 +87,12 @@ def main():
     parser.add_option('--energy', type='float', dest='energy', default=100.0)
     parser.add_option('--pos', type='string', dest='pos', default='(0,0,0)')
     parser.add_option('--dir', type='string', dest='dir', default='(1,0,0)')
+    parser.add_option('--save-photon-start', action='store_true', 
+                      dest='save_photon_start', default=False,
+                      help='Save initial photon vertices to disk')
+    parser.add_option('--save-photon-stop', action='store_true', 
+                      dest='save_photon_stop', default=False,
+                      help='Save final photon vertices to disk')
 
     options, args = parser.parse_args()
     if len(args) != 1:
@@ -125,8 +154,18 @@ def main():
         gpu_worker.propagate()
         gpu_worker.run_daq()
         hit_times = gpu_worker.get_hits()
-        root.fill_event(T, ev, i-1, position, len(hit_times), hit_times)
-        
+
+        if options.save_photon_start:
+            photon_start = photons
+        else:
+            photon_start = None
+        if options.save_photon_stop:
+            photon_stop = gpu_worker.get_photons()
+        else:
+            photon_stop = None
+            
+        write_event(T, ev, i, hit_times, 
+                    photon_start=photon_start, photon_stop=photon_stop)
         if i % 10 == 0:
             print >>sys.stderr, "\rEvent:", i,