def __init__(self, geometry, filename, **kwargs):
Camera.__init__(self, geometry, **kwargs)
# This is really slow, so we do it here in the constructor to
# avoid slowing down the import of this module
from chroma.io.root import RootReader
self.rr = RootReader(filename)
self.ev = next(self.rr)
self.display_mode_iter = itertools.cycle(
['geo', 'charge', 'time', 'hit', 'dichroicon'])
self.display_mode = next(self.display_mode_iter)
self.sum_mode = False
self.photon_display_mode_iter = itertools.cycle(['none', 'beg', 'end'])
self.photon_display_mode = next(self.photon_display_mode_iter)
self.track_display_mode_iter = itertools.cycle(
['none', 'geant4', 'chroma', 'both'])
self.track_display_mode = next(self.track_display_mode_iter)
''' photons_max will randomly select at most that many photons
photons_max_steps truncates all photon tracks to that number of steps
photons_only_type can be set to 'cher', 'scint', or 'reemit'
photons_detected_only will show only detected photon tracks
photons_track_size controls the track size of the photons'''
self.photons_max = 1000
self.photons_max_steps = 20
self.photons_only_type = None
self.photons_detected_only = False
self.photons_track_size = 0.1
def __init__(self, geometry, filename, **kwargs):
Camera.__init__(self, geometry, **kwargs)
# This is really slow, so we do it here in the constructor to
# avoid slowing down the import of this module
from chroma.io.root import RootReader
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 __init__(self, geometry, filename, **kwargs):
Camera.__init__(self, geometry, **kwargs)
# This is really slow, so we do it here in the constructor to
# avoid slowing down the import of this module
from chroma.io.root import RootReader
self.rr = RootReader(filename)
self.ev = next(self.rr)
self.display_mode_iter = itertools.cycle(['geo','charge','time','hit','dichroicon'])
self.display_mode = next(self.display_mode_iter)
self.sum_mode = False
self.photon_display_mode_iter = itertools.cycle(['none','beg','end'])
self.photon_display_mode = next(self.photon_display_mode_iter)
self.track_display_mode_iter = itertools.cycle(['none','geant4','chroma','both'])
self.track_display_mode = next(self.track_display_mode_iter)
class EventViewer(Camera):
def __init__(self, geometry, filename, **kwargs):
Camera.__init__(self, geometry, **kwargs)
# This is really slow, so we do it here in the constructor to
# avoid slowing down the import of this module
from chroma.io.root import RootReader
self.rr = RootReader(filename)
self.ev = next(self.rr)
self.display_mode_iter = itertools.cycle(
['geo', 'charge', 'time', 'hit', 'dichroicon'])
self.display_mode = next(self.display_mode_iter)
self.sum_mode = False
self.photon_display_mode_iter = itertools.cycle(['none', 'beg', 'end'])
self.photon_display_mode = next(self.photon_display_mode_iter)
self.track_display_mode_iter = itertools.cycle(
['none', 'geant4', 'chroma', 'both'])
self.track_display_mode = next(self.track_display_mode_iter)
''' photons_max will randomly select at most that many photons
photons_max_steps truncates all photon tracks to that number of steps
photons_only_type can be set to 'cher', 'scint', or 'reemit'
photons_detected_only will show only detected photon tracks
photons_track_size controls the track size of the photons'''
self.photons_max = 1000
self.photons_max_steps = 20
self.photons_only_type = None
self.photons_detected_only = False
self.photons_track_size = 0.1
def render_photon_track(self,
geometry,
photon_track,
sz=1.0,
color='wavelength'):
origin = photon_track.pos[:-1]
extent = photon_track.pos[1:] - photon_track.pos[:-1]
perp1 = np.cross(origin, extent)
perp1 = np.inner(sz / 2.0 / np.linalg.norm(perp1, axis=1), perp1.T).T
perp2 = np.cross(perp1, extent)
perp2 = np.inner(sz / 2.0 / np.linalg.norm(perp2, axis=1), perp2.T).T
verts = [perp1 + perp2, -perp1 + perp2, perp1 - perp2, -perp1 - perp2]
bot = [vert + origin for vert in verts]
top = [vert + origin + extent for vert in verts]
vertices = [
origin, origin + extent, bot[0], top[0], bot[1], top[1], bot[2],
top[2], bot[3], top[3]
]
vertices = np.transpose(np.asarray(vertices, np.float32), (1, 0, 2))
triangles = np.asarray(
[[1, 3, 5], [1, 5, 7], [1, 7, 9], [1, 9, 3], [3, 2, 4], [5, 4, 6],
[7, 6, 8], [9, 8, 2], [2, 0, 0], [4, 0, 0], [6, 0, 0], [8, 0, 0],
[1, 5, 1], [1, 7, 1], [1, 9, 1], [1, 3, 1], [3, 4, 5], [5, 6, 7],
[7, 8, 9], [9, 2, 3], [2, 0, 4], [4, 0, 6], [6, 0, 8], [8, 0, 2]],
dtype=np.int32)
if color == 'wavelength':
r = np.asarray(np.interp(photon_track.wavelengths[:-1],
[300, 550, 800], [0, 0, 255]),
dtype=np.uint32)
g = np.asarray(np.interp(photon_track.wavelengths[:-1],
[300, 550, 800], [0, 255, 0]),
dtype=np.uint32)
b = np.asarray(np.interp(photon_track.wavelengths[:-1],
[300, 550, 800], [255, 0, 0]),
dtype=np.uint32)
colors = np.bitwise_or(
b, np.bitwise_or(np.left_shift(g, 8), np.left_shift(r, 16)))
else:
r, g, b = color
colors = np.full_like(photon_track.wavelengths[:-1],
((r << 16) | (g << 8) | b),
dtype=np.uint32)
markers = [
Solid(Mesh(v, triangles), vacuum, vacuum, color=c)
for v, c in zip(vertices, colors)
]
[geometry.add_solid(marker) for marker in markers]
def render_vertex(
self,
geometry,
vertex,
children=2,
sz=5.0,
colors={
'mu+': 0x50E0FF,
'mu-': 0xFFE050,
'e+': 0x0000FF,
'e-': 0xFF0000,
'gamma': 0x00FF00
}):
steps = np.vstack((vertex.steps.x, vertex.steps.y, vertex.steps.z)).T
for index in range(len(steps) - 1):
vec = steps[index + 1] - steps[index]
mag = np.linalg.norm(vec)
u = vec / mag
axis = np.cross(u, [0, 0, 1])
ang = np.arccos(np.dot(u, [0, 0, 1]))
rotmat = make_rotation_matrix(ang, axis)
x = sz / 2
y = x * np.sqrt(3) / 2
segment = make.linear_extrude([-x, 0, x], [-y, y, -y], mag,
[0, 0, 0], [0, 0, 0])
segment.vertices[:, 2] += mag / 2.0
marker = Solid(segment,
vacuum,
vacuum,
color=colors[vertex.particle_name]
if vertex.particle_name in colors else 0xAAAAAA)
geometry.add_solid(marker,
displacement=steps[index],
rotation=rotmat)
if children and vertex.children:
for child in vertex.children:
if isinstance(children, bool):
self.render_vertex(geometry, child, children)
else:
self.render_vertex(geometry, child, children - 1)
def render_photons(self, geometry, photons):
x = 10.0
h = x * np.sqrt(3) / 2
pyramid = make.linear_extrude([-x / 2, 0, x / 2],
[-h / 2, h / 2, -h / 2], h, [0] * 3,
[0] * 3)
marker = Solid(pyramid, vacuum, vacuum)
sample_factor = 1
subset = photons[::sample_factor]
for p, d, w in zip(subset.pos, subset.dir, subset.wavelengths):
geometry.add_solid(marker,
displacement=p,
rotation=gen_rot([0, 1, 0], d))
def render_mc_info(self):
#need to render photon tracking info if available
self.gpu_geometries = [self.gpu_geometry]
if self.sum_mode or self.ev is None:
return
if self.photon_display_mode == 'beg':
photons = self.ev.photons_beg
elif self.photon_display_mode == 'end':
photons = self.ev.photons_end
else:
photons = None
if photons is not None:
geometry = Geometry()
self.render_photons(geometry, photons)
geometry = create_geometry_from_obj(geometry)
gpu_geometry = gpu.GPUGeometry(geometry)
self.gpu_geometries.append(gpu_geometry)
if self.track_display_mode in ['geant4', 'both'
] and self.ev.vertices is not None:
geometry = Geometry()
any = False
for vertex in self.ev.vertices:
if vertex.steps:
any = True
self.render_vertex(geometry, vertex, children=True)
if any:
geometry = create_geometry_from_obj(geometry)
gpu_geometry = gpu.GPUGeometry(geometry)
self.gpu_geometries.append(gpu_geometry)
if self.track_display_mode in ['chroma', 'both'
] and self.ev.photon_tracks is not None:
geometry = Geometry()
print('Total Photons', len(self.ev.photon_tracks))
def has(flags, test):
return flags & test == test
tracks = self.ev.photon_tracks
if self.photons_detected_only:
detected = np.asarray([
has(track.flags[-1], event.SURFACE_DETECT)
for track in tracks
])
tracks = [t for t, m in zip(tracks, detected) if m]
cherenkov = np.asarray([
has(track.flags[0], event.CHERENKOV)
and not has(track.flags[-1], event.BULK_REEMIT)
for track in tracks
])
scintillation = np.asarray([
has(track.flags[0], event.SCINTILLATION)
and not has(track.flags[-1], event.BULK_REEMIT)
for track in tracks
])
reemission = np.asarray(
[has(track.flags[-1], event.BULK_REEMIT) for track in tracks])
if self.photons_only_type is not None:
if self.photons_only_type == 'cher':
selector = cherenkov
elif self.photons_only_type == 'scint':
selector = scintillation
elif self.photons_only_type == 'reemit':
selector = reemission
else:
raise Exception('Unknown only type: %s' % only)
tracks = [t for t, m in zip(tracks, selector) if m]
cherenkov = cherenkov[selector]
scintillation = scintillation[selector]
reemission = reemission[selector]
nphotons = len(tracks)
prob = self.photons_max / nphotons if self.photons_max is not None else 1.0
selector = np.random.random(len(tracks)) < prob
nphotons = np.count_nonzero(selector)
for i, track in ((i, t)
for i, (s, t) in enumerate(zip(selector, tracks))
if s):
if cherenkov[i]:
color = [255, 0, 0]
elif scintillation[i]:
color = [0, 0, 255]
elif reemission[i]:
color = [0, 255, 0]
else:
color = [255, 255, 255]
steps = min(
len(track), self.photons_max_steps
) if self.photons_max_steps is not None else len(track)
self.render_photon_track(geometry,
track[:steps],
sz=self.photons_track_size,
color=color)
if nphotons > 0:
print('Rendered Photons', nphotons)
geometry = create_geometry_from_obj(geometry)
gpu_geometry = gpu.GPUGeometry(geometry)
self.gpu_geometries.append(gpu_geometry)
def sum_events(self):
print('Summing events in file...')
nchannels = self.geometry.num_channels()
sum_hit = np.zeros(shape=nchannels, dtype=np.float)
sum_t = np.zeros(shape=nchannels, dtype=np.float)
sum_q = np.zeros(shape=nchannels, dtype=np.float)
nevents = len(self.rr)
for i, ev in enumerate(self.rr):
sum_hit += ev.channels.hit
sum_t[ev.channels.hit] += ev.channels.t[ev.channels.hit]
sum_q[ev.channels.hit] += ev.channels.q[ev.channels.hit]
if i % (nevents / 100 + 1) == 0:
print('.', end=' ', file=sys.stderr)
self.sum_hit = sum_hit
self.sum_t = sum_t / sum_hit
self.sum_q = sum_q / sum_hit
print('Done.')
def color_hit_pmts(self):
from chroma.color import map_to_color
self.gpu_geometry.reset_colors()
if self.display_mode == 'geo' or self.ev is None or self.ev.channels is None:
return
if self.sum_mode:
hit = self.sum_hit
t = self.sum_t
q = self.sum_q
select = hit > 0
else:
hit = self.ev.channels.hit
t = self.ev.channels.t
q = self.ev.channels.q
select = hit.copy()
if np.count_nonzero(select) == 0:
return
# Important: Compute range only with HIT channels
if self.display_mode == 'charge':
channel_color = map_to_color(q,
range=(q[select].min(),
q[select].max()))
elif self.display_mode == 'time':
if self.sum_mode:
crange = (t[select].min(), t[select].max())
else:
crange = (t[select].min(), t[select].mean())
channel_color = map_to_color(t, range=crange)
elif self.display_mode == 'hit':
channel_color = map_to_color(hit, range=(hit.min(), hit.max()))
elif self.display_mode == 'dichroicon':
if len(select) % 2 != 0:
return
channel_color = np.zeros_like(hit, dtype=np.uint32)
channel_color[::2] |= (255 * hit[::2] / np.max(hit[::2])).astype(
np.uint32)
channel_color[::2] |= (255 * hit[1::2] / np.max(hit[1::2])).astype(
np.uint32) << 16
select[0::2] |= select[1::2]
select[1::2] = 0
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] = select
#all but hit PMTs transparent
solid_hit[:] = True
solid_color[:] = 0xFF000000
channel_color[np.logical_not(select)] = 0xFF000000
solid_color[self.geometry.channel_index_to_solid_id] = channel_color
self.gpu_geometry.color_solids(solid_hit, solid_color)
def update(self):
Camera.update(self)
def process_event(self, event):
if event.type == KEYDOWN:
if event.key == K_p:
self.photon_display_mode = next(self.photon_display_mode_iter)
print(self.photon_display_mode)
self.render_mc_info()
self.update()
return
if event.key == K_t:
self.track_display_mode = next(self.track_display_mode_iter)
print(self.track_display_mode)
self.render_mc_info()
self.update()
return
if event.key == K_RIGHT and not self.sum_mode:
try:
self.ev = next(self.rr)
except StopIteration:
pass
else:
self.color_hit_pmts()
self.render_mc_info()
self.update()
return
elif event.key == K_LEFT and not self.sum_mode:
try:
self.ev = self.rr.prev()
except StopIteration:
pass
else:
self.color_hit_pmts()
self.render_mc_info()
self.update()
return
elif event.key == K_PERIOD:
self.display_mode = next(self.display_mode_iter)
self.color_hit_pmts()
self.update()
return
elif event.key == K_s:
self.sum_mode = not self.sum_mode
if self.sum_mode and not hasattr(self, 'sum_hit'):
self.sum_events()
elif not self.sum_mode and not hasattr(self, 'ev'):
return
self.color_hit_pmts()
self.render_mc_info()
self.update()
return
Camera.process_event(self, event)
class EventViewer(Camera):
# Constants for display_mode
CHARGE = 0
TIME = 1
HIT = 2
def __init__(self, geometry, filename, **kwargs):
Camera.__init__(self, geometry, **kwargs)
# This is really slow, so we do it here in the constructor to
# avoid slowing down the import of this module
from chroma.io.root import RootReader
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
h = x*np.sqrt(3)/2
pyramid = make.linear_extrude([-x/2,0,x/2], [-h/2,h/2,-h/2], h,
[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(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)
gpu_geometry = gpu.GPUGeometry(geometry)
self.gpu_geometries = [self.gpu_geometry, gpu_geometry]
def sum_events(self):
print 'Summing events in file...'
nchannels = self.geometry.num_channels()
sum_hit = np.zeros(shape=nchannels, dtype=np.float)
sum_t = np.zeros(shape=nchannels, dtype=np.float)
sum_q = np.zeros(shape=nchannels, dtype=np.float)
nevents = len(self.rr)
for i, ev in enumerate(self.rr):
sum_hit += ev.channels.hit
sum_t[ev.channels.hit] += ev.channels.t[ev.channels.hit]
sum_q[ev.channels.hit] += ev.channels.q[ev.channels.hit]
if i % (nevents / 100 + 1) == 0:
print >>sys.stderr, '.',
self.sum_hit = sum_hit
self.sum_t = sum_t / sum_hit
self.sum_q = sum_q / sum_hit
print 'Done.'
def color_hit_pmts(self):
from chroma.color import map_to_color
self.gpu_geometry.reset_colors()
if self.ev.channels is None:
return
if self.sum_mode:
hit = self.sum_hit
t = self.sum_t
q = self.sum_q
select = hit > 0
else:
hit = self.ev.channels.hit
t = self.ev.channels.t
q = self.ev.channels.q
select = hit
# Important: Compute range only with HIT channels
if self.display_mode == EventViewer.CHARGE:
channel_color = map_to_color(q, range=(q[select].min(),q[select].max()))
print 'charge'
elif self.display_mode == EventViewer.TIME:
if self.sum_mode:
crange = (t[select].min(), t[select].max())
else:
crange = (t[select].min(), t[select].mean())
channel_color = map_to_color(t, range=crange)
print 'time'#, crange
elif self.display_mode == EventViewer.HIT:
channel_color = map_to_color(hit, range=(hit.min(), hit.max()))
print 'hit'#, hit.min(), hit.max()
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] = select
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:
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()
except StopIteration:
pass
else:
self.color_hit_pmts()
if self.ev.photons_beg is not None:
self.render_particle_track()
self.update()
return
elif event.key == K_PAGEDOWN and not self.sum_mode:
try:
self.ev = self.rr.prev()
except StopIteration:
pass
else:
self.color_hit_pmts()
if self.ev.photons_beg is not None:
self.render_particle_track()
self.update()
return
elif event.key == K_PERIOD:
self.display_mode = (self.display_mode + 1) % 3
self.color_hit_pmts()
if not self.sum_mode and self.ev.photons_beg is not None:
self.render_particle_track()
self.update()
return
elif event.key == K_s:
self.sum_mode = not self.sum_mode
if self.sum_mode and not hasattr(self, 'sum_hit'):
self.sum_events()
elif not self.sum_mode and not hasattr(self, 'ev'):
return
self.color_hit_pmts()
self.update()
return
Camera.process_event(self, event)