diff options
| -rwxr-xr-x | camera.py | 42 | ||||
| -rw-r--r-- | event.py | 24 |
2 files changed, 33 insertions, 33 deletions
@@ -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) @@ -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''' |