class DataSourceWizard(HasTraits):
data_sources = Dict
_data_sources_names = Property(depends_on='data_sources')
def _get__data_sources_names(self):
names = []
for name in self.data_sources:
try:
self.data_sources[name] + 1
names.append(name)
except TypeError:
pass
names.sort()
return names
# Dictionnary mapping the views
data_type = Trait(
'point', {
'A surface': 'surface',
'A set of points, that can be connected by lines': 'point',
'A set of vectors': 'vector',
'Volumetric data': 'volumetric',
})
position_type = Trait(
'image data', {
'Specified explicitly': 'explicit',
'Implicitely positioned on a regular grid': 'image data',
'On an orthogonal grid with varying spacing': 'orthogonal grid',
})
# The array that is used for finding out the shape of the grid,
# when creating an ImageData
grid_shape_source = Property(depends_on='grid_shape_source_')
def _get_grid_shape_source(self):
if self.grid_shape_source_ == '':
# catter for improperly initialized view
keys = self._data_sources_names
if not self.grid_shape.any():
self.grid_shape = \
self.data_sources[keys[0]].shape
return keys[0]
elif self.grid_shape_source_[:16] == 'Shape of array: ':
return self.grid_shape_source_[17:-1]
else:
return ""
# Shadow traits for grid_shape_source
grid_shape_source_ = Str
def _grid_shape_source_changed(self):
if not self.grid_shape_source == '':
array_shape = \
atleast_3d(self.data_sources[self.grid_shape_source]).shape
grid_shape = ones((3, ))
grid_shape[:len(array_shape)] = array_shape
self.grid_shape = grid_shape
_grid_shape_source_labels = Property(depends_on='_data_sources_names')
def _get__grid_shape_source_labels(self):
values = [
'Shape of array: "%s"' % name for name in self._data_sources_names
]
values.sort
values.append('Specified explicitly')
return values
# The shape of the grid array. Used when position is implicit
grid_shape = CArray(shape=(3, ), dtype='i')
# Whether or not the data points should be connected.
lines = false
# The scalar data selection
scalar_data = Str('',
help="Select the array that gives the value of the "
"scalars plotted.")
position_x = Str(help="Select the array that gives the x "
"position of the data points")
position_y = Str(help="Select the array that gives the y "
"position of the data points")
position_z = Str(help="Select the array that gives the z "
"position of the data points")
connectivity_triangles = Str
has_vector_data = false(help="""Do you want to plot vector components?""")
# A boolean to ask the user if he wants to load scalar data
has_scalar_data = false
vector_u = Str
vector_v = Str
vector_w = Str
#----------------------------------------------------------------------
# Public interface
#----------------------------------------------------------------------
def init_arrays(self):
# Force all the array names to be properly initialized
array_names = set(self.data_sources.keys())
if len(array_names) == 0:
# We should probably bail out here.
return False
for attr in (
'position_x',
'position_y',
'position_z',
'scalar_data',
'vector_u',
'vector_v',
'vector_w',
'connectivity_triangles',
):
if len(array_names) > 0:
array_name = array_names.pop()
setattr(self, attr, array_name)
def guess_arrays(self):
""" Do some guess work on the arrays to find sensible default.
"""
array_names = set(self._data_sources_names)
found_some = False
if set(('x', 'y', 'z')).issubset(array_names):
self.position_x = 'x'
self.position_y = 'y'
self.position_z = 'z'
array_names = array_names.difference(('x', 'y', 'z'))
found_some = True
elif set(('X', 'Y', 'Z')).issubset(array_names):
self.position_x = 'X'
self.position_y = 'Y'
self.position_z = 'Z'
array_names = array_names.difference(('X', 'Y', 'Z'))
found_some = True
if set(('u', 'v', 'w')).issubset(array_names):
self.vector_u = 'u'
self.vector_v = 'v'
self.vector_w = 'w'
array_names = array_names.difference(('u', 'v', 'w'))
found_some = True
elif set(('U', 'V', 'W')).issubset(array_names):
self.vector_u = 'U'
self.vector_v = 'V'
self.vector_w = 'W'
array_names = array_names.difference(('U', 'V', 'W'))
found_some = True
if found_some:
# Need to re-attribute the guessed names.
for attr in ('scalar_data', 'vector_u', 'vector_v', 'vector_w',
'connectivity_triangles'):
if len(array_names) > 0:
setattr(self, attr, array_names.pop())
else:
break
def build_data_source(self):
""" This is where we apply the selections made by the user in
in the wizard to build the data source.
"""
factory = DataSourceFactory()
# Keep a reference to the factory to be able to replay it, say
# on other data.
self._factory = factory
if self.data_type_ == 'point':
# The user wants to explicitly position vector,
# thus only sensible data structures for points is with
# explicit positioning.
self.position_type_ == 'explicit'
# In addition, this view does not allow for
# connectivity.
factory.unstructured = True
factory.connected = False
else:
factory.connected = True
if (self.position_type_ == "image data"
and not self.data_type_ == "point"):
if not self.has_scalar_data and not self.vector_u == '':
# With image data we need a scalar array always:
factory.scalar_data = ones(self.grid_shape)
factory.position_implicit = True
else:
factory.position_x = self.get_sdata(self.position_x)
factory.position_y = self.get_sdata(self.position_y)
factory.position_z = self.get_sdata(self.position_z)
if self.position_type_ == "orthogonal grid":
factory.orthogonal_grid = True
if self.position_type_ == "explicit" and self.data_type_ == "surface":
factory.connectivity_triangles = self.get_data(
self.connectivity_triangles)
if self.lines and self.data_type_ == "point":
factory.lines = True
if self.has_vector_data or self.data_type_ == 'vector':
# In the vector view, the user is not explicitly asked to
# Enable vectors.
factory.has_vector_data = True
factory.vector_u = self.get_sdata(self.vector_u)
factory.vector_v = self.get_sdata(self.vector_v)
factory.vector_w = self.get_sdata(self.vector_w)
if self.has_scalar_data or self.data_type_ == 'volumetric':
# In the volumetric view, the user is not explicitly asked to
# Enable scalars.
factory.scalar_data = self.get_sdata(self.scalar_data)
if self.connectivity_triangles == '':
factory.connectivity_triangles = None
self.data_source = factory.build_data_source()
if self.has_scalar_data:
if hasattr(self.data_source, 'scalar_name'):
self.data_source.scalar_name = self.scalar_data
elif hasattr(self.data_source, 'point_scalar_name'):
self.data_source.point_scalar_name = self.scalars
#----------------------------------------------------------------------
# Private interface
#----------------------------------------------------------------------
def get_data(self, name):
return self.data_sources[name]
def get_sdata(self, name):
ary = self.data_sources[name]
if not self.data_type_ == 'point':
ary = resize(ary, self.grid_shape)
return ary
def active_arrays(self):
""" Return the list of the active array-selection drop-downs.
"""
arrays = []
if self.data_type_ == 'point' or self.position_type_ == 'explicit':
arrays.extend([
'position_x',
'position_y',
'position_z',
])
if self.data_type_ == 'vector' or self.has_vector_data:
arrays.extend(['vector_u', 'vector_v', 'vector_w'])
if self.has_scalar_data or self.data_type_ == 'volumetric':
arrays.extend(['scalar_data'])
return arrays
def check_arrays(self):
""" Checks that all the array have the right size.
"""
arrays_to_check = self.active_arrays()
if len(arrays_to_check) == 0:
return True
size = self.get_data(getattr(self, arrays_to_check.pop())).size
for attr in arrays_to_check:
if not self.get_data(getattr(self, attr)).size == size:
return False
if (self.data_type_ == 'surface'
and self.position_type_ == "explicit"):
if not self.connectivity_triangles.size / 3 == size:
return False
return True
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 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 = RevPrefixMap({'point_picker': 1,
'cell_picker': 2,
'world_picker': 3},
default_value='point_picker',
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.1, 0.025)
# Raise the GUI on pick ?
auto_raise = false(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),
),
resizable=True,
buttons=['OK'])
#################################################################
# `object` interface.
#################################################################
def __init__(self, renwin, **traits):
super(Picker, self).__init__(**traits)
self.pointpicker = tvtk.PointPicker()
self.cellpicker = tvtk.CellPicker()
self.worldpicker = tvtk.WorldPointPicker()
self.probe_data = tvtk.PolyData()
# 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
configure_input(self.p_mapper, self.p_source)
self.p_actor.mapper = self.p_mapper
self.probe_data.points = [[0.0, 0.0, 0.0]]
self.text_rep = tvtk.TextRepresentation()
self.text_widget = tvtk.TextWidget()
self.data = PickedData(renwin=renwin)
self.data.text_actor = tvtk.TextActor()
self.text_setup()
self.widgets = False
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.
"""
if self.pick_type_ == 1:
self.data = self.pick_point(x, y)
elif self.pick_type_ == 2:
self.data = self.pick_cell(x, y)
elif self.pick_type_ == 3:
self.data = self.pick_world(x, y)
if self.widgets is False:
self.setup_widgets()
if self.data.valid_:
self.text_widget.enabled = 1
self.pick_handler.handle_pick(self.data)
self.data.text_actor._get_text_property().trait_set(
justification="left"
)
else:
self.close_picker()
text_color = self.data.text_actor._get_text_property().color
if not self.data.renwin.background == text_color:
pass
else:
self.set_text_color()
def pick_point(self, x, y):
""" Picks the nearest point. Returns a `PickedData` instance."""
self.pointpicker.pick((float(x), float(y), 0.0),
self.data.renwin.renderer)
pp = self.pointpicker
id = pp.point_id
picked_data = PickedData(renwin=self.data.renwin,
text_actor=self.data.text_actor)
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.data.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.data.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.data.renwin.renderer)
cp = self.cellpicker
id = cp.cell_id
picked_data = PickedData(renwin=self.data.renwin,
text_actor=self.data.text_actor)
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.data.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.data.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.data.renwin.renderer)
except TypeError:
self.cellpicker.pick(float(x), float(y), 0.0,
self.data.renwin.renderer)
wp = self.worldpicker
cp = self.cellpicker
coord = wp.pick_position
self.probe_data.points = [list(coord)]
picked_data = PickedData(renwin=self.data.renwin,
text_actor=self.data.text_actor)
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.set_source_data(data)
probe.set_input_data(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.data.renwin.render()
return picked_data
def close_picker(self):
"""This method makes the picker actor invisible when a non
data point is selected"""
if self.widgets:
self.p_actor.visibility = 0
self.data.renwin.renderer.remove_actor(self.p_actor)
self.data.text_actor.visibility = 0
self.data.renwin.renderer.remove_actor(self.data.text_actor)
self.text_widget.enabled = 0
self.widgets = False
self.data.renwin.render()
#################################################################
# Non-public interface.
#################################################################
def text_setup(self):
"""Sets the properties of the text widget"""
self.data.text_actor._get_text_property().font_size = 100
self.text_rep._get_position_coordinate().trait_set(value=(.15, .15, 0))
self.text_rep._get_position2_coordinate().trait_set(value=(.3, .2, 0))
self.text_widget.trait_set(representation=self.text_rep)
self.text_widget.trait_set(text_actor=self.data.text_actor)
self.text_widget.selectable = 0
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
self.data.text_actor.visibility = 1
def setup_widgets(self):
"""Sets up the picker actor and text actor"""
self.text_widget.trait_set(interactor=self.data.renwin.interactor)
self.data.renwin.renderer.add_actor(self.p_actor)
self.data.renwin.renderer.add_actor(self.data.text_actor)
self.p_actor.visibility = 1
self.data.text_actor.visibility = 1
self.widgets = True
def set_picker_props(self, pick_type, tolerance, text_color):
"""Lets you set the picker properties"""
self.pick_type = pick_type
self.tolerance = tolerance
self.set_text_color(text_color)
def set_text_color(self, color=None):
if color is None:
bgcolor = self.data.renwin.background
text_color = [0, 0, 0]
for i in range(3):
text_color[i] = abs(bgcolor[i]-1)
self.data.text_actor._get_text_property().color = tuple(text_color)
else:
self.data.text_actor._get_text_property().color = color