class ContourSurf(Pipeline):
"""
Plots a the contours of a surface using grid-spaced data for
elevation supplied as a 2D array.
**Function signatures**::
contour_surf(s, ...)
contour_surf(x, y, s, ...)
contour_surf(x, y, f, ...)
s is the elevation matrix, a 2D array. The contour lines plotted
are lines of equal s value.
x and y can be 1D or 2D arrays (such as returned by numpy.ogrid or
numpy.mgrid), but the points should be located on an orthogonal grid
(possibly non-uniform). In other words, all the points sharing a same
index in the s array need to have the same x or y value. For
arbitrary-shaped position arrays (non-orthogonal grids), see the mesh
function.
If only 1 array s is passed, the x and y arrays are assumed to be
made from the indices of arrays, and an uniformly-spaced data set is
created.
If 3 positional arguments are passed the last one must be an array s,
or a callable, f, that returns an array. x and y give the
coordinates of positions corresponding to the s values."""
_source_function = Callable(array2d_source)
_pipeline = [WarpScalarFactory, ContourSurfaceFactory]
class Quiver3D(Points3d):
"""
Plots glyphs (like arrows) indicating the direction of the vectors
at the positions supplied.
**Function signatures**::
quiver3d(u, v, w, ...)
quiver3d(x, y, z, u, v, w, ...)
quiver3d(x, y, z, f, ...)
u, v, w are numpy arrays giving the components of the vectors.
If only 3 arrays, u, v, and w are passed, they must be 3D arrays, and
the positions of the arrows are assumed to be the indices of the
corresponding points in the (u, v, w) arrays.
If 6 arrays, (x, y, z, u, v, w) are passed, the 3 first arrays give
the position of the arrows, and the 3 last the components. They
can be of any shape.
If 4 positional arguments, (x, y, z, f) are passed, the last one must be
a callable, f, that returns vectors components (u, v, w) given the
positions (x, y, z)."""
scalars = Array(help="""optional scalar data.""")
_source_function = Callable(vector_scatter)
_pipeline = [
VectorsFactory,
]
class Contour3d(Pipeline):
"""
Plots iso-surfaces for a 3D volume of data suplied as arguments.
**Function signatures**::
contour3d(scalars, ...)
contour3d(x, y, z, scalars, ...)
scalars is a 3D numpy arrays giving the data on a grid.
If 4 arrays, (x, y, z, scalars) are passed, the 3 first arrays give
the position of the arrows, and the last the scalar value. The x, y
and z arrays are then supposed to have been generated by
`numpy.mgrid`, in other words, they are 3D arrays, with positions
lying on a 3D orthogonal and regularly spaced grid with nearest
neighbor in space matching nearest neighbor in the array. The
function builds a scalar field assuming the points are regularly
spaced.
If 4 positional arguments, (x, y, z, f) are passed, the last one
can also be a callable, f, that returns vectors components (u, v, w)
given the positions (x, y, z)."""
_source_function = Callable(scalar_field)
_pipeline = [
IsoSurfaceFactory,
]
class MayaviViewer(IVTK):
""" A viewer window for mlab.
"""
_scene_factory = Callable(mayavi_scene_factory)
def _size_default(self):
return (400, 300)
class SceneEditor(BasicEditorFactory):
""" A TraitsUI editor factory for SceneModel instances.
"""
# The class of the editor object to be constructed.
klass = _SceneEditor
# The class or factory function for creating the actual scene object.
scene_class = Callable(DecoratedScene)
class NullEngine(Engine):
"""
This class represents a NullEngine which creates a DummyViewer with
a scene set to None. This allows us to write full mayavi scripts
without the need for a UI and this is perfect for testing, or to
use Mayavi (and VTK) as a numerical engine.
This engine does not allow for rendring.
"""
scene_factory = Callable(dummy_viewer_factory)
class CallbackCSVLoader(CSVLoaderController):
""" Simple handler for a TraitsUI view to call a given callback on
exit.
"""
# A callable called when the TraitsUI view is closed. The instance
# that the view was editing is passed to this callable.
callback = Callable()
def closed(self, info, is_ok):
print "CallbackHander"
if is_ok:
self.callback(self.model)
class Flow(Pipeline):
"""
Creates a trajectory of particles following the flow of a vector field.
**Function signatures**::
flow(u, v, w, ...)
flow(x, y, z, u, v, w, ...)
flow(x, y, z, f, ...)
u, v, w are numpy arrays giving the components of the vectors.
If only 3 arrays, u, v, and w are passed, they must be 3D arrays, and
the positions of the arrows are assumed to be the indices of the
corresponding points in the (u, v, w) arrays.
If 6 arrays, (x, y, z, u, v, w) are passed, the 3 first arrays give
the position of the arrows, and the 3 last the components. The x, y
and z arrays are then supposed to have been generated by
`numpy.mgrid`, in other words, they are 3D arrays, with positions
lying on a 3D orthogonal and regularly spaced grid with nearest
neighbor in space matching nearest neighbor in the array. The
function builds a vector field assuming the points are regularly
spaced.
If 4 positional arguments, (x, y, z, f) are passed, the last one must be
a callable, f, that returns vectors components (u, v, w) given the
positions (x, y, z)."""
scalars = Array(help="""optional scalar data.""")
_source_function = Callable(vector_field)
_pipeline = [
ExtractVectorNormFactory,
StreamlineFactory,
]
def __call_internal__(self, *args, **kwargs):
""" Override the call to be able to choose whether to apply an
ExtractVectorNorm filter.
"""
self.source = self._source_function(*args, **kwargs)
kwargs.pop('name', None)
self.store_kwargs(kwargs)
# Copy the pipeline so as not to modify it for the next call
self.pipeline = self._pipeline[:]
if tools._has_scalar_data(self.source):
self.pipeline.pop(0)
return self.build_pipeline()
class Points3d(Pipeline):
"""
Plots glyphs (like points) at the position of the supplied data.
**Function signatures**::
points3d(x, y, z...)
points3d(x, y, z, s, ...)
points3d(x, y, z, f, ...)
x, y and z are numpy arrays, or lists, all of the same shape, giving
the positions of the points.
If only 3 arrays x, y, z are given, all the points are drawn with the
same size and color.
In addition, you can pass a fourth array s of the same
shape as x, y, and z giving an associated scalar value for each
point, or a function f(x, y, z) returning the scalar value. This
scalar value can be used to modulate the color and the size of the
points."""
_source_function = Callable(scalar_scatter)
_pipeline = [
GlyphFactory,
]
scale_factor = Any('auto',
help='The scaling applied to the glyphs. '
'the size of the glyph is by default calculated '
'from the inter-glyph spacing. Specify a float to '
'give the maximum glyph size in drawing units')
def __call_internal__(self, *args, **kwargs):
""" Override the call to be able to scale automatically the glyphs.
"""
scale_factor = kwargs.get('scale_factor', 'auto')
if scale_factor == 'auto':
kwargs['scale_factor'] = 1
g = Pipeline.__call_internal__(self, *args, **kwargs)
if scale_factor == 'auto':
g.glyph.glyph.scale_factor = \
tools._typical_distance(g.mlab_source.dataset)
g.glyph.glyph.clamping = True
else:
g.glyph.glyph.clamping = False
return g
class ImShow(Pipeline):
"""
View a 2D array as an image.
**Function signatures**::
imshow(s, ...)
s is a 2 dimension array. The values of s are mapped to a color using
the colormap."""
_source_function = Callable(array2d_source)
_pipeline = [
ImageActorFactory,
]
class ActorEditor(BasicEditorFactory):
""" An editor factory for TVTK scenes.
"""
# The class of the editor object to be constructed.
klass = _ActorEditor
# The class or factory function for creating the actual scene object.
scene_class = Callable(DecoratedScene)
# Keyword arguments to pass to the scene factory.
scene_kwds = Dict()
# The name of the trait used for ITVTKActorModel.disable_render.
disable_render_name = Str('disable_render')
# The name of the trait used for ITVTKActorModel.do_render.
do_render_name = Str('do_render')
class IVTK(ApplicationWindow):
""" Provides an Scene along without an embedded Python shell.
This is useful when scripting from the vanilla Python or IPython
interpreter."""
# The `Scene` instance into which VTK renders.
scene = Instance(Scene)
# The callable (or class) to create the scene instance
_scene_factory = Callable(DecoratedScene)
###########################################################################
# 'object' interface.
###########################################################################
def __init__(self, **traits):
""" Creates a new application window. """
# Base class constructor.
super(IVTK, self).__init__(**traits)
self.title = 'TVTK Scene'
self.menu_bar_manager = create_ivtk_menu(self)
###########################################################################
# `IWindow` interface.
###########################################################################
def close(self):
if self.scene is not None:
self.scene.close()
super(IVTK, self).close()
###########################################################################
# Protected 'ApplicationWindow' interface.
###########################################################################
# The icon of the window
icon = Instance(ImageResource, scene_icon)
def _create_contents(self, parent):
""" Create the contents of the window. """
self.scene = self._scene_factory(parent)
return self.scene.control
class TriangularMesh(Mesh):
"""
Plots a surface using a mesh defined by the position of its vertices
and the triangles connecting them.
**Function signatures**::
triangular_mesh(x, y, z, triangles ...)
x, y, z are arrays giving the positions of the vertices of the surface.
triangles is a list of triplets (or an array) list the vertices in
each triangle. Vertices are indexes by their appearance number in the
position arrays.
For simple structures (such as rectangular grids) prefer the surf or
mesh functions, as they will create more efficient data structures.
"""
_source_function = Callable(triangular_mesh_source)
class Plot3d(Pipeline):
"""
Draws lines between points.
**Function signatures**::
plot3d(x, y, z, ...)
plot3d(x, y, z, s, ...)
x, y, z and s are numpy arrays or lists of the same shape. x, y and z
give the positions of the successive points of the line. s is an
optional scalar value associated with each point."""
tube_radius = Trait(0.025,
CFloat,
None,
adapts='filter.radius',
help="""radius of the tubes used to represent the
lines, If None, simple lines are used.
""")
_source_function = Callable(line_source)
_pipeline = [
StripperFactory,
TubeFactory,
SurfaceFactory,
]
def __call_internal__(self, *args, **kwargs):
""" Override the call to be able to choose whether to apply
filters.
"""
self.source = self._source_function(*args, **kwargs)
kwargs.pop('name', None)
self.store_kwargs(kwargs)
# Copy the pipeline so as not to modify it for the next call
self.pipeline = self._pipeline[:]
if self.kwargs['tube_radius'] == None:
self.pipeline.remove(TubeFactory)
self.pipeline.remove(StripperFactory)
return self.build_pipeline()
class CustomBarChart(Pipeline):
"""
2012.2.21
custom version of BarChart. It has two more keyword arguments, x_scale, y_scale,
which is in charge of the lateral_scale in the X and Y direction.
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')
def __call__(self, *args, **kwargs):
""" Override the call to be able to scale automaticaly the axis.
"""
g = Pipeline.__call__(self, *args, **kwargs)
gs = g.glyph.glyph_source
# Use a cube source for glyphs.
if not 'mode' in kwargs:
gs.glyph_source = gs.glyph_list[-1]
# 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
x, y, z = g.mlab_source.x, g.mlab_source.y, g.mlab_source.z
scale_factor = g.glyph.glyph.scale_factor* \
tools._min_axis_distance(x, y, z)
x_scale = kwargs.pop('x_scale', self.lateral_scale)
y_scale = kwargs.pop('y_scale', self.lateral_scale)
try:
g.glyph.glyph_source.glyph_source.y_length = \
y_scale/(scale_factor)
g.glyph.glyph_source.glyph_source.x_length = \
x_scale/(scale_factor)
except TraitError:
" Not all types of glyphs have controlable y_length and x_length"
return g
def get_all_traits(self):
"""
2012.2.21
this function determines the set of possible keys in kwargs.
add two keywords, x_scale, y_scale.
Returns all the traits of class, and the classes in the pipeline.
"""
#call the parental one first.
traits = Pipeline.get_all_traits(self)
traits['x_scale'] = Trait(0.9, desc='The scale of the glyph on the x-axis, '
'in units of the distance between nearest points')
traits['y_scale'] = Trait(0.9, desc='The scale of the glyph on the y-axis, '
'in units of the distance between nearest points')
return traits
class Mesh(Pipeline):
"""
Plots a surface using grid-spaced data supplied as 2D arrays.
**Function signatures**::
mesh(x, y, z, ...)
x, y, z are 2D arrays, all of the same shape, giving the positions of
the vertices of the surface. The connectivity between these points is
implied by the connectivity on the arrays.
For simple structures (such as orthogonal grids) prefer the `surf` function,
as it will create more efficient data structures. For mesh defined by
triangles rather than regular implicit connectivity, see the
`triangular_mesh` function.
"""
scale_mode = Trait('none', {
'none': 'data_scaling_off',
'scalar': 'scale_by_scalar',
'vector': 'scale_by_vector'
},
help="""the scaling mode for the glyphs
('vector', 'scalar', or 'none').""")
scale_factor = CFloat(0.05,
desc="""scale factor of the glyphs used to represent
the vertices, in fancy_mesh mode. """)
tube_radius = Trait(0.025,
CFloat,
None,
help="""radius of the tubes used to represent the
lines, in mesh mode. If None, simple lines are used.
""")
scalars = Array(help="""optional scalar data.""")
mask = Array(help="boolean mask array to suppress some data points.")
representation = Trait(
'surface',
'wireframe',
'points',
'mesh',
'fancymesh',
desc="""the representation type used for the surface.""")
_source_function = Callable(grid_source)
_pipeline = [
ExtractEdgesFactory, GlyphFactory, TubeFactory, SurfaceFactory
]
def __call_internal__(self, *args, **kwargs):
""" Override the call to be able to choose whether to apply
filters.
"""
self.source = self._source_function(*args, **kwargs)
kwargs.pop('name', None)
self.store_kwargs(kwargs)
# Copy the pipeline so as not to modify it for the next call
self.pipeline = self._pipeline[:]
if not self.kwargs['representation'] in ('mesh', 'fancymesh'):
self.pipeline.remove(ExtractEdgesFactory)
self.pipeline.remove(TubeFactory)
self.pipeline.remove(GlyphFactory)
self.pipeline = [
PolyDataNormalsFactory,
] + self.pipeline
else:
if self.kwargs['tube_radius'] == None:
self.pipeline.remove(TubeFactory)
if not self.kwargs['representation'] == 'fancymesh':
self.pipeline.remove(GlyphFactory)
self.kwargs['representation'] = 'surface'
return self.build_pipeline()
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 Engine(HasStrictTraits):
""" The Mayavi engine base class.
"""
# The version of this class. Used for persistence.
__version__ = 0
# The scenes associated with this project.
scenes = List(Scene, record=True)
# The list to provide to a TreeEditor. Always add on a AdderNode.
# TODO: It makes more sense to put the modification of the list
# in some other UI module, and not here.
children_ui_list = Property(record=False)
# Our name.
name = Str('Mayavi Engine')
# Current scene.
current_scene = Property(Instance(Scene), record=False)
# Current object.
current_object = Property(record=False)
# Current selection -- the currently selected object on the tree.
current_selection = Property(record=False)
# Has the Engine started? Use this event to do something after
# the engine has been started.
started = Event(record=False)
# An optional callable that will generate a usable new viewer
# containing a `enthought.tvtk.pyface.TVTKScene` instance. Ideally
# the viewer should have an interface like
# `enthought.tvtk.pyface.TVTKWindow` -- basically it must
# implement the `closing` and `activated` events, however, this is
# not necessary. The created viewer is used by the `new_scene`
# method to create a new Viewer. This is a mechanism to use a
# user specified scene with the Engine and have the ability to
# load saved visualizations using the new scene. Handy for things
# like off-screen rendering.
scene_factory = Callable(viewer_factory)
# Are we running?
running = Bool(False, record=False)
# The recorder for script recording.
recorder = Instance(Recorder, record=False)
########################################
# Private traits.
_current_scene = Instance(Scene)
_current_object = Instance(HasTraits)
_current_selection = Instance(HasTraits)
_viewer_ref = Dict
# View related traits.
current_selection_view = View(Item(
name='_current_selection',
enabled_when='_current_selection is not None',
style='custom',
springy=True,
show_label=False,
),
resizable=True,
scrollable=True)
######################################################################
# `object` interface
######################################################################
def __init__(self, **traits):
super(Engine, self).__init__(**traits)
# FIXME: This is tied to preferences. It really should not be
# we need to use bind_preferences here.
cbk = lambda: self.trait_property_changed('children_ui_list', [], self.
children_ui_list)
preference_manager.root.on_trait_change(cbk, 'show_helper_nodes')
def __get_pure_state__(self):
d = self.__dict__.copy()
for x in [
'_current_scene', '_current_object', '__sync_trait__',
'_viewer_ref', '__traits_listener__'
]:
d.pop(x, None)
return d
def __set_pure_state__(self, state):
# Current number of scenes.
n_scene = len(self.scenes)
# Number of scenes in saved state.
n_saved_scene = len(state.scenes)
# Remove extra ones.
for i in range(n_scene - n_saved_scene):
self.close_scene(self.scenes[-1])
# Add new ones.
for i in range(n_saved_scene - n_scene):
self.new_scene()
# Set the state.
state_pickler.set_state(self, state)
def __getstate__(self):
return state_pickler.dumps(self)
def __setstate__(self, str_state):
self.__init__()
state = state_pickler.loads_state(str_state)
state_pickler.update_state(state)
self.__set_pure_state__(state)
######################################################################
# `Engine` interface
######################################################################
def start(self):
"""This is called by the plugin when the plugin actually
starts."""
registry.register_engine(self)
# Notify any listeners that the engine is started.
self.started = self
self.running = True
def stop(self):
registry.unregister_engine(self)
self.running = False
@recordable
def add_source(self, src, scene=None):
"""Adds a source to the pipeline. Uses the current scene unless a
scene is given in the scene keyword argument."""
passed_scene = scene
if scene is not None:
tvtk_scene = scene.scene
for sc in self.scenes:
if sc.scene == tvtk_scene:
scene = sc
break
else:
error('This scene is not managed by mayavi')
return
else:
scene = self.current_scene
# Create a new scene if none is available.
if scene is None:
self.new_scene()
scene = self.current_scene
scene.add_child(src)
self.current_object = src
@recordable
def add_filter(self, fil, obj=None):
"""Adds a filter to the pipeline at an appropriate point. Adds it
to the selected object, or to an object passed as the
kwarg `obj`.
"""
passed_obj = obj
if obj is None:
obj = self.current_object
if not isinstance(obj, Base):
msg = 'No valid current object, '\
'please select an active object.'
error(msg)
return
if (obj is not None) and (not isinstance(obj, Scene)):
if obj.running:
obj.add_child(fil)
self.current_object = fil
else:
msg = 'Current object is not active, '\
'please select an active object.'
error(msg)
else:
if obj is None:
error('Please create a VTK scene and open some data first.')
else:
error('No data: cannot use a Filter/Module/ModuleManager.')
@recordable
def add_module(self, mod, obj=None):
"""Adds a module to the pipeline at an appropriate point. Adds it
to the selected object, or to an object passed through the
kwarg `obj`.
"""
self.add_filter(mod, obj=obj)
@recordable
def save_visualization(self, file_or_fname):
"""Given a file or a file name, this saves the current
visualization to the file.
"""
# Save the state of VTK's global warning display.
o = vtk.vtkObject
w = o.GetGlobalWarningDisplay()
o.SetGlobalWarningDisplay(0) # Turn it off.
try:
state_pickler.dump(self, file_or_fname)
finally:
# Reset the warning state.
o.SetGlobalWarningDisplay(w)
@recordable
def load_visualization(self, file_or_fname):
"""Given a file/file name this loads the visualization."""
# Save the state of VTK's global warning display.
o = vtk.vtkObject
w = o.GetGlobalWarningDisplay()
o.SetGlobalWarningDisplay(0) # Turn it off.
try:
# Get the state from the file.
state = state_pickler.load_state(file_or_fname)
state_pickler.update_state(state)
# Add the new scenes.
for scene_state in state.scenes:
self.new_scene()
scene = self.scenes[-1]
# Disable rendering initially.
if scene.scene is not None:
scene.scene.disable_render = True
# Update the state.
state_pickler.update_state(scene_state)
scene.__set_pure_state__(scene_state)
# Setting the state will automatically reset the
# disable_render.
scene.render()
finally:
# Reset the warning state.
o.SetGlobalWarningDisplay(w)
@recordable
def open(self, filename, scene=None):
"""Open a file given a filename if possible in either the
current scene or the passed `scene`.
"""
passed_scene = scene
reader = registry.get_file_reader(filename)
if reader is None:
msg = 'No suitable reader found for the file %s' % filename
error(msg)
else:
src = None
if scene is None:
scene = self.current_scene
if scene is None:
scene = self.new_scene()
try:
sc = scene.scene
if sc is not None:
sc.busy = True
callable = reader.get_callable()
if reader.factory is None:
src = callable()
src.initialize(filename)
else:
# Factory functions are passed the filename and a
# reference to the engine.
src = callable(filename, self)
if src is not None:
self.add_source(src, passed_scene)
finally:
if sc is not None:
sc.busy = False
if src is not None:
return src
def record(self, msg):
"""This is merely a convenience method to record messages to the
script recorder.
"""
r = self.recorder
if r is not None:
r.record(msg)
######################################################################
# Scene creation/deletion related methods.
######################################################################
def add_scene(self, scene, name=None):
"""Add given `scene` (a `pyface.tvtk.scene.Scene` instance) to
the mayavi engine so that mayavi can manage the scene. This
is used when the user creates a scene. Note that for the
`EnvisageEngine` this is automatically taken care of when you
create a new scene using the TVTK scene plugin.
Parameters:
-----------
scene - `pyface.tvtk.scene.Scene`
The scene that needs to be managed from mayavi.
name - `str`
The name assigned to the scene. It tries to determine the
name of the scene from the passed scene instance. If this
is not possible it defaults to 'Mayavi Scene'.
"""
if name is None:
if hasattr(scene, 'name'):
name = scene.name
else:
name = 'Mayavi Scene %d' % scene_id_generator.next()
s = Scene(scene=scene, name=name, parent=self)
s.start()
# We don't want the startup setup to be recorded.
recorder = self.recorder
self.scenes.append(s)
self.current_scene = s
if recorder is not None:
recorder.register(s)
@recordable
def remove_scene(self, scene, **kwargs):
"""Remove a given `scene` (a `pyface.tvtk.scene.Scene`
instance) from the mayavi engine if it is already being
managed by mayavi. Note that for the `EnvisageEngine` this is
automatically taken care of when you close a scene started
using the TVTK scene plugin.
Parameters:
-----------
scene - `pyface.tvtk.scene.Scene`
The scene that needs to be removed from mayavi.
"""
s = None
for index, x in enumerate(self.scenes):
if x.scene is scene:
s = x
break
if s is not None:
s.stop()
self.scenes.remove(s)
# Don't record it shutting down. To do this we must
# unregister it here so we don't record unnecessary calls.
recorder = self.recorder
if recorder is not None:
recorder.unregister(s)
# Remove the reference to the viewer if any.
if scene in self._viewer_ref:
del self._viewer_ref[scene]
# Clear the current scene if it has been removed.
if scene is self._current_scene:
self._current_scene = None
@recordable
def new_scene(self, viewer=None, name=None, **kwargs):
"""Create or manage a new VTK scene window. If no `viewer`
argument is provided, the method creates a new viewer using
`self.scene_factory`. If `self.scene_factory` is `None` then
it creates an `ivtk` viewer. This code requires that the
`viewer` has a `scene` attribute/trait that is a
`pyface.tvtk.scene.Scene`. It also works best if the viewer
supports `closing` and `activated` events.
The method returns the created viewer.
Parameters:
-----------
viewer - The viewer object, if None, one is created for you.
name - The name attribute of the viewer
``**kwargs`` - The extra keyword arguments are passed along to
the scene factory.
"""
if viewer is None:
factory_kwargs = {}
factory_kwargs_names = get_args(self.scene_factory)
for arg, value in kwargs.iteritems():
if arg in factory_kwargs_names:
factory_kwargs[arg] = value
viewer = self.scene_factory(**factory_kwargs)
process_ui_events()
if name is not None:
viewer.name = name
# Hang on to a reference to this viewer, if not done this will cause a
# crash with Qt4. This because the viewer will be closed and gc'd if
# there isn't a reference to it. When the viewer is gc'd the scene is
# also closed and the engine will have a dead scene causing a crash.
self._viewer_ref[viewer.scene] = viewer
self.add_scene(viewer.scene)
if hasattr(viewer, 'on_trait_change'):
viewer.on_trait_change(self._on_scene_closed, 'closing')
viewer.on_trait_change(self._on_scene_activated, 'activated')
if hasattr(viewer, 'title'):
self.current_scene.sync_trait('name', viewer, 'title')
return viewer
@recordable
def close_scene(self, scene):
"""Given a scene created from new_scene, this method closes it
and removes the scene from the list of scenes we manage.
Parameters:
-----------
scene - `pyface.tvtk.scene.Scene` or an object that holds a
reference to a `pyface.tvtk.scene.Scene` in a `scene`
attribute.
"""
viewer = self.get_viewer(scene)
self.remove_scene(scene.scene)
if hasattr(scene, 'close'):
scene.close()
elif scene.scene is not None:
scene.scene.close()
if viewer is not None and hasattr(viewer, 'close'):
viewer.close()
def get_viewer(self, scene):
"""Return the viewer associated with a given scene.
Parameters:
-----------
scene - An `enthought.mayavi.core.scene.Scene` instance.
"""
return self._viewer_ref.get(scene.scene)
def dialog_view(self):
""" Default dialog view for Engine objects.
"""
return None
######################################################################
# Non-public interface
######################################################################
def _on_select(self, object):
"""Called by the EngineTree when an object on the view is
selected. This basically sets the current object and current
scene."""
self.current_selection = object
self._current_object = object
try:
scene = object.scene
for s in self.scenes:
if s.scene == scene:
self._current_scene = s
break
except AttributeError:
pass
def _get_current_scene(self):
n_scene = len(self.scenes)
if n_scene == 0:
return None
elif n_scene == 1:
return self.scenes[0]
elif self._current_scene is not None:
return self._current_scene
elif n_scene > 1:
return self.scenes[-1]
else:
return None
def _set_current_scene(self, scene):
old = self._current_scene
self._current_scene = scene
self.trait_property_changed('current_scene', old, scene)
def _get_current_object(self):
if self._current_object is not None:
return self._current_object
elif self.current_scene is not None:
return self.current_scene
else:
return None
def _set_current_object(self, object):
old = self._current_object
self._current_object = object
self.trait_property_changed('current_object', old, object)
def _get_current_selection(self):
return self._current_selection
def _set_current_selection(self, object):
old = self._current_selection
if not isinstance(object, (Base, AdderNode)):
object = None
self._current_selection = object
self.trait_property_changed('current_selection', old, object)
def _on_scene_closed(self, obj, name, old, new):
self.remove_scene(obj.scene)
def _on_scene_activated(self, obj, name, old, new):
for scene in self.scenes:
if scene.scene is obj.scene:
self.current_scene = scene
break
def _get_children_ui_list(self):
""" Trait getter for children_ui_list Property.
"""
if preference_manager.root.show_helper_nodes \
and len(self.scenes) == 0:
return [SceneAdderNode(object=self)]
else:
return self.scenes
@on_trait_change('scenes[]')
def _trigger_children_ui_list(self, old, new):
""" Trigger a children_ui_list change when scenes changed.
"""
self.trait_property_changed('children_ui_list', old, new)
def _recorder_changed(self, old, new):
if new is not None:
new.record('# Recorded script from Mayavi2')
new.record('from numpy import array')
new.record('try:')
new.record(' engine = mayavi.engine')
new.record('except NameError:')
new.record(' from enthought.mayavi.api import Engine')
new.record(' engine = Engine()')
new.record(' engine.start()')
new.record('if len(engine.scenes) == 0:')
new.record(' engine.new_scene()')
new.record('# ------------------------------------------- ')
elif old is not None:
old.record('# ------------------------------------------- ')
old.record('from enthought.mayavi.tools.show import show')
old.record('show()')
class OffScreenEngine(Engine):
# Overriding the scene factory trait of Engine.
scene_factory = Callable(off_screen_viewer_factory)
# Our name.
name = Str('Mayavi offscreen Engine')
id_trait = Str(desc='the name of the view')
# Contents of the view trait (i.e., a single Group object)
content_trait = Instance(Group, desc='the content of the view')
# The menu bar for the view
#menubar_trait = Instance( 'enthought.pyface.action.MenuBarManager',
# desc = 'the menu bar for the view' )
# The tool bar for the view
#toolbar_trait = Instance( 'enthought.pyface.action.ToolBarManager',
# desc = 'the tool bar for the view' )
# An optional model/view factory for converting the model into a viewable
# 'model_view' object
model_view_trait = Callable(desc='the factory function for converting a'
'model into a model/view object')
# Reference to a Handler object trait
handler_trait = Any(desc='the handler for the view')
# Dialog window title trait
title_trait = Str(desc='the window title for the view')
# Dialog window icon trait
#icon_trait = Instance( 'enthought.pyface.image_resource.ImageResource',
# desc = 'the ImageResource of the icon file for the view' )
# User interface 'kind' trait. The values have the following meanings:
#
# * 'panel': An embeddable panel. This type of window is intended to be used as
# part of a larger interface.
class TableFilter ( HasPrivateTraits ):
""" Filter for items displayed in a table.
"""
#---------------------------------------------------------------------------
# Trait definitions:
#---------------------------------------------------------------------------
# UI name of this filter (so the user can identify it in the UI)
name = Str( 'Default filter' )
# Default name that can be automatically overridden
_name = Str( 'Default filter' )
# A user-readable description of what kind of object satisfies the filter
desc = Str( 'All items' )
# A callable function that returns whether the passed object is allowed
# by the filter
allowed = Callable( lambda object: True )
# Is the filter a template (i.e., non-deletable, non-editable)?
template = false
#---------------------------------------------------------------------------
# Class constants:
#---------------------------------------------------------------------------
# Traits that are ignored by the _anytrait_changed() handler
ignored_traits = [ '_name', 'template', 'desc' ]
#---------------------------------------------------------------------------
# Traits view definitions:
#---------------------------------------------------------------------------
traits_view = View(
'name{Filter name}', '_',
Include( 'filter_view' ),
title = 'Edit Filter',
width = 0.2,
buttons = [ 'OK',
'Cancel',
Action(
name = 'Help',
action = 'show_help',
defined_when = "ui.view_elements.content['filter_view']"
".help_id != ''"
)
]
)
searchable_view = View( [
[ Include( 'search_view' ), '|[]' ],
[ 'handler.status~', '|[]<>' ],
[ 'handler.find_next`Find the next matching item`',
'handler.find_previous`Find the previous matching item`',
'handler.select`Select all matching items`',
'handler.OK`Exit search`', '-<>' ],
'|<>' ],
title = 'Search for',
kind = 'livemodal',
undo = False,
revert = False,
ok = False,
cancel = False,
help = False,
width = 0.25 )
search_view = Group( Include( 'filter_view' ) )
filter_view = Group()
#---------------------------------------------------------------------------
# Returns whether a specified object meets the filter/search criteria:
# (Should normally be overridden)
#---------------------------------------------------------------------------
def filter ( self, object ):
""" Returns whether a specified object meets the filter or search
criteria.
"""
return self.allowed( object )
#---------------------------------------------------------------------------
# Returns a user readable description of what kind of object will
# satisfy the filter:
# (Should normally be overridden):
#---------------------------------------------------------------------------
def description ( self ):
""" Returns a user-readable description of what kind of object
satisfies the filter.
"""
return self.desc
#---------------------------------------------------------------------------
# Edits the contents of the filter:
#---------------------------------------------------------------------------
def edit ( self, object ):
""" Edits the contents of the filter.
The ''object'' parameter is a sample object for the table that the
filter will be applied to. It is supplied in case the filter needs to
extract data or metadata from the object. If the table is empty, the
''object'' argument is None.
"""
return self.edit_traits( kind = 'livemodal' )
#---------------------------------------------------------------------------
# 'object' interface:
#---------------------------------------------------------------------------
def __str__ ( self ):
return self.name
#---------------------------------------------------------------------------
# Event handlers:
#---------------------------------------------------------------------------
def _anytrait_changed ( self, name, old, new ):
if ((name not in self.ignored_traits) and
((self.name == self._name) or (self.name == ''))):
self.name = self._name = self.description()
class Pipeline(HasTraits):
""" Function used to build pipelines for helper functions """
#doc = ''
_source_function = Callable()
_pipeline = List()
# Traits here only for documentation purposes
figure = Instance('enthought.mayavi.core.scene.Scene',
help='Figure to populate.')
def __call__(self, *args, **kwargs):
""" Calls the logics of the factory, but only after disabling
rendering, if needed.
"""
# First retrieve the scene, if any.
if 'figure' in kwargs:
figure = kwargs['figure']
assert isinstance(figure, (Scene, None))
scene = figure.scene
else:
scene = tools.gcf().scene
if scene is not None:
self._do_redraw = not scene.disable_render
scene.disable_render = True
# Then call the real logic
output = self.__call_internal__(*args, **kwargs)
# And re-enable the rendering, if needed.
if scene is not None:
scene.disable_render = not self._do_redraw
return output
def __call_internal__(self, *args, **kwargs):
""" Builds the source and runs through the pipeline, returning
the last object created by the pipeline."""
self.store_kwargs(kwargs)
self.source = self._source_function(*args, **kwargs)
# Copy the pipeline so as not to modify it for the next call
self.pipeline = self._pipeline[:]
return self.build_pipeline()
def store_kwargs(self, kwargs):
""" Merges the given keyword argument, with traits default and
store the resulting dictionary in self.kwargs."""
kwargs = kwargs.copy()
name = kwargs.pop('name', None)
all_traits = self.get_all_traits()
if not set(kwargs.keys()).issubset(all_traits.keys()):
raise ValueError, "Invalid keyword arguments : %s" % \
', '.join(str(k) for k in
set(kwargs.keys()).difference(all_traits.keys()) )
traits = self.get(self.class_trait_names())
[traits.pop(key) for key in traits.keys() if key[0] == '_']
traits.update(kwargs)
self.kwargs = traits
def build_pipeline(self):
""" Runs through the pipeline, applying pipe after pipe. """
object = self.source
for pipe in self.pipeline:
keywords = set(pipe.class_trait_names())
keywords.remove('trait_added')
keywords.remove('trait_modified')
this_kwargs = {}
for key, value in self.kwargs.iteritems():
if key in keywords:
this_kwargs[key] = value
object = pipe(object, **this_kwargs)._target
return object
def get_all_traits(self):
""" Returns all the traits of class, and the classes in the pipeline.
"""
traits = {}
for pipe in self._pipeline:
traits.update(pipe.class_traits())
traits.update(self.class_traits())
traits.pop('trait_added')
traits.pop('trait_modified')
return traits
class Surf(Pipeline):
"""
Plots a surface using regularly-spaced elevation data supplied as a 2D
array.
**Function signatures**::
surf(s, ...)
surf(x, y, s, ...)
surf(x, y, f, ...)
s is the elevation matrix, a 2D array, where indices along the first
array axis represent x locations, and indices along the second array
axis represent y locations.
x and y can be 1D or 2D arrays such as returned by numpy.ogrid or
numpy.mgrid. Arrays returned by numpy.meshgrid require a transpose
first to obtain correct indexing order.
The points should be located on an orthogonal grid (possibly
non-uniform). In other words, all the points sharing a same
index in the s array need to have the same x or y value. For
arbitrary-shaped position arrays (non-orthogonal grids), see the mesh
function.
If only 1 array s is passed, the x and y arrays are assumed to be
made from the indices of arrays, and an uniformly-spaced data set is
created.
If 3 positional arguments are passed the last one must be an array s,
or a callable, f, that returns an array. x and y give the
coordinates of positions corresponding to the s values."""
_source_function = Callable(array2d_source)
_pipeline = [WarpScalarFactory, PolyDataNormalsFactory, SurfaceFactory]
warp_scale = Any(1,
help="""scale of the z axis (warped from
the value of the scalar). By default this scale
is a float value.
If you specify 'auto', the scale is calculated to
give a pleasant aspect ratio to the plot,
whatever the bounds of the data.
If you specify a value for warp_scale in
addition to an extent, the warp scale will be
determined by the warp_scale, and the plot be
positioned along the z axis with the zero of the
data centered on the center of the extent. If you
are using explicit extents, this is the best way
to control the vertical scale of your plots.
If you want to control the extent (or range)
of the surface object, rather than its scale,
see the `extent` keyword argument.
""")
mask = Array(help="boolean mask array to suppress some data points.")
def __call_internal__(self, *args, **kwargs):
""" Override the call to be able to scale automatically the axis.
"""
self.source = self._source_function(*args, **kwargs)
kwargs.pop('name', None)
# Deal with both explicit warp scale and extent, this is
# slightly hairy. The wigner example is a good test case for
# this.
if not 'warp_scale' in kwargs and not 'extent' in kwargs:
try:
xi, xf, yi, yf, _, _ = self.source.data.bounds
zi, zf = self.source.data.scalar_range
except AttributeError:
xi, xf, yi, yf, _, _ = self.source.image_data.bounds
zi, zf = self.source.image_data.scalar_range
aspect_ratios = [(zf - zi) / (xf - xi), (zf - zi) / (yf - yi)]
if min(aspect_ratios) < 0.01 or max(aspect_ratios) > 100:
print 'Warning: the range of your scalar values differs by ' \
'more than a factor 100 than the range of the grid values ' \
'and you did not '\
'specify a warp_scale. You could try warp_scale="auto".'
if 'warp_scale' in kwargs and not kwargs['warp_scale']=='auto' \
and 'extent' in kwargs:
# XXX: I should use the logging module.
print 'Warning: both warp_scale and extent keyword argument ' \
'specified, the z bounds of the extents will be overridden'
xi, xf, yi, yf, zi, zf = kwargs['extent']
zo = 0.5 * (zi + zf)
try:
si, sf = self.source.data.scalar_range
except AttributeError:
si, sf = self.source.image_data.scalar_range
z_span = kwargs['warp_scale'] * abs(sf - si)
zi = zo + si * kwargs['warp_scale']
zf = zi + z_span
kwargs['extent'] = (xi, xf, yi, yf, zi, zf)
kwargs['warp_scale'] = 1
elif kwargs.get('warp_scale', 1) == 'auto':
if 'extent' in kwargs:
if 'warp_scale' in kwargs:
print "Warning: extent specified, warp_scale='auto' " \
"ignored."
else:
try:
xi, xf, yi, yf, _, _ = self.source.data.bounds
zi, zf = self.source.data.scalar_range
except AttributeError:
xi, xf, yi, yf, _, _ = self.source.image_data.bounds
zi, zf = self.source.image_data.scalar_range
z0 = zf - zi
dz = 0.3 * ((xf - xi) + (yf - yi))
zi = z0 - 0.5 * dz
zf = z0 + 0.5 * dz
kwargs['extent'] = (xi, xf, yi, yf, zi, zf)
kwargs['warp_scale'] = 1.
self.store_kwargs(kwargs)
# Copy the pipeline so as not to modify it for the next call
self.pipeline = self._pipeline[:]
return self.build_pipeline()
