Exemplo n.º 1
0
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)
Exemplo n.º 2
0
 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!')
Exemplo n.º 3
0
    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)
Exemplo n.º 4
0
    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)
Exemplo n.º 5
0
    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)
Exemplo n.º 6
0
    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)
Exemplo n.º 7
0
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)
Exemplo n.º 8
0
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')