solvcon.io.vtkxml

1 # -*- coding: UTF-8 -*- 2 # 3 # Copyright (C) 2008-2010 Yung-Yu Chen <yyc@solvcon.net>. 4 # 5 # This program is free software; you can redistribute it and/or modify 6 # it under the terms of the GNU General Public License as published by 7 # the Free Software Foundation; either version 2 of the License, or 8 # (at your option) any later version. 9 # 10 # This program is distributed in the hope that it will be useful, 11 # but WITHOUT ANY WARRANTY; without even the implied warranty of 12 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 13 # GNU General Public License for more details. 14 # 15 # You should have received a copy of the GNU General Public License along 16 # with this program; if not, write to the Free Software Foundation, Inc., 17 # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA. 18 19 """ 20 VTK XML file. 21 22 For a detailed description of the file format, see: 23 http://www.geophysik.uni-muenchen.de/intranet/it-service/applications/paraview/vtk-file-formats/ 24 """

25 26 -class VtkXmlWriter(object):

27 """ 28 Base class for VTK XML format data writer. 29 30 @cvar cltpn_map: mapper for cell type number. 31 @ctype cltpn_map: numpy.ndarray 32 33 @ivar appended: append data at the end of XML file. Forces binary. 34 @itype appended: bool 35 @ivar binary: flag for BINARY (False for ASCII). 36 @itype binary: bool 37 @ivar encoding: encoding for binary data ('base64' or 'raw'). Default is 38 raw. Must be base64 if inline (not appended). 39 @itype encoding: str 40 @ivar compressor: compressor for binary data. Can only be 'gz' or ''. 41 @itype compressor: str 42 @ivar fpdtype: floating-point data type (single/double). 43 @itype fpdtype: numpy.dtype 44 @ivar blk: corresponding block object. 45 @itype blk: solvcon.block.Block 46 """ 47 from numpy import array, float32, float64, int32 48 cltpn_map = array([1, 3, 9, 5, 12, 10, 13, 14], dtype='int32') 49 dtype_map = dict( 50 float32='Float32', 51 float64='Float64', 52 int32='Int32', 53 ) 54 del array, float32, float64, int32 55

56 - def __init__(self, blk, *args, **kw):

57 self.appended = kw.pop('appended', True) 58 binary = kw.pop('binary', False) 59 self.binary = True if self.appended else binary 60 encoding = kw.pop('encoding', 'raw') 61 if self.binary and not self.appended: 62 self.encoding = 'base64' 63 else: 64 self.encoding = encoding 65 self.compressor = kw.pop('compressor', 'gz') 66 fpdtype = kw.pop('fpdtype', None) 67 if fpdtype == None: 68 fpdtype = blk.fpdtype 69 self.fpdtype = fpdtype 70 self.blk = blk 71 super(VtkXmlWriter, self).__init__() 72 self.lvl = 0

73

74 - def _tag_open(self, tag, attr=None, newline=True, close=False):

75 """ 76 Create an opening XML tag. 77 78 @param tag: the name of tag. 79 @type tag: str 80 @keyword attr: attribute to be assigned to the tag in the format 81 ('name1', value1), ('name2', value2) ... 82 @type attr: list 83 @keyword newline: add newline at the end. Default is True. 84 @type newline: bool 85 @keyword close: close the tag while opening. Default is False. 86 @type close: 87 @return: contructed opening tag. 88 @rtype: str 89 """ 90 if self.lvl < 0: 91 raise ValueError('lvl %d < 0 when opening %s' % (lvl, tag)) 92 lvl = self.lvl 93 msg = [' '*(lvl*2)] 94 finish = ' />' if close else '>' 95 if attr: 96 msg.append('<%s %s%s' % ( 97 tag, 98 ' '.join(['%s="%s"'%(key, str(val)) for key, val in attr]), 99 finish, 100 )) 101 else: 102 msg.append('<%s%s' % (tag, finish)) 103 if newline: 104 msg.append('\n') 105 if not close: 106 self.lvl += 1 107 return ''.join(msg)

108 - def _tag_close(self, tag, newline=True):

109 """ 110 Create a closing XML tag. 111 112 @param tag: the name of tag. 113 @type tag: str 114 @keyword newline: add newline at the end. Default is True. 115 @type newline: bool 116 @return: contructed closing tag. 117 @rtype: str 118 """ 119 self.lvl -= 1 120 if self.lvl < 0: 121 raise ValueError('lvl %d < 0 when closing %s' % (lvl, tag)) 122 lvl = self.lvl 123 msg = [' '*(lvl*2)] 124 msg.append('</%s>'%tag) 125 if newline: 126 msg.append('\n') 127 return ''.join(msg)

128

129 - def _create_data(self, arr):

130 """ 131 Create data buffer from array. 132 133 @param arr: input array. 134 @type arr: numpy.ndarray 135 @return: the created data buffer. 136 @rtype: buffer 137 """ 138 from struct import pack 139 from base64 import standard_b64encode 140 from zlib import compress 141 data = arr.tostring() 142 if self.compressor == 'gz': 143 osize = len(data) 144 data = compress(data) 145 zsize = len(data) 146 if self.encoding == 'base64': 147 data = standard_b64encode(data) 148 if self.compressor == 'gz': 149 size = pack('iiii', 1, osize, osize, zsize) 150 else: 151 size = pack('i', len(data)) 152 if self.encoding == 'base64': 153 size = standard_b64encode(size) 154 return ''.join([size, data])

155

156 - def _write_darr(self, arr, outf, aplist, attr):

157 """ 158 Write data array to a stream. 159 160 @param arr: data array to be written. 161 @type arr: numpy.ndarray 162 @param outf: output file stream. 163 @type outf: file 164 @param aplist: appended data list. 165 @type aplist: list 166 @param attr: additional attributes to the DataArray tag. 167 @type attr: list 168 @return: nothing 169 """ 170 # craft attributes. 171 mattr = [ 172 ('type', self.dtype_map[str(arr.dtype)]), 173 ] 174 if self.binary: 175 if aplist is None: 176 mattr.append(('format', 'binary')) 177 mattr.append(('encoding', self.encoding)) 178 else: 179 mattr.append(('format', 'appended')) 180 mattr.append(('offset', sum([len(dat) for dat in aplist]))) 181 else: 182 mattr.append(('format', 'ascii')) 183 mattr.extend(attr) 184 attr = mattr 185 # write. 186 if self.binary: 187 outf.write(self._tag_open('DataArray', attr, close=True)) 188 data = self._create_data(arr) 189 if aplist is None: 190 outf.write(data) 191 outf.write('\n') 192 else: 193 aplist.append(data) 194 else: 195 outf.write(self._tag_open('DataArray', attr)) 196 arr.tofile(outf, sep=' ') 197 outf.write('\n') 198 outf.write(self._tag_close('DataArray'))

199

200 - def _write_appended(self, aplist, outf):

201 """ 202 Write appended data list to a stream. 203 204 @param aplist: the appendedn data list to be outputted. 205 @type aplist: list 206 @param outf: the file stream to be written. 207 @type outf: file 208 @return: nothing 209 """ 210 outf.write(self._tag_open('AppendedData', [ 211 ('encoding', self.encoding), 212 ])) 213 outf.write('_') 214 for data in aplist: 215 outf.write(data) 216 outf.write('\n') 217 outf.write(self._tag_close('AppendedData'))

218

219 -class VtkXmlUstGridWriter(VtkXmlWriter):

220 """ 221 VTK XML unstructured mesh file format. Capable for ASCII or BINARY. 222 223 @ivar cache_grid: flag to cache the grid data (True) or not (False). 224 Default True. 225 @itype cache_grid: bool 226 @ivar scalars: dictionary holding scalar data. 227 @itype scalars: dict 228 @ivar vectors: dictionary holding vector data. 229 @itype vectors: dict 230 @ivar griddata: cached grid data. 231 @itype griddata: str 232 """

233 - def __init__(self, blk, *args, **kw):

234 self.cache_grid = kw.pop('cache_grid', True) 235 self.scalars = kw.pop('scalars', dict()) 236 self.vectors = kw.pop('vectors', dict()) 237 super(VtkXmlUstGridWriter, self).__init__(blk, *args, **kw) 238 self.griddata = None

239

240 - def _convert_varr(self, arr):

241 """ 242 Helper to convert vector data array from a block. 243 244 @param arr: array to be converted. It will be copied and remain 245 untouched. 246 @type arr: numpy.ndarray 247 @return: converted array. 248 @rtype: numpy.ndarray 249 """ 250 from numpy import empty 251 arr = arr.copy() 252 ndim = self.blk.ndim 253 nit = arr.shape[0] 254 if ndim == 2: 255 arrn = empty((nit, ndim+1), dtype=arr.dtype) 256 arrn[:,2] = 0.0 257 arrn[:,:2] = arr[:,:] 258 arr = arrn 259 return arr

260 261 @staticmethod

262 - def _convert_clnds(clnds):

263 """ 264 Convert nodes in cells into a compressed array format. 265 266 @param clnds: the array to be compressed. 267 @type clnds: numpy.ndarray 268 @return: the compressed array. 269 @rtype: numpy.ndarray 270 """ 271 from numpy import empty 272 arr = empty(clnds[:,0].sum(), dtype='int32') 273 ncell = clnds.shape[0] 274 icl = 0 275 it = 0 276 while icl < ncell: 277 ncl = clnds[icl,0] 278 arr[it:it+ncl] = clnds[icl,1:ncl+1] 279 it += ncl 280 icl += 1 281 return arr

282

283 - def write(self, outf, close_on_finish=False):

284 """ 285 Write to file. 286 287 @param outf: output file object or file name. 288 @type outf: file str 289 @keyword close_on_finish: flag close on finishing (True). Default 290 False. If outf is file name, the output file will be close no 291 matter what is set in this flag. 292 @type close_on_finish: bool 293 @return: nothing 294 """ 295 if isinstance(outf, str): 296 mode = 'wb' if self.binary else 'w' 297 outf = open(outf, mode) 298 close_on_finish = True 299 # write header. 300 outf.write('<?xml version="1.0"?>\n') 301 attr = [ 302 ('type', 'UnstructuredGrid'), 303 ('version', '0.1'), 304 ('byte_order', 'LittleEndian'), 305 ] 306 if self.compressor == 'gz': 307 attr.append(('compressor', 'vtkZLibDataCompressor')) 308 outf.write(self._tag_open('VTKFile', attr)) 309 outf.write(self._tag_open('UnstructuredGrid')) 310 outf.write(self._tag_open('Piece', [ 311 ('NumberOfPoints', self.blk.nnode), 312 ('NumberOfCells', self.blk.ncell), 313 ])) 314 aplist = list() if self.appended else None 315 # write points. 316 outf.write(self._tag_open('Points')) 317 self._write_darr(self._convert_varr(self.blk.ndcrd), outf, aplist, [ 318 ('NumberOfComponents', 3), 319 ]) 320 outf.write(self._tag_close('Points')) 321 # write cells. 322 outf.write(self._tag_open('Cells')) 323 self._write_darr(self._convert_clnds(self.blk.clnds), outf, aplist, [ 324 ('Name', 'connectivity'), 325 ]) 326 self._write_darr(self.blk.clnds[:,0].cumsum(dtype='int32'), outf, 327 aplist, [ 328 ('Name', 'offsets'), 329 ]) 330 self._write_darr(self.cltpn_map[self.blk.cltpn], outf, aplist, [ 331 ('Name', 'types'), 332 ]) 333 outf.write(self._tag_close('Cells')) 334 # write footer. 335 outf.write(self._tag_close('Piece')) 336 outf.write(self._tag_close('UnstructuredGrid')) 337 if aplist: 338 self._write_appended(aplist, outf) 339 outf.write(self._tag_close('VTKFile')) 340 if close_on_finish: 341 outf.close() 342 # discard griddata if not cached. 343 if not self.cache_grid: 344 self.griddata = None

345

Source Code for Module solvcon.io.vtkxml