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')