def write_vtk(data, filename, scalars=None, vectors={'V':('u','v','w')}, tensors={},
coords=('x','y','z'), dims=None, **kwargs):
''' write data in to vtk file
Parameters
----------
data : dict
mapping of variable name to their numpy array
filename : str
the file to write to (can be any recognized vtk extension)
if extension is missing .vts extension is appended
scalars : list
list of arrays to write as scalars (defaults to data.keys())
vectors : dict
mapping of vector name to vector component names to take from data
tensors : dict
mapping of tensor name to tensor component names to take from data
coords : list
the name of coordinate data arrays (default=('x','y','z'))
dims : 3 tuple
the size along the dimensions for (None means x.shape)
**kwargs : extra arguments for the file writer
example file_type=binary/ascii
'''
x = data[coords[0]]
y = data.get(coords[1], zeros_like(x))
z = data.get(coords[2], zeros_like(x))
if dims is None:
dims = array([1,1,1])
dims[:x.ndim] = x.shape
else:
dims = array(dims)
sg = tvtk.StructuredGrid(points=c_[x.flat,y.flat,z.flat],dimensions=array(dims))
pd = tvtk.PointData()
if scalars is None:
scalars = [i for i in data.keys() if i not in coords]
for v in scalars:
pd.scalars = ravel(data[v])
pd.scalars.name = v
sg.point_data.add_array(pd.scalars)
for vec,vec_vars in vectors.iteritems():
u,v,w = [data[i] for i in vec_vars]
pd.vectors = c_[ravel(u),ravel(v),ravel(w)]
pd.vectors.name = vec
sg.point_data.add_array(pd.vectors)
for ten,ten_vars in tensors.iteritems():
vars = [data[i] for i in ten_vars]
tensors = c_[[ravel(i) for i in vars]].T
pd.tensors = tensors
pd.tensors.name = ten
sg.point_data.add_array(pd.tensors)
write_data(sg, filename, **kwargs)
def save_output(self, fname):
"""Save our output (by default the first of our outputs) to the
specified filename as a VTK file. Both old style and new style
XML files are supported.
"""
if len(self.outputs) > 0:
write_data(self.outputs[0], fname)
else:
error('Object has no outputs to save!')
def test_write_data_xml(self):
"XML file writing with specified extension"
datasets = self.datasets
for d in datasets:
fh, fname = tempfile.mkstemp('.xml')
os.close(fh)
os.remove(fname)
self.assertEqual(os.path.exists(fname), False)
write_data(d, fname)
self.assertEqual(os.path.exists(fname), True)
os.remove(fname)
def test_write_data_xml_kwargs(self):
"XML file writing with extra keyword arguments"
datasets = self.datasets
exts = self.exts
for d, ext in zip(datasets, exts):
fh, fname = tempfile.mkstemp(ext)
fbase = os.path.splitext(fname)[0]
os.close(fh)
os.remove(fname)
# Test if passing extra keyword args is supported.
write_data(d, fbase, compressor=None, data_mode='ascii')
self.assertEqual(os.path.exists(fname), True)
os.remove(fname)
def test_write_data_vtk(self):
"Old-style VTK file writing with specified extension"
datasets = self.datasets
for d in datasets:
fh, fname = tempfile.mkstemp('.vtk')
os.close(fh)
os.remove(fname)
self.assertEqual(os.path.exists(fname), False)
write_data(d, fname)
self.assertEqual(os.path.exists(fname), True)
r = tvtk.DataSetReader(file_name=fname)
r.update()
self.assertEqual(isinstance(r.output, d.__class__), True)
os.remove(fname)
def test_write_data_xml_noext(self):
"XML file writing without extensions"
# Check if write_data writes out XML files with the correct
# extension when none is specified.
datasets = self.datasets
exts = self.exts
for d, ext in zip(datasets, exts):
fh, fname = tempfile.mkstemp(ext)
fbase = os.path.splitext(fname)[0]
os.close(fh)
os.remove(fname)
write_data(d, fbase)
self.assertEqual(os.path.exists(fname), True)
os.remove(fname)
def write_vtk(data,
filename,
scalars=None,
vectors={'V': ('u', 'v', 'w')},
tensors={},
coords=('x', 'y', 'z'),
dims=None,
**kwargs):
''' write data in to vtk file
Parameters
----------
data : dict
mapping of variable name to their numpy array
filename : str
the file to write to (can be any recognized vtk extension)
if extension is missing .vts extension is appended
scalars : list
list of arrays to write as scalars (defaults to data.keys())
vectors : dict
mapping of vector name to vector component names to take from data
tensors : dict
mapping of tensor name to tensor component names to take from data
coords : list
the name of coordinate data arrays (default=('x','y','z'))
dims : 3 tuple
the size along the dimensions for (None means x.shape)
**kwargs : extra arguments for the file writer
example file_type=binary/ascii
'''
x = data[coords[0]]
y = data.get(coords[1], zeros_like(x))
z = data.get(coords[2], zeros_like(x))
if dims is None:
dims = array([1, 1, 1])
dims[:x.ndim] = x.shape
else:
dims = array(dims)
sg = tvtk.StructuredGrid(points=c_[x.flat, y.flat, z.flat],
dimensions=array(dims))
pd = tvtk.PointData()
if scalars is None:
scalars = [i for i in data.keys() if i not in coords]
for v in scalars:
pd.scalars = ravel(data[v])
pd.scalars.name = v
sg.point_data.add_array(pd.scalars)
for vec, vec_vars in vectors.items():
u, v, w = [data[i] for i in vec_vars]
pd.vectors = c_[ravel(u), ravel(v), ravel(w)]
pd.vectors.name = vec
sg.point_data.add_array(pd.vectors)
for ten, ten_vars in tensors.items():
vars = [data[i] for i in ten_vars]
tensors = c_[[ravel(i) for i in vars]].T
pd.tensors = tensors
pd.tensors.name = ten
sg.point_data.add_array(pd.tensors)
write_data(sg, filename, **kwargs)
import numpy as np
from enthought.tvtk.api import tvtk, write_data
data = np.random.random((10, 10, 10))
grid = tvtk.ImageData(spacing=(10, 5, -10),
origin=(100, 350, 200),
dimensions=data.shape)
grid.point_data.scalars = np.ravel(order='F')
grid.point_data.scalars.name = 'Test Data'
# Writes legacy ".vtk" format if filename ends with "vtk", otherwise
# this will write data using the newer xml-based format.
write_data(grid, 'test.vtk')
