def _finalize_source(fld, arr, grid, dat_target):
if isinstance(fld, field.ScalarField):
dat_target.scalars = arr
dat_target.scalars.name = fld.name
elif isinstance(fld, field.VectorField):
dat_target.vectors = arr
dat_target.vectors.name = fld.name
src = VTKDataSource(data=grid)
src.name = fld.name
return src
def add_dataset(dataset, name='', **kwargs):
"""Add a dataset object to the Mayavi pipeline.
**Parameters**
:dataset: a tvtk dataset, or a Mayavi source.
The dataset added to the Mayavi pipeline
:figure: a figure identifier number or string, None or False, optionnal.
If no `figure` keyword argument is given, the data
is added to the current figure (a new figure if created if
necessary).
If a `figure` keyword argument is given, it should either the name
the number of the figure the dataset should be added to, or None,
in which case the data is not added to the pipeline.
If figure is False, a null engine is created. This null
engine does not create figures, and is mainly usefull for
tensting, or using the VTK algorithms without visualization.
**Returns**
The corresponding Mayavi source is returned.
"""
if isinstance(dataset, tvtk.Object):
d = VTKDataSource()
d.data = dataset
elif isinstance(dataset, Source):
d = dataset
else:
raise TypeError(
"first argument should be either a TVTK object"\
" or a mayavi source")
if len(name) > 0:
d.name = name
if not 'figure' in kwargs:
# No figure has been specified, retrieve the default one.
gcf()
engine = get_engine()
elif kwargs['figure'] is False:
# Get a null engine that we can use.
engine = get_null_engine()
elif kwargs['figure'] is not None:
figure = kwargs['figure']
engine = engine_manager.find_figure_engine(figure)
engine.current_scene = figure
else:
# Return early, as we don't want to add the source to an engine.
return d
engine.add_source(d)
return d
def add_dataset(dataset, name='', **kwargs):
"""Add a dataset object to the Mayavi pipeline.
**Parameters**
:dataset: a tvtk/vtk dataset/tvtk/VTK Algorithm, or a Mayavi source. The
dataset added to the Mayavi pipeline
:figure: a figure identifier number or string, None or False, optionnal.
If no `figure` keyword argument is given, the data
is added to the current figure (a new figure if created if
necessary).
If a `figure` keyword argument is given, it should either the name
the number of the figure the dataset should be added to, or None,
in which case the data is not added to the pipeline.
If figure is False, a null engine is created. This null
engine does not create figures, and is mainly usefull for
tensting, or using the VTK algorithms without visualization.
**Returns**
The corresponding Mayavi source is returned.
"""
if isinstance(dataset, (tvtk.DataSet, vtk.vtkDataSet)):
d = VTKDataSource()
d.data = tvtk.to_tvtk(dataset)
elif isinstance(dataset, (tvtk.DataObject, vtk.vtkDataObject)):
d = VTKObjectSource()
tp = tvtk.TrivialProducer()
tp.set_output(tvtk.to_tvtk(dataset))
d.object = tp
elif isinstance(dataset, (tvtk.Object, vtk.vtkObject)):
d = VTKObjectSource()
d.object = tvtk.to_tvtk(dataset)
elif isinstance(dataset, Source):
d = dataset
else:
raise TypeError(
"first argument should be either a TVTK object"
" or a mayavi source")
if len(name) > 0:
d.name = name
engine = _get_engine_from_kwarg(kwargs)
if engine is None:
# Return early, as we don't want to add the source to an engine.
return d
engine.add_source(d)
return d
def add_dataset(dataset, name='', **kwargs):
"""Add a dataset object to the Mayavi pipeline.
**Parameters**
:dataset: a tvtk/vtk dataset/tvtk/VTK Algorithm, or a Mayavi source. The
dataset added to the Mayavi pipeline
:figure: a figure identifier number or string, None or False, optional.
If no `figure` keyword argument is given, the data
is added to the current figure (a new figure if created if
necessary).
If a `figure` keyword argument is given, it should either the name
the number of the figure the dataset should be added to, or None,
in which case the data is not added to the pipeline.
If figure is False, a null engine is created. This null
engine does not create figures, and is mainly usefull for
tensting, or using the VTK algorithms without visualization.
**Returns**
The corresponding Mayavi source is returned.
"""
if isinstance(dataset, (tvtk.DataSet, vtk.vtkDataSet)):
d = VTKDataSource()
d.data = tvtk.to_tvtk(dataset)
elif isinstance(dataset, (tvtk.DataObject, vtk.vtkDataObject)):
d = VTKObjectSource()
tp = tvtk.TrivialProducer()
tp.set_output(tvtk.to_tvtk(dataset))
d.object = tp
elif isinstance(dataset, (tvtk.Object, vtk.vtkObject)):
d = VTKObjectSource()
d.object = tvtk.to_tvtk(dataset)
elif isinstance(dataset, Source):
d = dataset
else:
raise TypeError(
"first argument should be either a TVTK object"
" or a mayavi source")
if len(name) > 0:
d.name = name
engine = _get_engine_from_kwarg(kwargs)
if engine is None:
# Return early, as we don't want to add the source to an engine.
return d
engine.add_source(d)
return d
def fancy_axes(figure=None, target=None, nb_labels=5, xl=None, xh=None,
tight=False, symmetric=False, padding=0.05, opacity=0.7,
face_color=None, line_width=2.0, grid_color=None,
labels=True, label_color=None, label_shadow=True,
consolidate_labels=True):
"""Make axes with 3 shaded walls and a grid similar to matplotlib
Args:
figure (mayavi.core.scene.Scene): specific figure, or None for
:py:func:`mayavi.mlab.gcf`
target (Mayavi Element): If either xl or xh are not given, then
get that limit from a bounding box around `target`
nb_labels (int, sequence): number of labels in all, or each
(x, y, z) directions
xl (float, sequence): lower corner of axes
xh (float, sequence): upper corner of axes
tight (bool): If False, then let xl and xh expand to make nicer
labels. This uses matplotlib to determine new extrema
symmetric (bool): If True, then xl + xh = 0
padding (float): add padding as a fraction of the total length
opacity (float): opacity of faces
face_color (sequence): color (r, g, b) of faces
line_width (float): Width of grid lines
grid_color (sequence): Color of grid lines
labels (bool): Whether or not to put axis labels on
label_color (sequence): color of axis labels
label_shadow (bool): Add shadows to all labels
consolidate_labels (bool): if all nb_labels are the same, then
only make one axis for the labels
Returns:
VTKDataSource: source to which 2 surfaces and 3 axes belong
"""
if figure is None:
figure = mlab.gcf()
# setup xl and xh
if xl is None or xh is None:
_outline = mlab.outline(target, figure=figure)
if xl is None:
xl = _outline.bounds[0::2]
if xh is None:
xh = _outline.bounds[1::2]
_outline.remove()
nb_labels = np.broadcast_to(nb_labels, (3,))
xl = np.array(np.broadcast_to(xl, (3,)))
xh = np.array(np.broadcast_to(xh, (3,)))
L = xh - xl
xl -= padding * L
xh += padding * L
# now adjust xl and xh to be prettier
if symmetric:
tight = False
if not tight:
from matplotlib.ticker import AutoLocator
for i in range(len(xl)): # pylint: disable=consider-using-enumerate
l = AutoLocator()
l.create_dummy_axis()
l.set_view_interval(xl[i], xh[i])
locs = l()
xl[i] = locs[0]
xh[i] = locs[-1]
dx = (xh - xl) / (nb_labels - 1)
grid = tvtk.ImageData(dimensions=nb_labels, origin=xl, spacing=dx)
src = VTKDataSource(data=grid)
src.name = "fancy_axes"
if face_color is None:
face_color = figure.scene.background
if grid_color is None:
grid_color = figure.scene.foreground
if label_color is None:
label_color = grid_color
face = mlab.pipeline.surface(src, figure=figure, opacity=opacity,
color=face_color)
face.actor.property.frontface_culling = True
if line_width:
grid = mlab.pipeline.surface(src, figure=figure, opacity=1.0,
color=grid_color, line_width=line_width,
representation='wireframe')
grid.actor.property.frontface_culling = True
if labels:
def _make_ax_for_labels(_i, all_axes=False):
if all_axes:
_ax = Axes(name='axes-labels')
else:
_ax = Axes(name='{0}-axis-labels'.format('xyz'[_i]))
# VTK bug... y_axis and z_axis are flipped... how is VTK still
# the de-facto 3d plotting library?
if _i == 0:
_ax.axes.x_axis_visibility = True
_ax.axes.y_axis_visibility = False
_ax.axes.z_axis_visibility = False
elif _i == 1:
_ax.axes.x_axis_visibility = False
_ax.axes.y_axis_visibility = False
_ax.axes.z_axis_visibility = True # VTK bug
elif _i == 2:
_ax.axes.x_axis_visibility = False
_ax.axes.y_axis_visibility = True # VTK bug
_ax.axes.z_axis_visibility = False
else:
raise ValueError()
_ax.property.opacity = 0.0
_ax.axes.number_of_labels = nb_labels[_i]
# import IPython; IPython.embed()
_ax.title_text_property.color = label_color
_ax.title_text_property.shadow = label_shadow
_ax.label_text_property.color = label_color
_ax.label_text_property.shadow = label_shadow
src.add_module(_ax)
if consolidate_labels and np.all(nb_labels[:] == nb_labels[0]):
_make_ax_for_labels(0, all_axes=True)
else:
_make_ax_for_labels(0, all_axes=False)
_make_ax_for_labels(1, all_axes=False)
_make_ax_for_labels(2, all_axes=False)
return src
