class Axes(AxesLikeModuleFactory):
""" Creates axes for the current (or given) object."""
xlabel = String(None,
adapts='axes.x_label',
help='the label of the x axis')
ylabel = String(None,
adapts='axes.y_label',
help='the label of the y axis')
zlabel = String(None,
adapts='axes.z_label',
help='the label of the z axis')
nb_labels = Range(0,
50,
2,
adapts='axes.number_of_labels',
desc='The number of labels along each direction')
ranges = Trait(
None,
None,
CArray(shape=(6, )),
help="""[xmin, xmax, ymin, ymax, zmin, zmax]
Ranges of the labels displayed on the axes.
Default is the object's extents.""",
)
x_axis_visibility = true(
adapts='axes.x_axis_visibility',
help="Whether or not the x axis is visible (boolean)")
y_axis_visibility = true(
adapts='axes.y_axis_visibility',
help="Whether or not the y axis is visible (boolean)")
z_axis_visibility = true(
adapts='axes.z_axis_visibility',
help="Whether or not the z axis is visible (boolean)")
_target = Instance(modules.Axes, ())
def _extent_changed(self):
""" Code to modify the extents for
"""
axes = self._target
axes.axes.use_data_bounds = False
axes.axes.bounds = self.extent
if self.ranges is None:
axes.axes.ranges = \
axes.module_manager.source.outputs[0].bounds
def _ranges_changed(self):
if self.ranges is not None:
self._target.axes.ranges = self.ranges
self._target.axes.use_ranges = True
class Text3D(ModuleFactory):
""" Positions text at a 3D location in the scene.
**Function signature**::
text3d(x, y, z, text, ...)
x, y, and z are the position of the origin of the text. The
text is positionned in 3D, in figure coordinnates.
"""
_target = Instance(modules.Text3D, ())
scale = Either(CFloat(1),
CArray(shape=(3, )),
help="""The scale of the text, in figure units.
Either a float, or 3-tuple of floats.""")
orientation = CArray(shape=(3, ),
adapts='orientation',
desc="""the angles giving the orientation of the
text. If the text is oriented to the camera,
these angles are referenced to the axis of the
camera. If not, these angles are referenced to
the z axis.""")
orient_to_camera = true(adapts='orient_to_camera',
desc="""if the text is kept oriented to the
camera, or is pointing in a specific direction,
regardless of the camera position.""")
def __init__(self, x, y, z, text, **kwargs):
""" Override init as for different positional arguments."""
if not 'scale' in kwargs:
kwargs['scale'] = 1
super(Text3D, self).__init__(None, **kwargs)
self._target.text = text
self._target.position = (x, y, z)
def _scale_changed(self):
scale = self.scale
if operator.isNumberType(scale):
scale = scale * np.ones((3, ))
self._target.scale = scale
class DataModuleFactory(ModuleFactory):
""" Base class for all the module factories operating on data (ie not
text and outline) """
reset_zoom = true(help="""Reset the zoom to accomodate the data newly
added to the scene. Defaults to True.""")
extent = CArray(shape=(6,),
help="""[xmin, xmax, ymin, ymax, zmin, zmax]
Default is the x, y, z arrays extent. Use
this to change the extent of the object
created.""", )
def _extent_changed(self):
tools.set_extent(self._target, self.extent)
transparent = false(help="""make the opacity of the actor depend on the
scalar.""")
def _transparent_changed(self):
if self.transparent:
data_range = \
self._target.module_manager.scalar_lut_manager.data_range
self._target.module_manager.scalar_lut_manager.lut.alpha_range = \
(0.2, 0.8)
data_range = ( numpy.mean(data_range)
+ 0.4 * ( data_range.max() - data_range.min())
* numpy.array([-1, 1]))
self._target.module_manager.scalar_lut_manager.data_range = \
data_range
colormap = Trait('blue-red', lut_mode_list(),
help="""type of colormap to use.""")
def _colormap_changed(self):
colormap = self.colormap
if colormap[-2:] == "_r":
colormap = colormap[:-2]
self._target.module_manager.scalar_lut_manager.reverse_lut = True
self._target.module_manager.vector_lut_manager.reverse_lut = True
self._target.module_manager.scalar_lut_manager.lut_mode = colormap
self._target.module_manager.vector_lut_manager.lut_mode = colormap
vmin = Trait(None, None, CFloat,
help="""vmin is used to scale the colormap.
If None, the min of the data will be used""")
vmax = Trait(None, None, CFloat,
help="""vmax is used to scale the colormap.
If None, the max of the data will be used""")
def _vmin_changed(self):
if self.vmin == None and self.vmax == None:
self._target.module_manager.scalar_lut_manager.use_default_range\
= True
return
self._target.module_manager.scalar_lut_manager.use_default_range \
= False
vmin, vmax = \
self._target.module_manager.scalar_lut_manager.data_range
if self.vmin is not None:
vmin = self.vmin
if self.vmax is not None:
vmax = self.vmax
self._target.module_manager.scalar_lut_manager.data_range = \
(vmin, vmax)
_vmax_changed = _vmin_changed
def __init__(self, *args, **kwargs):
super(DataModuleFactory, self).__init__(*args, **kwargs)
# We are adding data to the scene, reset the zoom:
scene = self._scene.scene
if scene is not None and self.reset_zoom:
scene.reset_zoom()
class BarChart(Pipeline):
"""
Plots vertical glyphs (like bars) scaled vertical, to do
histogram-like plots.
This functions accepts a wide variety of inputs, with positions given
in 2-D or in 3-D.
**Function signatures**::
barchart(s, ...)
barchart(x, y, s, ...)
barchart(x, y, f, ...)
barchart(x, y, z, s, ...)
barchart(x, y, z, f, ...)
If only one positional argument is passed, it can be a 1-D, 2-D, or 3-D
array giving the length of the vectors. The positions of the data
points are deducted from the indices of array, and an
uniformly-spaced data set is created.
If 3 positional arguments (x, y, s) are passed the last one must be
an array s, or a callable, f, that returns an array. x and y give the
2D coordinates of positions corresponding to the s values.
If 4 positional arguments (x, y, z, s) are passed, the 3 first are
arrays giving the 3D coordinates of the data points, and the last one
is an array s, or a callable, f, that returns an array giving the
data value.
"""
_source_function = Callable(vertical_vectors_source)
_pipeline = [
VectorsFactory,
]
mode = Trait('cube',
bar_mode_dict,
desc='The glyph used to represent the bars.')
lateral_scale = CFloat(0.9,
desc='The lateral scale of the glyph, '
'in units of the distance between nearest points')
auto_scale = true(desc='whether to compute automatically the '
'lateral scaling of the glyphs. This might be '
'computationally expensive.')
def __call_internal__(self, *args, **kwargs):
""" Override the call to be able to scale automatically the axis.
"""
g = Pipeline.__call_internal__(self, *args, **kwargs)
gs = g.glyph.glyph_source
# Use a cube source for glyphs.
if not 'mode' in kwargs:
gs.glyph_source = gs.glyph_dict['cube_source']
# Position the glyph tail on the point.
gs.glyph_position = 'tail'
gs.glyph_source.center = (0.0, 0.0, 0.5)
g.glyph.glyph.orient = False
if not 'color' in kwargs:
g.glyph.color_mode = 'color_by_scalar'
if not 'scale_mode' in kwargs:
g.glyph.scale_mode = 'scale_by_vector_components'
g.glyph.glyph.clamping = False
# The auto-scaling code. It involves finding the minimum
# distance between points, which can be very expensive. We
# shortcut this calculation for structured data
if len(args) == 1 or self.auto_scale:
min_axis_distance = 1
else:
x, y, z = g.mlab_source.x, g.mlab_source.y, g.mlab_source.z
min_axis_distance = \
tools._min_axis_distance(x, y, z)
scale_factor = g.glyph.glyph.scale_factor * min_axis_distance
lateral_scale = kwargs.pop('lateral_scale', self.lateral_scale)
try:
g.glyph.glyph_source.glyph_source.y_length = \
lateral_scale / (scale_factor)
g.glyph.glyph_source.glyph_source.x_length = \
lateral_scale / (scale_factor)
except TraitError:
" Not all types of glyphs have controlable y_length and x_length"
return g
class Picker(HasTraits):
"""This module creates a 'Picker' that can interactively select a
point and/or a cell in the data. It also can use a world point
picker (i.e. a generic point in space) and will probe for the data
at that point.
The Picker is usually called via a callback from the GUI
interactor window. After performing a pick on the VTK scene, a
Picker object creates a `PickedData` object and passes it on to
the pick_handler trait for further handling.
"""
# The version of this class. Used for persistence.
__version__ = 0
# Speficifies the pick type. The 'point_picker' and 'cell_picker'
# options are self-explanatory. The 'world_picker' picks a point
# using a WorldPointPicker and additionally uses a ProbeFilter to
# probe the data at the picked point.
pick_type = Trait('point',
TraitRevPrefixMap({
'point_picker': 1,
'cell_picker': 2,
'world_picker': 3
}),
desc='specifies the picker type to use')
# The pick_handler. Set this to your own subclass if you want do
# do something different from the default.
pick_handler = Trait(DefaultPickHandler(), Instance(PickHandler))
# Picking tolerance.
tolerance = Range(0.0, 0.25, 0.025)
# show the GUI on pick ?
show_gui = true(desc="whether to show the picker GUI on pick")
# Raise the GUI on pick ?
auto_raise = true(desc="whether to raise the picker GUI on pick")
default_view = View(Group(
Group(Item(name='pick_type'), Item(name='tolerance'),
show_border=True),
Group(Item(name='pick_handler', style='custom'),
show_border=True,
show_labels=False),
Group(Item(name='show_gui'), Item(name='auto_raise'),
show_border=True),
),
resizable=True,
buttons=['OK'],
handler=CloseHandler())
#################################################################
# `object` interface.
#################################################################
def __init__(self, renwin, **traits):
super(Picker, self).__init__(**traits)
self.renwin = renwin
self.pointpicker = tvtk.PointPicker()
self.cellpicker = tvtk.CellPicker()
self.worldpicker = tvtk.WorldPointPicker()
self.probe_data = tvtk.PolyData()
self._tolerance_changed(self.tolerance)
# Use a set of axis to show the picked point.
self.p_source = tvtk.Axes()
self.p_mapper = tvtk.PolyDataMapper()
self.p_actor = tvtk.Actor()
self.p_source.symmetric = 1
self.p_actor.pickable = 0
self.p_actor.visibility = 0
prop = self.p_actor.property
prop.line_width = 2
prop.ambient = 1.0
prop.diffuse = 0.0
self.p_mapper.input = self.p_source.output
self.p_actor.mapper = self.p_mapper
self.probe_data.points = [[0.0, 0.0, 0.0]]
self.ui = None
def __get_pure_state__(self):
d = self.__dict__.copy()
for x in [
'renwin', 'ui', 'pick_handler', '__sync_trait__',
'__traits_listener__'
]:
d.pop(x, None)
return d
def __getstate__(self):
return state_pickler.dumps(self)
def __setstate__(self, str_state):
# This method is unnecessary since this object will almost
# never be pickled by itself and only via the scene, therefore
# __init__ will be called when the scene is constructed.
# However, setstate is defined just for completeness.
state_pickler.set_state(self, state_pickler.loads_state(str_state))
#################################################################
# `Picker` interface.
#################################################################
def pick(self, x, y):
"""Calls one of the current pickers and then passes the
obtained data to the `self.pick_handler` object's
`handle_pick` method.
Parameters
----------
- x : X position of the mouse in the window.
- y : Y position of the mouse in the window.
Note that the origin of x, y must be at the left bottom
corner of the window. Thus, for most GUI toolkits, y must
be flipped appropriately such that y=0 is the bottom of the
window.
"""
data = None
if self.pick_type_ == 1:
data = self.pick_point(x, y)
elif self.pick_type_ == 2:
data = self.pick_cell(x, y)
elif self.pick_type_ == 3:
data = self.pick_world(x, y)
self.pick_handler.handle_pick(data)
if self.show_gui:
self._setup_gui()
def pick_point(self, x, y):
""" Picks the nearest point. Returns a `PickedData` instance."""
self.pointpicker.pick((float(x), float(y), 0.0), self.renwin.renderer)
pp = self.pointpicker
id = pp.point_id
picked_data = PickedData()
coord = pp.pick_position
picked_data.coordinate = coord
if id > -1:
data = pp.mapper.input.point_data
bounds = pp.mapper.input.bounds
picked_data.valid = 1
picked_data.point_id = id
picked_data.data = data
self._update_actor(coord, bounds)
else:
self.p_actor.visibility = 0
self.renwin.render()
return picked_data
def pick_cell(self, x, y):
""" Picks the nearest cell. Returns a `PickedData` instance."""
try:
self.cellpicker.pick(float(x), float(y), 0.0, self.renwin.renderer)
except TypeError:
# On old versions of VTK, the signature used to be different
self.cellpicker.pick((float(x), float(y), 0.0),
self.renwin.renderer)
cp = self.cellpicker
id = cp.cell_id
picked_data = PickedData()
coord = cp.pick_position
picked_data.coordinate = coord
if id > -1:
data = cp.mapper.input.cell_data
bounds = cp.mapper.input.bounds
picked_data.valid = 1
picked_data.cell_id = id
picked_data.data = data
self._update_actor(coord, bounds)
else:
self.p_actor.visibility = 0
self.renwin.render()
return picked_data
def pick_world(self, x, y):
""" Picks a world point and probes for data there. Returns a
`PickedData` instance."""
self.worldpicker.pick((float(x), float(y), 0.0), self.renwin.renderer)
# Use the cell picker to get the data that needs to be probed.
try:
self.cellpicker.pick((float(x), float(y), 0.0),
self.renwin.renderer)
except TypeError:
self.cellpicker.pick(float(x), float(y), 0.0, self.renwin.renderer)
wp = self.worldpicker
cp = self.cellpicker
coord = wp.pick_position
self.probe_data.points = [list(coord)]
picked_data = PickedData()
picked_data.coordinate = coord
if cp.mapper:
data = get_last_input(cp.mapper.input)
# Need to create the probe each time because otherwise it
# does not seem to work properly.
probe = tvtk.ProbeFilter()
probe.source = data
probe.input = self.probe_data
probe.update()
data = probe.output.point_data
bounds = cp.mapper.input.bounds
picked_data.valid = 1
picked_data.world_pick = 1
picked_data.point_id = 0
picked_data.data = data
self._update_actor(coord, bounds)
else:
self.p_actor.visibility = 0
self.renwin.render()
return picked_data
def on_ui_close(self):
"""This method makes the picker actor invisible when the GUI
dialog is closed."""
self.p_actor.visibility = 0
self.renwin.renderer.remove_actor(self.p_actor)
self.ui = None
#################################################################
# Non-public interface.
#################################################################
def _tolerance_changed(self, val):
""" Trait handler for the tolerance trait."""
self.pointpicker.tolerance = val
self.cellpicker.tolerance = val
def _update_actor(self, coordinate, bounds):
"""Updates the actor by setting its position and scale."""
dx = 0.3 * (bounds[1] - bounds[0])
dy = 0.3 * (bounds[3] - bounds[2])
dz = 0.3 * (bounds[5] - bounds[4])
scale = max(dx, dy, dz)
self.p_source.origin = coordinate
self.p_source.scale_factor = scale
self.p_actor.visibility = 1
def _setup_gui(self):
"""Pops up the GUI control widget."""
# Popup the GUI control.
if self.ui is None:
self.ui = self.edit_traits()
# Note that we add actors to the renderer rather than to
# renwin to prevent event notifications on actor
# additions.
self.renwin.renderer.add_actor(self.p_actor)
elif self.auto_raise:
try:
self.ui.control.Raise()
except AttributeError:
pass
