class Line2D(HasTraits):
points = List(point_2d)
line_color = RGBColor("black")
updated = Event
def redraw(self):
pass # Not implemented for this example
def _points_changed(self):
self.updated = True
def _updated_fired(self):
self.redraw()
class RGBColorEditorDemo(HasTraits):
""" Defines the main RGBColorEditor demo. """
# Define a Color trait to view:
color_trait = RGBColor()
# Items are used to define the demo display, one item per editor style:
color_group = Group(
Item('color_trait', style='simple', label='Simple'), Item('_'),
Item('color_trait', style='custom', label='Custom'), Item('_'),
Item('color_trait', style='text', label='Text'), Item('_'),
Item('color_trait', style='readonly', label='ReadOnly'))
# Demo view
traits_view = View(color_group,
title='RGBColorEditor',
buttons=['OK'],
resizable=True)
class Object(HasPrivateTraits):
"""Represent a 3d object in a mayavi scene."""
points = Array(float, shape=(None, 3))
nn = Array(float, shape=(None, 3))
name = Str
scene = Instance(MlabSceneModel, ())
src = Instance(VTKDataSource)
# This should be Tuple, but it is broken on Anaconda as of 2016/12/16
color = RGBColor((1., 1., 1.))
# Due to a MESA bug, we use 0.99 opacity to force alpha blending
opacity = Range(low=0., high=1., value=0.99)
visible = Bool(True)
def _update_points(self):
"""Update the location of the plotted points."""
if hasattr(self.src, 'data'):
self.src.data.points = self.points
return True
class MyEntry(HasTraits):
name = Str()
value = Int(0)
color = RGBColor()
entry_view = View(Group(Item("name"), Item("value"), Item("color")))
class PointObject(Object):
"""Represent a group of individual points in a mayavi scene."""
label = Bool(False)
label_scale = Float(0.01)
projectable = Bool(False) # set based on type of points
orientable = Property(depends_on=['nearest'])
text3d = List
point_scale = Float(10, label='Point Scale')
# projection onto a surface
nearest = Instance(_DistanceQuery)
check_inside = Instance(_CheckInside)
project_to_trans = ArrayOrNone(float, shape=(4, 4))
project_to_surface = Bool(False,
label='Project',
desc='project points '
'onto the surface')
orient_to_surface = Bool(False,
label='Orient',
desc='orient points '
'toward the surface')
scale_by_distance = Bool(False,
label='Dist.',
desc='scale points by '
'distance from the surface')
mark_inside = Bool(False,
label='Mark',
desc='mark points inside the '
'surface in a different color')
inside_color = RGBColor((0., 0., 0.))
glyph = Instance(Glyph)
resolution = Int(8)
view = View(
HGroup(Item('visible', show_label=False),
Item('color', show_label=False), Item('opacity')))
def __init__(self, view='points', has_norm=False, *args, **kwargs):
"""Init.
Parameters
----------
view : 'points' | 'cloud'
Whether the view options should be tailored to individual points
or a point cloud.
has_norm : bool
Whether a norm can be defined; adds view options based on point
norms (default False).
"""
assert view in ('points', 'cloud', 'arrow')
self._view = view
self._has_norm = bool(has_norm)
super(PointObject, self).__init__(*args, **kwargs)
def default_traits_view(self): # noqa: D102
color = Item('color', show_label=False)
scale = Item('point_scale',
label='Size',
width=_SCALE_WIDTH,
editor=laggy_float_editor_headscale)
orient = Item('orient_to_surface',
enabled_when='orientable and not project_to_surface',
tooltip='Orient points toward the surface')
dist = Item('scale_by_distance',
enabled_when='orientable and not project_to_surface',
tooltip='Scale points by distance from the surface')
mark = Item('mark_inside',
enabled_when='orientable and not project_to_surface',
tooltip='Mark points inside the surface using a different '
'color')
if self._view == 'arrow':
visible = Item('visible', label='Show', show_label=False)
return View(HGroup(visible, scale, 'opacity', 'label', Spring()))
elif self._view == 'points':
visible = Item('visible', label='Show', show_label=True)
views = (visible, color, scale, 'label')
else:
assert self._view == 'cloud'
visible = Item('visible', show_label=False)
views = (visible, color, scale)
if not self._has_norm:
return View(HGroup(*views))
group2 = HGroup(dist,
Item('project_to_surface',
show_label=True,
enabled_when='projectable',
tooltip='Project points onto the surface '
'(for visualization, does not affect '
'fitting)'),
orient,
mark,
Spring(),
show_left=False)
return View(HGroup(HGroup(*views), group2))
@on_trait_change('label')
def _show_labels(self, show):
_toggle_mlab_render(self, False)
while self.text3d:
text = self.text3d.pop()
text.remove()
if show and len(self.src.data.points) > 0:
fig = self.scene.mayavi_scene
if self._view == 'arrow': # for axes
x, y, z = self.src.data.points[0]
self.text3d.append(
text3d(x,
y,
z,
self.name,
scale=self.label_scale,
color=self.color,
figure=fig))
else:
for i, (x, y, z) in enumerate(np.array(self.src.data.points)):
self.text3d.append(
text3d(x,
y,
z,
' %i' % i,
scale=self.label_scale,
color=self.color,
figure=fig))
_toggle_mlab_render(self, True)
@on_trait_change('visible')
def _on_hide(self):
if not self.visible:
self.label = False
@on_trait_change('scene.activated')
def _plot_points(self):
"""Add the points to the mayavi pipeline"""
if self.scene is None:
return
if hasattr(self.glyph, 'remove'):
self.glyph.remove()
if hasattr(self.src, 'remove'):
self.src.remove()
_toggle_mlab_render(self, False)
x, y, z = self.points.T
fig = self.scene.mayavi_scene
scatter = pipeline.scalar_scatter(x, y, z, fig=fig)
if not scatter.running:
# this can occur sometimes during testing w/ui.dispose()
return
# fig.scene.engine.current_object is scatter
mode = 'arrow' if self._view == 'arrow' else 'sphere'
glyph = pipeline.glyph(scatter,
color=self.color,
figure=fig,
scale_factor=self.point_scale,
opacity=1.,
resolution=self.resolution,
mode=mode)
glyph.actor.property.backface_culling = True
glyph.glyph.glyph.vector_mode = 'use_normal'
glyph.glyph.glyph.clamping = False
if mode == 'arrow':
glyph.glyph.glyph_source.glyph_position = 'tail'
glyph.actor.mapper.color_mode = 'map_scalars'
glyph.actor.mapper.scalar_mode = 'use_point_data'
glyph.actor.mapper.use_lookup_table_scalar_range = False
self.src = scatter
self.glyph = glyph
self.sync_trait('point_scale', self.glyph.glyph.glyph, 'scale_factor')
self.sync_trait('color', self.glyph.actor.property, mutual=False)
self.sync_trait('visible', self.glyph)
self.sync_trait('opacity', self.glyph.actor.property)
self.sync_trait('mark_inside', self.glyph.actor.mapper,
'scalar_visibility')
self.on_trait_change(self._update_points, 'points')
self._update_marker_scaling()
self._update_marker_type()
self._update_colors()
_toggle_mlab_render(self, True)
# self.scene.camera.parallel_scale = _scale
def _nearest_default(self):
return _DistanceQuery(np.zeros((1, 3)))
def _get_nearest(self, proj_rr):
idx = self.nearest.query(proj_rr)[1]
proj_pts = apply_trans(self.project_to_trans, self.nearest.data[idx])
proj_nn = apply_trans(self.project_to_trans,
self.check_inside.surf['nn'][idx],
move=False)
return proj_pts, proj_nn
@on_trait_change('points,project_to_trans,project_to_surface,mark_inside,'
'nearest')
def _update_projections(self):
"""Update the styles of the plotted points."""
if not hasattr(self.src, 'data'):
return
if self._view == 'arrow':
self.src.data.point_data.normals = self.nn
self.src.data.point_data.update()
return
# projections
if len(self.nearest.data) <= 1 or len(self.points) == 0:
return
# Do the projections
pts = self.points
inv_trans = np.linalg.inv(self.project_to_trans)
proj_rr = apply_trans(inv_trans, self.points)
proj_pts, proj_nn = self._get_nearest(proj_rr)
vec = pts - proj_pts # point to the surface
if self.project_to_surface:
pts = proj_pts
nn = proj_nn
if self.mark_inside and not self.project_to_surface:
scalars = (~self.check_inside(proj_rr, verbose=False)).astype(int)
else:
scalars = np.ones(len(pts))
# With this, a point exactly on the surface is of size point_scale
dist = np.linalg.norm(vec, axis=-1, keepdims=True)
self.src.data.point_data.normals = (250 * dist + 1) * nn
self.src.data.point_data.scalars = scalars
self.glyph.actor.mapper.scalar_range = [0., 1.]
self.src.data.points = pts # projection can change this
self.src.data.point_data.update()
@on_trait_change('color,inside_color')
def _update_colors(self):
if self.glyph is None:
return
# inside_color is the surface color, let's try to get far
# from that
inside = np.array(self.inside_color)
# if it's too close to gray, just use black:
if np.mean(np.abs(inside - 0.5)) < 0.2:
inside.fill(0.)
else:
inside = 1 - inside
colors = np.array([tuple(inside) + (1, ),
tuple(self.color) + (1, )]) * 255.
self.glyph.module_manager.scalar_lut_manager.lut.table = colors
@on_trait_change('project_to_surface,orient_to_surface')
def _update_marker_type(self):
# not implemented for arrow
if self.glyph is None or self._view == 'arrow':
return
defaults = DEFAULTS['coreg']
gs = self.glyph.glyph.glyph_source
res = getattr(gs.glyph_source, 'theta_resolution',
getattr(gs.glyph_source, 'resolution', None))
if self.project_to_surface or self.orient_to_surface:
gs.glyph_source = tvtk.CylinderSource()
gs.glyph_source.height = defaults['eegp_height']
gs.glyph_source.center = (0., -defaults['eegp_height'], 0)
gs.glyph_source.resolution = res
else:
gs.glyph_source = tvtk.SphereSource()
gs.glyph_source.phi_resolution = res
gs.glyph_source.theta_resolution = res
@on_trait_change('scale_by_distance,project_to_surface')
def _update_marker_scaling(self):
if self.glyph is None:
return
if self.scale_by_distance and not self.project_to_surface:
self.glyph.glyph.scale_mode = 'scale_by_vector'
else:
self.glyph.glyph.scale_mode = 'data_scaling_off'
def _resolution_changed(self, new):
if not self.glyph:
return
gs = self.glyph.glyph.glyph_source.glyph_source
if isinstance(gs, tvtk.SphereSource):
gs.phi_resolution = new
gs.theta_resolution = new
elif isinstance(gs, tvtk.CylinderSource):
gs.resolution = new
else: # ArrowSource
gs.tip_resolution = new
gs.shaft_resolution = new
@cached_property
def _get_orientable(self):
return len(self.nearest.data) > 1