photon object initializes any arguments not passed when initialized. camera.EventViewer uses new fileio.root.RootReader class.

2 files changed, 33 insertions, 33 deletions

diff --git a/camera.py b/camera.py
index e54c3b8..53c028e 100755
--- a/camera.py
+++ b/camera.py
@@ -14,6 +14,7 @@ from chroma.tools import timeit
 from chroma.transform import rotate
 from chroma.optics import vacuum, lambertian_surface
 import chroma.project as project
+from chroma.fileio.root import RootReader
 
 import pygame
 from pygame.locals import *
@@ -620,27 +621,14 @@ class Camera(Thread):
 class EventViewer(Camera):
     def __init__(self, geometry, filename, **kwargs):
         Camera.__init__(self, geometry, **kwargs)
-
-        import ROOT
-
-        self.f = ROOT.TFile(filename)
-        self.T = self.f.Get('T')
-        self.T.GetEntry(0)
-        self.nsolids = geometry.solid_id.max() + 1
+        self.rr = RootReader(filename)
 
     def color_hit_pmts(self):
         self.gpu.reset_colors()
 
-        hit = np.empty(self.nsolids, np.int32)
-        t = np.empty(self.nsolids, np.float32)
-        q = np.empty(self.nsolids, np.float32)
-
-        # self.nsolids has a weird data type that PyROOT doesn't understand
-        self.T.ev.get_channels(int(self.nsolids), hit, t, q)
-
-        # PyROOT prints warnings when we try to pass a bool array directly
-        # so we convert afterward
-        hit = hit.astype(np.bool)
+        hit = self.ev.channels.hit
+        t = self.ev.channels.t
+        q = self.ev.channels.q
 
         # Important: Compute range only with HIT channels
         solid_colors = map_to_color(q, range=(q[hit].min(),q[hit].max()))
@@ -650,18 +638,22 @@ class EventViewer(Camera):
     def process_event(self, event):
         if event.type == KEYDOWN:
             if event.key == K_PAGEUP:
-                entry = self.T.GetReadEntry()
-                if entry < self.T.GetEntries() - 1:
-                    self.T.GetEntry(entry+1)
+                try:
+                    self.ev = self.rr.next()
+                except StopIteration:
+                    pass
+                else:
                     self.color_hit_pmts()
-                    return
+                return
 
             elif event.key == K_PAGEDOWN:
-                entry = self.T.GetReadEntry()
-                if entry > 0:
-                    self.T.GetEntry(entry-1)
+                try:
+                    self.ev = self.rr.prev()
+                except StopIteration:
+                    pass
+                else:
                     self.color_hit_pmts()
-                    return
+                return
 
         Camera.process_event(self, event)
 
diff --git a/event.py b/event.py
index 657f414..633fe26 100644
--- a/event.py
+++ b/event.py
@@ -4,14 +4,10 @@ from chroma.sample import uniform_sphere
 
 class Photons(object):
     def __init__(self, positions, directions, wavelengths, polarizations=None, times=None, last_hit_triangles=None, histories=None):
-        self.positions = positions
-
         nphotons = len(positions)
 
-        assert len(directions) == nphotons
+        self.positions = positions
         self.directions = directions
-
-        assert len(wavelengths) == nphotons
         self.wavelengths = wavelengths
 
         if polarizations is not None:
@@ -24,9 +20,21 @@ class Photons(object):
             polarizations /= np.tile(np.apply_along_axis(np.linalg.norm,1,polarizations),[3,1]).transpose()
             self.polarizations = np.asarray(polarizations, order='C')
 
-        self.times = times
-        self.last_hit_triangles = last_hit_triangles
-        self.histories = histories
+        if times is None:
+            self.times = np.zeros(nphotons, dtype=np.float32)
+        else:
+            self.times = times
+
+        if last_hit_triangles is None:
+            self.last_hit_triangles = np.empty(nphotons, dtype=np.int32)
+            self.last_hit_triangles.fill(-1)
+        else:
+            self.last_hit_triangles = last_hit_triangles
+
+        if histories is None:
+            self.histories = np.zeros(nphotons, dtype=np.uint32)
+        else:
+            self.histories = histories
 
 def concatenate_photons(photons):
     '''Merge a list of Photons objects into one long list of photons'''