add ability to toggle between viewing initial and final photon positions. try to use Image module to save screenshot if pygame screenshot doesn't work.

1 files changed, 26 insertions, 3 deletions

diff --git a/chroma/camera.py b/chroma/camera.py
index 9f69f1a..54e8ba7 100644
--- a/chroma/camera.py
+++ b/chroma/camera.py
@@ -249,7 +249,15 @@ class Camera(multiprocessing.Process):
             if not os.path.exists(path):
                 break
 
-        pygame.image.save(self.screen, path)
+        try:
+            pygame.image.save(self.screen, path)
+        except ImportError:
+            import Image
+            mode = 'RGBA'
+            data = self.screen.get_buffer()
+            im = Image.frombuffer(mode,self.size,data,'raw',mode,0,1)
+            im.save(path)
+
         print 'image saved to %s' % path
 
     def rotate(self, phi, n):
@@ -672,6 +680,8 @@ class EventViewer(Camera):
         self.rr = RootReader(filename)
         self.display_mode = EventViewer.CHARGE
         self.sum_mode = False
+        self.photon_display_iter = itertools.cycle(['beg','end'])
+        self.photon_display_mode = self.photon_display_iter.next()
 
     def render_particle_track(self):
         x = 10.0
@@ -680,9 +690,14 @@ class EventViewer(Camera):
                                       [0]*3, [0]*3)
         marker = Solid(pyramid, vacuum, vacuum)
 
+        if self.photon_display_mode == 'beg':
+            photons = self.ev.photons_beg
+        else:
+            photons = self.ev.photons_end
+
         geometry = Geometry()
-        sample_factor = max(1, len(self.ev.photons_beg.pos) / 10000)
-        for pos in self.ev.photons_beg.pos[::sample_factor]:
+        sample_factor = max(1, len(photons.pos) / 10000)
+        for pos in photons.pos[::sample_factor]:
             geometry.add_solid(marker, displacement=pos, rotation=make_rotation_matrix(np.random.uniform(0,2*np.pi), uniform_sphere()))
 
         geometry = create_geometry_from_obj(geometry)
@@ -715,6 +730,9 @@ class EventViewer(Camera):
     def color_hit_pmts(self):
         self.gpu_geometry.reset_colors()
 
+        if self.ev.channels is None:
+            return
+
         if self.sum_mode:
             hit = self.sum_hit
             t = self.sum_t
@@ -751,6 +769,11 @@ class EventViewer(Camera):
 
     def process_event(self, event):
         if event.type == KEYDOWN:
+            if event.key == K_t:
+                self.photon_display_mode = self.photon_display_iter.next()
+                self.render_particle_track()
+                self.update()
+
             if event.key == K_PAGEUP and not self.sum_mode:
                 try:
                     self.ev = self.rr.next()