Period key in event viewer flips between charge and time display

1 files changed, 23 insertions, 2 deletions

diff --git a/chroma/camera.py b/chroma/camera.py
index a7f768b..aa4983c 100644
--- a/chroma/camera.py
+++ b/chroma/camera.py
@@ -681,6 +681,7 @@ class EventViewer(Camera):
     def __init__(self, geometry, filename, **kwargs):
         Camera.__init__(self, geometry, **kwargs)
         self.rr = RootReader(filename)
+        self.show_charge = True
 
     def render_particle_track(self):
         x = 10.0
@@ -706,8 +707,19 @@ class EventViewer(Camera):
         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()))
-        self.gpu_geometry.color_solids(hit, solid_colors)
+        if self.show_charge:
+            channel_color = map_to_color(q, range=(q[hit].min(),q[hit].max()))
+            print 'charge'
+        else:
+            channel_color = map_to_color(t, range=(t[hit].min(),t[hit].mean()))
+            print 'time'
+        solid_hit = np.zeros(len(self.geometry.mesh.triangles), dtype=np.bool)
+        solid_color = np.zeros(len(self.geometry.mesh.triangles), dtype=np.uint32)
+
+        solid_hit[self.geometry.channel_index_to_solid_id] = hit
+        solid_color[self.geometry.channel_index_to_solid_id] = channel_color
+
+        self.gpu_geometry.color_solids(solid_hit, solid_color)
 
     def process_event(self, event):
         if event.type == KEYDOWN:
@@ -738,6 +750,15 @@ class EventViewer(Camera):
 
                     self.update()
                 return
+            elif event.key == K_PERIOD:
+                self.show_charge = not self.show_charge # flip bit
+                self.color_hit_pmts()
+
+                if self.ev.photons_beg is not None:
+                    self.render_particle_track()
+
+                self.update()
+                return
 
         Camera.process_event(self, event)