Source code for pyamg.vis.vtk_writer

"""VTK output functions.

Create coarse grid views and write meshes/primitives to .vtu files.  Use the
XML VTK format for unstructured meshes (.vtu)

This will use the XML VTK format for unstructured meshes, .vtu

See here for a guide:  http://www.vtk.org/pdf/file-formats.pdf
"""

__docformat__ = "restructuredtext en"

__all__ = ['write_vtu', 'write_basic_mesh']

import xml.dom.minidom
import numpy


[docs]def write_vtu(Verts, Cells, pdata=None, pvdata=None, cdata=None, cvdata=None, fname='output.vtk'): """ Write a .vtu file in xml format Parameters ---------- fname : {string} file to be written, e.g. 'mymesh.vtu' Verts : {array} Ndof x 3 (if 2, then expanded by 0) list of (x,y,z) point coordinates Cells : {dictionary} Dictionary of with the keys pdata : {array} Ndof x Nfields array of scalar values for the vertices pvdata : {array} Nfields*3 x Ndof array of vector values for the vertices cdata : {dictionary} scalar valued cell data cvdata : {dictionary} vector valued cell data Returns ------- writes a .vtu file for use in Paraview Notes ----- - Poly data not supported - Non-Poly data is stored in Numpy array: Ncell x vtk_cell_info - Each I1 must be >=3 - pdata = Ndof x Nfields - pvdata = 3*Ndof x Nfields - cdata,cvdata = list of dictionaries in the form of Cells ===== =================== ============= === keys type n points dim ===== =================== ============= === 1 VTK_VERTEX: 1 point 2d 2 VTK_POLY_VERTEX: n points 2d 3 VTK_LINE: 2 points 2d 4 VTK_POLY_LINE: n+1 points 2d 5 VTK_TRIANGLE: 3 points 2d 6 VTK_TRIANGLE_STRIP: n+2 points 2d 7 VTK_POLYGON: n points 2d 8 VTK_PIXEL: 4 points 2d 9 VTK_QUAD: 4 points 2d 10 VTK_TETRA: 4 points 3d 11 VTK_VOXEL: 8 points 3d 12 VTK_HEXAHEDRON: 8 points 3d 13 VTK_WEDGE: 6 points 3d 14 VTK_PYRAMID: 5 points 3d ===== =================== ============= === Examples -------- >>> from pyamg.vis import write_vtu >>> import numpy >>> Verts = numpy.array([[0.0,0.0], ... [1.0,0.0], ... [2.0,0.0], ... [0.0,1.0], ... [1.0,1.0], ... [2.0,1.0], ... [0.0,2.0], ... [1.0,2.0], ... [2.0,2.0], ... [0.0,3.0], ... [1.0,3.0], ... [2.0,3.0]]) >>> E2V = numpy.array([[0,4,3], ... [0,1,4], ... [1,5,4], ... [1,2,5], ... [3,7,6], ... [3,4,7], ... [4,8,7], ... [4,5,8], ... [6,10,9], ... [6,7,10], ... [7,11,10], ... [7,8,11]]) >>> E2edge = numpy.array([[0,1]]) >>> E2point = numpy.array([2,3,4,5]) >>> Cells = {5:E2V,3:E2edge,1:E2point} >>> pdata=numpy.ones((12,2)) >>> pvdata=numpy.ones((12*3,2)) >>> cdata={5:numpy.ones((12,2)),3:numpy.ones((1,2)),1:numpy.ones((4,2))} >>> cvdata={5:numpy.ones((3*12,2)),3:numpy.ones((3*1,2)), 1:numpy.ones((3*4,2))} >>> write_vtu(Verts=Verts, Cells=Cells, fname='test.vtu') See Also -------- write_mesh """ # number of indices per cell for each cell type vtk_cell_info = [-1, 1, None, 2, None, 3, None, None, 4, 4, 4, 8, 8, 6, 5] # check fname if type(fname) is str: try: fname = open(fname, 'w') except IOError, (errno, strerror): print ".vtu error (%s): %s" % (errno, strerror) else: raise ValueError('fname is assumed to be a string') # check Verts # get dimension and verify that it's 3d data Ndof, dim = Verts.shape if dim == 2: # always use 3d coordinates (x,y) -> (x,y,0) Verts = numpy.hstack((Verts, numpy.zeros((Ndof, 1)))) # check Cells # keys must ve valid (integer and not "None" in vtk_cell_info) # Cell data can't be empty for a non empty key for key in Cells: if ((type(key) != int) or (key not in range(1, 15))): raise ValueError('cell array must have positive integer keys\ in [1,14]') if (vtk_cell_info[key] is None) and (Cells[key] is not None): # Poly data raise NotImplementedError('Poly Data not implemented yet') if Cells[key] is None: raise ValueError('cell array cannot be empty for\ key %d' % (key)) if numpy.ndim(Cells[key]) != 2: Cells[key] = Cells[key].reshape((Cells[key].size, 1)) if vtk_cell_info[key] != Cells[key].shape[1]: raise ValueError('cell array has %d columns, expected %d' % (Cells[key].shape[1], vtk_cell_info[key])) # check pdata # must be Ndof x n_pdata n_pdata = 0 if pdata is not None: if numpy.ndim(pdata) > 1: n_pdata = pdata.shape[1] else: n_pdata = 1 pdata = pdata.reshape((pdata.size, 1)) if pdata.shape[0] != Ndof: raise ValueError('pdata array should be length %d (it is %d)' % (Ndof, pdata.shape[0])) # check pvdata # must be 3*Ndof x n_pvdata n_pvdata = 0 if pvdata is not None: if numpy.ndim(pvdata) > 1: n_pvdata = pvdata.shape[1] else: n_pvdata = 1 pvdata = pvdata.reshape((pvdata.size, 1)) if pvdata.shape[0] != 3*Ndof: raise ValueError('pvdata array should be of size %d (or multiples)\ (it is now %d)' % (Ndof*3, pvdata.shape[0])) # check cdata # must be NCells x n_cdata for each key n_cdata = 0 if cdata is not None: for key in Cells: # all valid now if numpy.ndim(cdata[key]) > 1: if n_cdata == 0: n_cdata = cdata[key].shape[1] elif n_cdata != cdata[key].shape[1]: raise ValueError('cdata dimension problem') else: n_cdata = 1 cdata[key] = cdata[key].reshape((cdata[key].size, 1)) if cdata[key].shape[0] != Cells[key].shape[0]: raise ValueError('size mismatch with cdata %d and Cells %d' % (cdata[key].shape[0], Cells[key].shape[0])) if cdata[key] is None: raise ValueError('cdata array cannot be empty for key %d' % (key)) # check cvdata # must be NCells*3 x n_cdata for each key n_cvdata = 0 if cvdata is not None: for key in Cells: # all valid now if numpy.ndim(cvdata[key]) > 1: if n_cvdata == 0: n_cvdata = cvdata[key].shape[1] elif n_cvdata != cvdata[key].shape[1]: raise ValueError('cvdata dimension problem') else: n_cvdata = 1 cvdata[key] = cvdata[key].reshape((cvdata[key].size, 1)) if cvdata[key].shape[0] != 3*Cells[key].shape[0]: raise ValueError('size mismatch with cvdata and Cells') if cvdata[key] is None: raise ValueError('cvdata array cannot be empty for key %d' % (key)) Ncells = 0 cell_ind = [] cell_offset = [] # numpy.zeros((Ncells,1),dtype=uint8) # zero indexed cell_type = [] # numpy.zeros((Ncells,1),dtype=uint8) cdata_all = None cvdata_all = None for key in Cells: # non-Poly data sz = Cells[key].shape[0] offset = Cells[key].shape[1] Ncells += sz uu = numpy.ones((sz,), dtype='uint8') cell_ind = numpy.hstack((cell_ind, Cells[key].ravel())) cell_offset = numpy.hstack((cell_offset, offset*uu)) cell_type = numpy.hstack((cell_type, key*uu)) if cdata is not None: if cdata_all is None: cdata_all = cdata[key] else: cdata_all = numpy.vstack((cdata_all, cdata[key])) if cvdata is not None: if cvdata_all is None: cvdata_all = cvdata[key] else: cvdata_all = numpy.vstack((cvdata_all, cvdata[key])) # doc element doc = xml.dom.minidom.Document() # vtk element root = doc.createElementNS('VTK', 'VTKFile') d = {'type': 'UnstructuredGrid', 'version': '0.1', 'byte_order': 'LittleEndian'} set_attributes(d, root) # unstructured element grid = doc.createElementNS('VTK', 'UnstructuredGrid') # piece element piece = doc.createElementNS('VTK', 'Piece') d = {'NumberOfPoints': str(Ndof), 'NumberOfCells': str(Ncells)} set_attributes(d, piece) # POINTS # points element points = doc.createElementNS('VTK', 'Points') # data element points_data = doc.createElementNS('VTK', 'DataArray') d = {'type': 'Float32', 'Name': 'vertices', 'NumberOfComponents': '3', 'format': 'ascii'} set_attributes(d, points_data) # string for data element points_data_str = doc.createTextNode(a2s(Verts)) # CELLS # points element cells = doc.createElementNS('VTK', 'Cells') # data element cells_data = doc.createElementNS('VTK', 'DataArray') d = {'type': 'Int32', 'Name': 'connectivity', 'format': 'ascii'} set_attributes(d, cells_data) # string for data element cells_data_str = doc.createTextNode(a2s(cell_ind)) # offset data element cells_offset_data = doc.createElementNS('VTK', 'DataArray') d = {'type': 'Int32', 'Name': 'offsets', 'format': 'ascii'} set_attributes(d, cells_offset_data) # string for data element cells_offset_data_str = doc.createTextNode(a2s(cell_offset.cumsum())) # offset data element cells_type_data = doc.createElementNS('VTK', 'DataArray') d = {'type': 'UInt8', 'Name': 'types', 'format': 'ascii'} set_attributes(d, cells_type_data) # string for data element cells_type_data_str = doc.createTextNode(a2s(cell_type)) # POINT DATA pointdata = doc.createElementNS('VTK', 'PointData') # pdata pdata_obj = [] pdata_str = [] for i in range(0, n_pdata): pdata_obj.append(doc.createElementNS('VTK', 'DataArray')) d = {'type': 'Float32', 'Name': 'pdata %d' % (i), 'NumberOfComponents': '1', 'format': 'ascii'} set_attributes(d, pdata_obj[i]) pdata_str.append(doc.createTextNode(a2s(pdata[:, i]))) # pvdata pvdata_obj = [] pvdata_str = [] for i in range(0, n_pvdata): pvdata_obj.append(doc.createElementNS('VTK', 'DataArray')) d = {'type': 'Float32', 'Name': 'pvdata %d' % (i), 'NumberOfComponents': '3', 'format': 'ascii'} set_attributes(d, pvdata_obj[i]) pvdata_str.append(doc.createTextNode(a2s(pvdata[:, i]))) # CELL DATA celldata = doc.createElementNS('VTK', 'CellData') # cdata cdata_obj = [] cdata_str = [] for i in range(0, n_cdata): cdata_obj.append(doc.createElementNS('VTK', 'DataArray')) d = {'type': 'Float32', 'Name': 'cdata %d' % (i), 'NumberOfComponents': '1', 'format': 'ascii'} set_attributes(d, cdata_obj[i]) cdata_str.append(doc.createTextNode(a2s(cdata_all[:, i]))) # cvdata cvdata_obj = [] cvdata_str = [] for i in range(0, n_cvdata): cvdata_obj.append(doc.createElementNS('VTK', 'DataArray')) d = {'type': 'Float32', 'Name': 'cvdata %d' % (i), 'NumberOfComponents': '3', 'format': 'ascii'} set_attributes(d, cvdata_obj[i]) cvdata_str.append(doc.createTextNode(a2s(cvdata_all[:, i]))) doc.appendChild(root) root.appendChild(grid) grid.appendChild(piece) piece.appendChild(points) points.appendChild(points_data) points_data.appendChild(points_data_str) piece.appendChild(cells) cells.appendChild(cells_data) cells.appendChild(cells_offset_data) cells.appendChild(cells_type_data) cells_data.appendChild(cells_data_str) cells_offset_data.appendChild(cells_offset_data_str) cells_type_data.appendChild(cells_type_data_str) piece.appendChild(pointdata) for i in range(0, n_pdata): pointdata.appendChild(pdata_obj[i]) pdata_obj[i].appendChild(pdata_str[i]) for i in range(0, n_pvdata): pointdata.appendChild(pvdata_obj[i]) pvdata_obj[i].appendChild(pvdata_str[i]) piece.appendChild(celldata) for i in range(0, n_cdata): celldata.appendChild(cdata_obj[i]) cdata_obj[i].appendChild(cdata_str[i]) for i in range(0, n_cvdata): celldata.appendChild(cvdata_obj[i]) cvdata_obj[i].appendChild(cvdata_str[i]) doc.writexml(fname, newl='\n') fname.close()
[docs]def write_basic_mesh(Verts, E2V=None, mesh_type='tri', pdata=None, pvdata=None, cdata=None, cvdata=None, fname='output.vtk'): """ Write mesh file for basic types of elements Parameters ---------- fname : {string} file to be written, e.g. 'mymesh.vtu' Verts : {array} coordinate array (N x D) E2V : {array} element index array (Nel x Nelnodes) mesh_type : {string} type of elements: tri, quad, tet, hex (all 3d) pdata : {array} scalar data on vertices (N x Nfields) pvdata : {array} vector data on vertices (3*Nfields x N) cdata : {array} scalar data on cells (Nfields x Nel) cvdata : {array} vector data on cells (3*Nfields x Nel) Returns ------- writes a .vtu file for use in Paraview Notes ----- The difference between write_basic_mesh and write_vtu is that write_vtu is more general and requires dictionaries of cell information. write_basic_mesh calls write_vtu Examples -------- >>> import numpy >>> from pyamg.vis import write_basic_mesh >>> Verts = numpy.array([[0.0,0.0], ... [1.0,0.0], ... [2.0,0.0], ... [0.0,1.0], ... [1.0,1.0], ... [2.0,1.0], ... [0.0,2.0], ... [1.0,2.0], ... [2.0,2.0], ... [0.0,3.0], ... [1.0,3.0], ... [2.0,3.0]]) >>> E2V = numpy.array([[0,4,3], ... [0,1,4], ... [1,5,4], ... [1,2,5], ... [3,7,6], ... [3,4,7], ... [4,8,7], ... [4,5,8], ... [6,10,9], ... [6,7,10], ... [7,11,10], ... [7,8,11]]) >>> pdata=numpy.ones((12,2)) >>> pvdata=numpy.ones((12*3,2)) >>> cdata=numpy.ones((12,2)) >>> cvdata=numpy.ones((3*12,2)) >>> write_basic_mesh(Verts, E2V=E2V, mesh_type='tri',pdata=pdata, pvdata=pvdata, cdata=cdata, cvdata=cvdata, fname='test.vtu') See Also -------- write_vtu """ if E2V is None: mesh_type = 'vertex' map_type_to_key = {'vertex': 1, 'tri': 5, 'quad': 9, 'tet': 10, 'hex': 12} if mesh_type not in map_type_to_key: raise ValueError('unknown mesh_type=%s' % mesh_type) key = map_type_to_key[mesh_type] if mesh_type == 'vertex': uidx = numpy.arange(0, Verts.shape[0]).reshape((Verts.shape[0], 1)) E2V = {key: uidx} else: E2V = {key: E2V} if cdata is not None: cdata = {key: cdata} if cvdata is not None: cvdata = {key: cvdata} write_vtu(Verts=Verts, Cells=E2V, pdata=pdata, pvdata=pvdata, cdata=cdata, cvdata=cvdata, fname=fname)
def set_attributes(d, elm): """ helper function: Set attributes from dictionary of values """ for key in d: elm.setAttribute(key, d[key]) def a2s(a): """ helper function: Convert to string """ str = '' return str.join(['%g ' % (v) for v in a.ravel()])