def write_to_vtk(conf, ar, filename):
"""
This function writes a numpy array into vtk format file.
It reads configuration into DispalyParams object.
The given array is cropped to the configured size if the crop parameter is defined in configuration file.
Based on configuration parameters the grid needed for the tvtk StructuredGrid is calculated.
Other StructuredGrid attributes are set.
The initialized StructuredGrid object is then written into vtk type file. If configuration parameter
save_two_files is set to True, one file will be saved for amplitudes with "_Amp.vtk" ending, and one for
phases with "_Phase.vtk" ending. Otherwise, a single file will be saved that contains double array
with amplitudes a nd phases.
Parameters
----------
conf : str
configuration file name
ar : numpy array
a complex array thatwill be saved
filename : str
a prefix to an output filename
Returns
-------
none
"""
params = DispalyParams(conf)
if params.crop is None:
dims = ar.shape
arr_cropped = ar.ravel()
else:
dims = params.crop
arr_cropped = crop_array_center(ar, params.crop).ravel()
amps = np.abs(arr_cropped)
phases = np.arctan2(arr_cropped.imag, arr_cropped.real)
geometry = params.get_geometry()
coordinates = get_coords(dims, geometry)
sg = tvtk.StructuredGrid()
sg.points = coordinates
sg.dimensions = (dims[2], dims[1], dims[0])
sg.extent = 0, dims[2] - 1, 0, dims[1] - 1, 0, dims[0] - 1
if params.save_two_files:
sg.point_data.scalars = amps
sg.point_data.scalars.name = "Amp"
write_array(sg, filename + "_Amp.vtk")
sg.point_data.scalars = phases
sg.point_data.scalars.name = "Phase"
write_array(sg, filename + "_Phase.vtk")
else:
sg.point_data.scalars = amps
sg.point_data.scalars.name = "Amp"
ph = tvtk.FloatArray()
ph.from_array(phases)
ph.name = "Phase"
sg.point_data.add_array(ph)
write_array(sg, filename + ".vtk")
def make_data():
points = N.array([[0,0,0], [1,0,0], [0,1,0], [0,0,1], # tets
[1,0,0], [2,0,0], [1,1,0], [1,0,1],
[2,0,0], [3,0,0], [2,1,0], [2,0,1],
], 'f')
tets = N.array([[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]])
tet_type = tvtk.Tetra().cell_type
ug = tvtk.UnstructuredGrid(points=points)
ug.set_cells(tet_type, tets)
# Setup the point attributes.
temp = N.random.random(12)
v = N.random.randn(12, 3)
ten = N.random.randn(12, 9)
a = tvtk.FloatArray(name='p')
a.from_array(N.random.randn(12))
ug.point_data.add_array(a)
ug.point_data.scalars = temp
ug.point_data.scalars.name = 't'
ug.point_data.vectors = v
ug.point_data.vectors.name = 'v'
ug.point_data.tensors = ten
ug.point_data.tensors.name = 'ten'
# Setup the cell attributes.
temp = N.random.random(3)
v = N.random.randn(3, 3)
ten = N.random.randn(3, 9)
ug.cell_data.scalars = temp
ug.cell_data.scalars.name = 't'
ug.cell_data.vectors = v
ug.cell_data.vectors.name = 'v'
ug.cell_data.tensors = ten
ug.cell_data.tensors.name = 'ten'
return ug
def update_data(path, polydata, mask, variables, t=0):
""" Add arrays for each variable to polydata """
ds = netCDF4.Dataset(path)
for key in variables:
arr = ds.variables[key][t, ...]
array = tvtk.FloatArray()
array.from_array(arr[~mask])
array.name = key
polydata.cell_data.add_array(array)
polydata.modified()
ds.close()
def test_array_conversion(self):
"""Test if Numeric/VTK array conversion works."""
# This is only a simple test.
data = numpy.array(
[[0, 0, 0, 10], [1, 0, 0, 20], [0, 1, 0, 20], [0, 0, 1, 30]], 'f')
triangles = numpy.array([[0, 1, 3], [0, 3, 2], [1, 2, 3], [0, 2, 1]])
points = data[:, :3]
temperature = data[:, -1]
mesh = tvtk.PolyData()
mesh.points = points
mesh.polys = triangles
mesh.point_data.scalars = temperature
# Test if a normal float array also works.
temp = tvtk.FloatArray()
temp.from_array(temperature)
mesh.point_data.scalars = temp
def test_data_array(self):
"""Test if vtkDataArrays behave in a Pythonic fashion."""
# Check a 3D array.
f = tvtk.FloatArray()
a = numpy.array([[0., 0, 0], [1, 1, 1]])
f.from_array(a)
self.assertEqual(f.number_of_components, 3)
self.assertEqual(f.number_of_tuples, 2)
self.assertEqual(mysum(f.to_array() - a), 0)
for i, j in zip(a, f):
self.assertEqual(mysum(i-j), 0.0)
self.assertEqual(f[0], (0.0, 0.0, 0.0))
self.assertEqual(f[-1], (1., 1., 1.))
self.assertEqual(repr(f), '[(0.0, 0.0, 0.0), (1.0, 1.0, 1.0)]')
f.append((2, 2, 2))
f.extend([[3, 3, 3], [4, 4, 4]])
self.assertEqual(len(f), 5)
self.assertEqual(f.number_of_components, 3)
self.assertEqual(f.number_of_tuples, 5)
f[1] = [-1, -1, -1]
self.assertEqual(f[1], (-1.0, -1.0, -1.0))
self.assertRaises(IndexError, f.__getitem__, 100)
self.assertRaises(IndexError, f.__setitem__, 100, 100)
# Check a 5D arrray
a = numpy.array([[0., 0, 0, 0, 0], [1, 1, 1, 1, 1]])
f.from_array(a)
self.assertEqual(mysum(f.to_array() - a), 0.0)
for i, j in zip(a, f):
self.assertEqual(mysum(i - j), 0.0)
self.assertEqual(f[0], (0.0, 0.0, 0.0, 0.0, 0.0))
self.assertEqual(f[-1], (1., 1., 1., 1., 1.))
self.assertEqual(
repr(f), '[(0.0, 0.0, 0.0, 0.0, 0.0), (1.0, 1.0, 1.0, 1.0, 1.0)]')
f.append((2, 2, 2, 2, 2))
f.extend([[3, 3, 3, 3, 3], [4, 4, 4, 4, 4]])
self.assertEqual(len(f), 5)
self.assertEqual(f.number_of_components, 5)
self.assertEqual(f.number_of_tuples, 5)
self.assertEqual(f[-1], (4., 4., 4., 4., 4.))
f[1] = [-1, -1, -1, -1, -1]
self.assertEqual(f[1], (-1.0, -1.0, -1.0, -1.0, -1.0))
self.assertRaises(IndexError, f.__getitem__, 100)
self.assertRaises(IndexError, f.__setitem__, 100, 100)
def screenshot(figure=None, mode='rgb', antialiased=False):
""" Return the current figure pixmap as an array.
**Parameters**
:figure: a figure instance or None, optional
If specified, the figure instance to capture the view of.
:mode: {'rgb', 'rgba'}
The color mode of the array captured.
:antialiased: {True, False}
Use anti-aliasing for rendering the screenshot.
Uses the number of aa frames set by
figure.scene.anti_aliasing_frames
**Notes**
On most systems, this works similarly to taking a screenshot of
the rendering window. Thus if it is hidden by another window, you
will capture the other window. This limitation is due to the
heavy use of the hardware graphics system.
**Examples**
This function can be useful for integrating 3D plotting with
Mayavi in a 2D plot created by matplotlib.
>>> from mayavi import mlab
>>> mlab.test_plot3d()
>>> arr = mlab.screenshot()
>>> import pylab as pl
>>> pl.imshow(arr)
>>> pl.axis('off')
>>> pl.show()
"""
if figure is None:
figure = gcf()
# don't use figure.scene.get_size() here because this will be incorrect
# for HiDPI systems
x, y = tuple(figure.scene.render_window.size)
# Try to lift the window
figure.scene._lift()
if mode == 'rgb':
out = tvtk.UnsignedCharArray()
shape = (y, x, 3)
pixel_getter = figure.scene.render_window.get_pixel_data
if vtk_major_version > 7:
pg_args = (0, 0, x - 1, y - 1, 1, out, 0)
else:
pg_args = (0, 0, x - 1, y - 1, 1, out)
elif mode == 'rgba':
out = tvtk.FloatArray()
shape = (y, x, 4)
pixel_getter = figure.scene.render_window.get_rgba_pixel_data
if vtk_major_version > 7:
pg_args = (0, 0, x - 1, y - 1, 1, out, 0)
else:
pg_args = (0, 0, x - 1, y - 1, 1, out)
else:
raise ValueError('mode type not understood')
if antialiased:
# save the current aa value to restore it later
render_window = figure.scene.render_window
if hasattr(render_window, 'aa_frames'):
old_aa = render_window.aa_frames
render_window.aa_frames = figure.scene.anti_aliasing_frames
else:
old_aa = render_window.multi_samples
render_window.multi_samples = figure.scene.anti_aliasing_frames
figure.scene.render()
pixel_getter(*pg_args)
if hasattr(render_window, 'aa_frames'):
render_window.aa_frames = old_aa
else:
render_window.multi_samples = old_aa
figure.scene.render()
else:
pixel_getter(*pg_args)
# Return the array in a way that pylab.imshow plots it right:
out = out.to_array()
out.shape = shape
out = np.flipud(out)
return out
def test_set_scalars(self):
"""Test if SetScalars works without a segfault."""
mesh = tvtk.PolyData()
sc = tvtk.FloatArray()
# If this does not segfault, we are OK.
mesh.point_data.scalars = sc
# self.start()
def add_actor(self, actor):
self.renderer.add_actor(actor)
# update window here.
### DATA
points = array([[0, 0, 0, 10], [1, 0, 0, 20], [0, 1, 0, 20], [0, 0, 1, 30]])
triangles = array([[0, 1, 3], [0, 3, 2], [1, 2, 3], [0, 2, 1]])
### PIPELINE
# Convert list of points to VTK structure
verts = tvtk.Points()
temperature = tvtk.FloatArray()
for p in points:
verts.insert_next_point(p[0], p[1], p[2])
temperature.insert_next_value(p[3])
# Define triangular cells from the vertex
# "ids" (index) and append to polygon list.
polygons = tvtk.CellArray()
for tri in triangles:
cell = tvtk.IdList()
cell.insert_next_id(tri[0])
cell.insert_next_id(tri[1])
cell.insert_next_id(tri[2])
polygons.insert_next_cell(cell)
# Create a mesh from these lists
