1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
|
"""
A helper function to retrieve the data from Mayavi structures on
arbitrary points.
"""
import numpy as np
from tvtk.api import tvtk
from . import tools
import tvtk.common as tvtk_common
def probe_data(mayavi_object, x, y, z, type='scalars', location='points'):
""" Retrieve the data from a described by Mayavi visualization object
at points x, y, z.
**Parameters**
:viz_obj: A Mayavi visualization object, or a VTK dataset
The object describing the data you are interested in.
:x: float or ndarray.
The x position where you want to retrieve the data.
:y: float or ndarray.
The y position where you want to retrieve the data.
:z: float or ndarray
The z position where you want to retrieve the data.
:type: 'scalars', 'vectors' or 'tensors', optional
The type of the data to retrieve.
:location: 'points' or 'cells', optional
The location of the data to retrieve.
**Returns**
The values of the data at the given point, as an ndarray
(or multiple arrays, in the case of vectors or tensors) of the
same shape as x, y, and z.
"""
dataset = tools.get_vtk_src(mayavi_object)[0]
assert type in ('scalars', 'vectors', 'cells'), (
"Invalid value for type: must be 'scalars', 'vectors' or "
"'cells', but '%s' was given" % type)
x = np.atleast_1d(x)
y = np.atleast_1d(y)
z = np.atleast_1d(z)
shape = x.shape
assert y.shape == z.shape == shape, \
'The x, y and z arguments must have the same shape'
probe_data = mesh = tvtk.PolyData(points=np.c_[x.ravel(),
y.ravel(),
z.ravel()])
shape = list(shape)
probe = tvtk.ProbeFilter()
tvtk_common.configure_input_data(probe, probe_data)
tvtk_common.configure_source_data(probe, dataset)
probe.update()
if location == 'points':
data = probe.output.point_data
elif location == 'cells':
data = probe.output.cell_data
else:
raise ValueError("Invalid value for data location, must be "
"'points' or 'cells', but '%s' was given."
% location)
values = getattr(data, type)
if values is None:
raise ValueError("The object given has no %s data of type %s"
% (location, type))
values = values.to_array()
if type == 'scalars':
values = np.reshape(values, shape)
elif type == 'vectors':
values = np.reshape(values, shape + [3, ])
values = np.rollaxis(values, -1)
else:
values = np.reshape(values, shape + [-1, ])
values = np.rollaxis(values, -1)
return values
|
