def move_seeds(seeds, vfield, dt):
"""
Move a list of seeds based on a velocity field.
.. warning:: WARNING: THIS IS HARD CODED FOR GRID SIZE!
Parameters
----------
seeds: tvtk.PolyData
Old seed points
vfield: mayavi.sources.array_source.ArraySource object
The velocity field
dt: float
The time step betweent the current and the previous step.
Returns
-------
seeds_arr: ndarray
New Seed points
"""
v_seed = tvtk.ProbeFilter()
tvtk_common.configure_input_data(v_seed, seeds)
tvtk_common.configure_source_data(v_seed, vfield)
v_seed.update()
int_vels = np.array(v_seed.output.point_data.vectors)[:,:2]/(15.625*1e3)
seed_arr = np.array(seeds.points)
seed_arr[:,:2] += int_vels * dt
#seeds.points = seed_arr
return seed_arr
def interpolate_vectors(image_data, poly_data):
""" Interpolate a imagedata vectors to a set points in polydata"""
surface_probe_filter = tvtk.ProbeFilter()
tvtk_common.configure_source_data(surface_probe_filter, image_data)
tvtk_common.configure_input_data(surface_probe_filter, poly_data)
surface_probe_filter.update()
# Calculate Vperp, Vpar, Vphi
surface_vectors = np.array(surface_probe_filter.output.point_data.vectors)
return surface_probe_filter, surface_vectors
def get_the_line(bfield, surf_seeds, n):
"""Generate the vertical line on the surface"""
the_line = tvtk.StreamTracer()
source=tvtk.PolyData(points=np.array([surf_seeds.points.get_point(n),[0,0,0]]))
tvtk_common.configure_input_data(the_line, bfield)
tvtk_common.configure_source_data(the_line, source)
the_line.integrator = tvtk.RungeKutta4()
the_line.maximum_propagation = 1000
the_line.integration_direction = 'backward'
the_line.update()
return the_line
def create_flux_surface(bfield, surf_seeds):
"""
Create a flux surface from an array of seeds and a tvtk vector field.
Parameters
----------
bfield: tvtk.ImageData
The vector field to use for streamline traceing
surf_seeds: numpy.ndarray
The array of seed points to start the fieldline tracing from
Returns
-------
surf_field_lines: tvtk.StreamTracer instance
The fieldline tracer with the fieldlines stored inside it.
surface: tvtk.RuledSurfaceFilter instance
The surface built from the StreamTracer instance
"""
#Make a streamline instance with the bfield
surf_field_lines = tvtk.StreamTracer()
# surf_field_lines.input_connection = bfield
tvtk_common.configure_input(surf_field_lines, bfield)
tvtk_common.configure_source_data(surf_field_lines, surf_seeds)
# surf_field_lines.source = surf_seeds
surf_field_lines.integrator = tvtk.RungeKutta4()
surf_field_lines.maximum_propagation = 1000
surf_field_lines.integration_direction = 'backward'
surf_field_lines.update()
#Create surface from 'parallel' lines
surface = tvtk.RuledSurfaceFilter()
tvtk_common.configure_connection(surface, surf_field_lines)
# surface.input = surf_field_lines.output
surface.close_surface = True
surface.pass_lines = True
surface.offset = 0
surface.distance_factor = 30
surface.ruled_mode = 'point_walk'
# surface.ruled_mode = 'resample'
# surface.resolution = (10,1)
surface.update()
return surf_field_lines, surface
def make_tracer_pipeline(polydata,
seed,
maximum_number_of_steps=200,
maximum_propagation=1000):
"""Make a tracer pileine based on a polydata and seed.
The polydata is a vtk object with point_data for the velocities.
Seed is a polydata with point_data for the seed coordinates.
Propagation can be tuned using the other parameters (in meters).
"""
# create elements of the pipeline
tracer = tvtk.StreamTracer()
# maximum 1km
tracer.maximum_propagation = maximum_propagation
# # # Maximum 200 steps
tracer.maximum_number_of_steps = maximum_number_of_steps
# # # In m
tracer.integration_step_unit = vtk.vtkStreamTracer.LENGTH_UNIT
# # # Minimum 5 m per step
tracer.minimum_integration_step = (maximum_propagation /
maximum_number_of_steps)
# # # Maximum 50m per step
tracer.maximum_integration_step = 10 * tracer.minimum_integration_step
# # # Maximum error 1cm
tracer.maximum_error = 1e-2
# # # We use a path integration. You could argue that you need a
# # # particle tracking algorithm that matches the numerical grid
# # # (in our case edge velocities
# # # and integration over a cell instead of over a line)
tracer.integrator_type = 'runge_kutta45'
tracer.debug = True
cell2point = tvtk.CellDataToPointData()
# setup the pipeline
configure_input(cell2point, polydata)
configure_input(tracer, cell2point)
configure_source_data(tracer, seed)
return tracer
def configure_source_data(self, obj, data):
""" Configure the source data for vtk pipeline object obj."""
tvtk_common.configure_source_data(obj, data)
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
def configure_source_data(self, obj, data):
""" Configure the source data for vtk pipeline object obj."""
tvtk_common.configure_source_data(obj, data)
# Orientation/scaling is as per the vector attribute.
vecs = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
pd.point_data.vectors = vecs
# Create the glyph3d and set up the pipeline.
g = tvtk.Glyph3D(scale_mode='data_scaling_off', vector_mode='use_vector')
configure_input_data(g, pd)
# Note that VTK's vtkGlyph.SetSource is special because it has two
# call signatures: SetSource(src) and SetSource(int N, src) (which
# sets the N'th source). In tvtk it is represented as both a property
# and as a method. Using the `source` property will work fine if all
# you want is the first `source`. OTOH if you want the N'th `source`
# use get_source(N).
# g.source = cs.output
configure_source_data(g, cs.output)
cs.update()
g.update()
m = tvtk.PolyDataMapper()
configure_input_data(m, g.output)
a = tvtk.Actor(mapper=m)
# Read the texture from image and set the texture on the actor. If
# you don't like this image, replace with your favorite -- any image
# will do (you must use a suitable reader though).
def image_from_array(ary):
""" Create a VTK image object that references the data in ary.
The array is either 2D or 3D with. The last dimension
# Orientation/scaling is as per the vector attribute.
vecs = [[1,0,0], [0,1,0], [0,0,1]]
pd.point_data.vectors = vecs
# Create the glyph3d and set up the pipeline.
g = tvtk.Glyph3D(scale_mode='data_scaling_off', vector_mode = 'use_vector')
configure_input_data(g, pd)
# Note that VTK's vtkGlyph.SetSource is special because it has two
# call signatures: SetSource(src) and SetSource(int N, src) (which
# sets the N'th source). In tvtk it is represented as both a property
# and as a method. Using the `source` property will work fine if all
# you want is the first `source`. OTOH if you want the N'th `source`
# use get_source(N).
# g.source = cs.output
configure_source_data(g, cs.output)
cs.update()
g.update()
m = tvtk.PolyDataMapper()
configure_input_data(m, g.output)
a = tvtk.Actor(mapper=m)
# Read the texture from image and set the texture on the actor. If
# you don't like this image, replace with your favorite -- any image
# will do (you must use a suitable reader though).
def image_from_array(ary):
""" Create a VTK image object that references the data in ary.
The array is either 2D or 3D with. The last dimension
is always the number of channels. It is only tested
# The angular par of the spherical harmonic (3, 2) x, y, z = np.mgrid[-.5:.5:100j, -.5:.5:100j, -.5:.5:100j] Phi = np.angle((x+y*1j)**2*z) field = mlab.pipeline.scalar_field(x, y, z, Phi) ipw = mlab.pipeline.image_plane_widget(field) mlab.outline(field) surface = mlab.pipeline.builtin_surface() surface.source = 'sphere' surface.data_source.radius = .4 surface.data_source.phi_resolution = 200 surface.data_source.theta_resolution = 200 probe_filter = tvtk.ProbeFilter() configure_source_data(probe_filter, field.outputs[0]) probe = mlab.pipeline.user_defined(surface, filter=probe_filter) surf = mlab.pipeline.surface(probe) fig = mlab.gcf() ################################################################################ # Finally, to inspect the VTK Pipeline (and not the Mayavi one, we # use the TVTK pipeline browser) # Note that for Mayavi version < 3.4.1, there is a bug in the # PipelineBrowser preventing a good display of this pipeline. from tvtk.pipeline.browser import PipelineBrowser browser = PipelineBrowser(fig.scene) browser.show()
def main(*anim_files):
fig = mlab.figure(bgcolor=(1, 1, 1))
all_verts = []
pds = []
actors = []
datasets = []
glyph_pds = []
glyph_actors = []
colors = cycle([(1, 1, 1), (1, 0, 0), (0, 1, 0), (0, 0, 1)])
for i, (f, color) in enumerate(zip(anim_files, colors)):
data = h5py.File(f, 'r')
verts = data['verts'].value
tris = data['tris'].value
print f
print " Vertices: ", verts.shape
print " Triangles: ", tris.shape
datasets.append(data)
# setup mesh
pd = tvtk.PolyData(points=verts[0], polys=tris)
normals = tvtk.PolyDataNormals(compute_point_normals=True,
splitting=False)
configure_input_data(normals, pd)
actor = tvtk.Actor(mapper=tvtk.PolyDataMapper())
configure_input(actor.mapper, normals)
actor.mapper.immediate_mode_rendering = True
actor.visibility = False
fig.scene.add_actor(actor)
actors.append(actor)
all_verts.append(verts)
pds.append(normals)
# setup arrows
arrow = tvtk.ArrowSource(tip_length=0.25,
shaft_radius=0.03,
shaft_resolution=32,
tip_resolution=4)
glyph_pd = tvtk.PolyData()
glyph = tvtk.Glyph3D()
scale_factor = verts.reshape(-1, 3).ptp(0).max() * 0.1
glyph.set(scale_factor=scale_factor,
scale_mode='scale_by_vector',
color_mode='color_by_scalar')
configure_input_data(glyph, glyph_pd)
configure_source_data(glyph, arrow)
glyph_actor = tvtk.Actor(mapper=tvtk.PolyDataMapper(),
visibility=False)
configure_input(glyph_actor.mapper, glyph)
fig.scene.add_actor(glyph_actor)
glyph_actors.append(glyph_actor)
glyph_pds.append(glyph_pd)
actors[0].visibility = True
glyph_actors[0].visibility = True
class Viewer(HasTraits):
animator = Instance(Animator)
visible = Enum(*range(len(pds)))
normals = Bool(True)
export_off = Button
restpose = Bool(True)
show_scalars = Bool(True)
show_bones = Bool(True)
show_actual_bone_centers = Bool(False)
def _export_off_changed(self):
fd = FileDialog(title='Export OFF',
action='save as',
wildcard='OFF Meshes|*.off')
if fd.open() == OK:
v = all_verts[self.visible][self.animator.current_frame]
save_off(fd.path, v, tris)
@on_trait_change('visible, normals, restpose, show_scalars, show_bones'
)
def _changed(self):
for a in actors + glyph_actors:
a.visibility = False
for d in pds:
d.compute_point_normals = self.normals
actors[self.visible].visibility = True
actors[self.visible].mapper.scalar_visibility = self.show_scalars
glyph_actors[self.visible].visibility = self.show_bones
#for i, visible in enumerate(self.visibilities):
# actors[i].visibility = visible
self.animator.render = True
def show_frame(self, frame):
v = all_verts[self.visible][frame]
dataset = datasets[self.visible]
if not self.restpose:
rbms_frame = dataset['rbms'][frame]
v = v * dataset.attrs['scale'] + dataset.attrs['verts_mean']
v = blend_skinning(v,
dataset['segments'].value,
rbms_frame,
method=dataset.attrs['skinning_method'])
pds[self.visible].input.points = v
if 'scalar' in dataset and self.show_scalars:
if dataset['scalar'].shape[0] == all_verts[
self.visible].shape[0]:
scalar = dataset['scalar'][frame]
else:
scalar = dataset['scalar'].value
pds[self.visible].input.point_data.scalars = scalar
else:
pds[self.visible].input.point_data.scalars = None
if 'bone_transformations' in dataset and self.show_bones:
W = dataset['bone_blendweights'].value
T = dataset['bone_transformations'].value
gpd = glyph_pds[self.visible]
if self.show_actual_bone_centers:
verts0 = dataset['verts_restpose'].value
mean_bonepoint = verts0[W.argmax(axis=1)] # - T[0,:,:,3]
#mean_bonepoint = np.array([
# np.average(verts0, weights=w, axis=0) for w in W])
#gpd.points = np.repeat(mean_bonepoint + T[frame,:,:,3], 3, 0)
#print np.tile(mean_bonepoint + T[frame,:,:,3], 3).reshape((-1, 3))
#gpd.points = np.tile(mean_bonepoint + T[frame,:,:,3], 3).reshape((-1, 3))
pts = []
for i in xrange(T.shape[1]):
#offset = V.transform(V.hom4(mean_bonepoint[i]), T[frame,i,:,:])
offset = np.dot(T[frame, i], V.hom4(mean_bonepoint[i]))
pts += [offset] * 3
gpd.points = np.array(pts)
else:
bonepoint = np.array(
[np.average(v, weights=w, axis=0) for w in W])
gpd.points = np.repeat(bonepoint, 3, 0)
gpd.point_data.vectors = \
np.array(map(np.linalg.inv, T[frame,:,:,:3])).reshape(-1, 3)
# color vertices
vert_colors = vertex_weights_to_colors(W)
pds[self.visible].input.point_data.scalars = vert_colors
bone_colors = hue_linspace_colors(W.shape[0],
sat=0.8,
light=0.7)
gpd.point_data.scalars = np.repeat(bone_colors, 3, 0)
view = View(
Group(
Group(Item('visible'),
Item('export_off'),
Item('normals'),
Item('restpose'),
Item('show_scalars'),
Item('show_bones'),
Item('show_actual_bone_centers'),
label="Viewer"),
Item('animator', style='custom', show_label=False),
))
app = Viewer()
animator = Animator(verts.shape[0], app.show_frame)
app.animator = animator
app.edit_traits()
mlab.show()
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
