hedge.vtk

1 """Generic support for new-style (XML) VTK visualization data files.""" 2 3 __copyright__ = "Copyright (C) 2007 Andreas Kloeckner" 4 5 __license__ = """ 6 This program is free software: you can redistribute it and/or modify 7 it under the terms of the GNU General Public License as published by 8 the Free Software Foundation, either version 3 of the License, or 9 (at your option) any later version. 10 11 This program is distributed in the hope that it will be useful, 12 but WITHOUT ANY WARRANTY; without even the implied warranty of 13 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 14 GNU General Public License for more details. 15 16 You should have received a copy of the GNU General Public License 17 along with this program. If not, see U{http://www.gnu.org/licenses/}. 18 """ 19 20 21 22 23 import numpy 24 25 26 27 28 VTK_INT8 = "Int8" 29 VTK_UINT8 = "UInt8" 30 VTK_INT16 = "Int16" 31 VTK_UINT16 = "UInt16" 32 VTK_INT32 = "Int32" 33 VTK_UINT32 = "UInt32" 34 VTK_INT64 = "Int64" 35 VTK_UINT64 = "UInt64" 36 VTK_FLOAT32 = "Float32" 37 VTK_FLOAT64 = "Float64" 38 39 VTK_VERTEX = 1 40 VTK_POLY_VERTEX = 2 41 VTK_LINE = 3 42 VTK_POLY_LINE = 4 43 VTK_TRIANGLE = 5 44 VTK_TRIANGLE_STRIP = 6 45 VTK_POLYGON = 7 46 VTK_PIXEL = 8 47 VTK_QUAD = 9 48 VTK_TETRA = 10 49 VTK_VOXEL = 11 50 VTK_HEXAHEDRON = 12 51 VTK_WEDGE = 13 52 VTK_PYRAMID = 14 53 54 CELL_NODE_COUNT = { 55 VTK_VERTEX: 1, 56 # VTK_POLY_VERTEX: no a-priori size 57 VTK_LINE: 2, 58 # VTK_POLY_LINE: no a-priori size 59 VTK_TRIANGLE: 3, 60 # VTK_TRIANGLE_STRIP: no a-priori size 61 # VTK_POLYGON: no a-priori size 62 VTK_PIXEL: 4, 63 VTK_QUAD: 4, 64 VTK_TETRA: 4, 65 VTK_VOXEL: 8, 66 VTK_HEXAHEDRON: 8, 67 VTK_WEDGE: 6, 68 VTK_PYRAMID: 5, 69 } 70 71 72 VF_LIST_OF_COMPONENTS = 0 # [[x0,y0,z0], [x1,y1,z1] 73 VF_LIST_OF_VECTORS = 1 # [[x0,x1], [y0,y1], [z0,z1]] 74 75 _U32CHAR = numpy.dtype(numpy.uint32).char 76 77 78 79 80 # Ah, the joys of home-baked non-compliant XML goodness.

81 -class XMLElementBase(object):

82 - def __init__(self):

83 self.children = []

84

85 - def copy(self, new_children=None):

86 result = self.__class__(self.tag, self.attributes) 87 if new_children is not None: 88 result.children = new_children 89 else: 90 result.children = self.children 91 return result

92

93 - def add_child(self, child):

94 self.children.append(child)

95 96 97 98

99 -class XMLElement(XMLElementBase):

100 - def __init__(self, tag, **attributes):

101 XMLElementBase.__init__(self) 102 self.tag = tag 103 self.attributes = attributes

104

105 - def write(self, file):

106 attr_string = "".join( 107 " %s=\"%s\"" % (key,value) 108 for key,value in self.attributes.iteritems()) 109 if self.children: 110 file.write("<%s%s>\n" % (self.tag, attr_string)) 111 for child in self.children: 112 if isinstance(child, XMLElement): 113 child.write(file) 114 else: 115 # likely a string instance, write it directly 116 file.write(child) 117 file.write("</%s>\n" % self.tag) 118 else: 119 file.write("<%s%s/>\n" % (self.tag, attr_string))

120 121 122 123 124

125 -class XMLRoot(XMLElementBase):

126 - def __init__(self, child=None):

127 XMLElementBase.__init__(self) 128 if child: 129 self.add_child(child)

130

131 - def write(self, file):

132 file.write("<?xml version=\"1.0\"?>\n") 133 for child in self.children: 134 if isinstance(child, XMLElement): 135 child.write(file) 136 else: 137 # likely a string instance, write it directly 138 file.write(child)

139 140 141

142 -class EncodedBuffer:

143 - def encoder(self):

144 """Return an identifier for the binary encoding used.""" 145 raise NotImplementedError

146

147 - def compressor(self):

148 """Return an identifier for the compressor used, or None.""" 149 raise NotImplementedError

150

151 - def raw_buffer(self):

152 """Reobtain the raw buffer string object that was used to 153 construct this encoded buffer.""" 154 155 raise NotImplementedError

156

157 - def add_to_xml_element(self, xml_element):

158 """Add encoded buffer to the given C{xml_element}. 159 Return total size of encoded buffer in bytes.""" 160 161 raise NotImplementedError

162 163 164 165 166

167 -class BinaryEncodedBuffer:

168 - def __init__(self, buffer):

169 self.buffer = buffer

170

171 - def encoder(self):

172 return "binary"

173

174 - def compressor(self):

175 return None

176

177 - def raw_buffer(self):

178 return self.buffer

179

180 - def add_to_xml_element(self, xml_element):

181 raise NotImplementedError

182 183 184 185

186 -class Base64EncodedBuffer:

187 - def __init__(self, buffer):

188 from struct import pack 189 from base64 import b64encode 190 self.b64header = b64encode( 191 pack(_U32CHAR, len(buffer))) 192 self.b64data = b64encode(buffer)

193

194 - def encoder(self):

195 return "base64"

196

197 - def compressor(self):

198 return None

199

200 - def raw_buffer(self):

201 from base64 import b64decode 202 return b64decode(self.b64data)

203

204 - def add_to_xml_element(self, xml_element):

205 """Add encoded buffer to the given C{xml_element}. 206 Return total size of encoded buffer in bytes.""" 207 208 xml_element.add_child(self.b64header) 209 xml_element.add_child(self.b64data) 210 211 return len(self.b64header) + len(self.b64data)

212 213 214 215

216 -class Base64ZLibEncodedBuffer:

217 - def __init__(self, buffer):

218 from struct import pack 219 from base64 import b64encode 220 from zlib import compress 221 comp_buffer = compress(buffer) 222 comp_header = [1, len(buffer), len(buffer), len(comp_buffer)] 223 self.b64header = b64encode( 224 pack(_U32CHAR*len(comp_header), *comp_header)) 225 self.b64data = b64encode(comp_buffer)

226

227 - def encoder(self):

228 return "base64"

229

230 - def compressor(self):

231 return "zlib"

232

233 - def raw_buffer(self):

234 from base64 import b64decode 235 from zlib import decompress 236 return decompress(b64decode(self.b64data))

237

238 - def add_to_xml_element(self, xml_element):

239 """Add encoded buffer to the given C{xml_element}. 240 Return total size of encoded buffer in bytes.""" 241 242 xml_element.add_child(self.b64header) 243 xml_element.add_child(self.b64data) 244 245 return len(self.b64header) + len(self.b64data)

246 247 248 249 250 251

252 -class DataArray(object):

253 - def __init__(self, name, container, vector_padding=3, 254 vector_format=VF_LIST_OF_COMPONENTS, components=None):

255 self.name = name 256 257 if isinstance(container, DataArray): 258 self.type = container.type 259 self.components = container.components 260 self.encoded_buffer = container.encoded_buffer 261 return 262 263 def vec_type(vec): 264 if vec.dtype == numpy.int8: return VTK_INT8 265 elif vec.dtype == numpy.uint8: return VTK_UINT8 266 elif vec.dtype == numpy.int16: return VTK_INT16 267 elif vec.dtype == numpy.uint16: return VTK_UINT16 268 elif vec.dtype == numpy.int32: return VTK_INT32 269 elif vec.dtype == numpy.uint32: return VTK_UINT32 270 elif vec.dtype == numpy.int64: return VTK_INT64 271 elif vec.dtype == numpy.uint64: return VTK_UINT64 272 elif vec.dtype == numpy.float32: return VTK_FLOAT32 273 elif vec.dtype == numpy.float64: return VTK_FLOAT64 274 else: 275 raise TypeError, "Unsupported vector type '%s' in VTK writer" % (vec.dtype)

276 277 if not isinstance(container, numpy.ndarray): 278 raise ValueError, "cannot convert object of type `%s' to DataArray" % type(container) 279 280 if container.dtype == object: 281 container = numpy.array(list(container)) 282 assert container.dtype != object 283 284 if len(container.shape) > 1: 285 if vector_format == VF_LIST_OF_COMPONENTS: 286 container = container.T.copy() 287 288 assert len(container.shape) == 2, "numpy vectors of rank >2 are not supported" 289 assert container.strides[1] == container.itemsize, "2D numpy arrays must be row-major" 290 if vector_padding > container.shape[1]: 291 container = numpy.asarray(numpy.hstack(( 292 container, 293 numpy.zeros(( 294 container.shape[0], 295 vector_padding-container.shape[1], 296 ), 297 container.dtype))), order="C") 298 self.components = container.shape[1] 299 else: 300 self.components = 1 301 self.type = vec_type(container) 302 buf = buffer(container) 303 304 self.encoded_buffer = BinaryEncodedBuffer(buf)

305

306 - def get_encoded_buffer(self, encoder, compressor):

307 have_encoder = self.encoded_buffer.encoder() 308 have_compressor = self.encoded_buffer.compressor() 309 310 if (encoder, compressor) != (have_encoder, have_compressor): 311 raw_buf = self.encoded_buffer.raw_buffer() 312 313 # avoid having three copies of the buffer around temporarily 314 del self.encoded_buffer 315 316 if (encoder, compressor) == ("binary", None): 317 self.encoded_buffer = BinaryEncodedBuffer(raw_buf) 318 elif (encoder, compressor) == ("base64", None): 319 self.encoded_buffer = Base64EncodedBuffer(raw_buf) 320 elif (encoder, compressor) == ("base64", "zlib"): 321 self.encoded_buffer = Base64ZLibEncodedBuffer(raw_buf) 322 else: 323 self.encoded_buffer = BinaryEncodedBuffer(raw_buf) 324 raise ValueError, "invalid encoder/compressor pair" 325 326 have_encoder = self.encoded_buffer.encoder() 327 have_compressor = self.encoded_buffer.compressor() 328 329 assert (encoder, compressor) == (have_encoder, have_compressor) 330 331 return self.encoded_buffer

332

333 - def encode(self, compressor, xml_element):

334 ebuf = self.get_encoded_buffer("base64", compressor) 335 return ebuf.add_to_xml_element(xml_element)

336

337 - def invoke_visitor(self, visitor):

338 return visitor.gen_data_array(self)

339 340 341 342

343 -class UnstructuredGrid(object):

344 - def __init__(self, points, cells, cell_types):

345 self.point_count = len(points) 346 self.cell_count = len(cells) 347 348 self.point_count, self.points = points 349 assert self.points.name == "points" 350 351 try: 352 self.cell_count, self.cell_connectivity, \ 353 self.cell_offsets = cells 354 except: 355 self.cell_count = len(cell_types) 356 357 offsets = numpy.cumsum(numpy.fromiter( 358 (CELL_NODE_COUNT[ct] for ct in cell_types), 359 dtype=numpy.int, 360 count=len(cell_types))) 361 362 self.cell_connectivity = DataArray("connectivity", cells) 363 self.cell_offsets = DataArray("offsets", offsets) 364 365 self.cell_types = DataArray("types", cell_types) 366 367 self.pointdata = [] 368 self.celldata = []

369

370 - def copy(self):

371 return UnstructuredGrid( 372 (self.point_count, self.points), 373 (self.cell_count, self.cell_connectivity, 374 self.cell_offsets), 375 self.cell_types)

376

377 - def vtk_extension(self):

378 return "vtu"

379

380 - def invoke_visitor(self, visitor):

381 return visitor.gen_unstructured_grid(self)

382

383 - def add_pointdata(self, data_array):

384 self.pointdata.append(data_array)

385 386 387 388 389

390 -def make_vtkfile(filetype, compressor):

391 import sys 392 if sys.byteorder == "little": 393 bo = "LittleEndian" 394 else: 395 bo = "BigEndian" 396 397 kwargs = {} 398 if compressor == "zlib": 399 kwargs["compressor"] = "vtkZLibDataCompressor" 400 401 return XMLElement("VTKFile", type=filetype, version="0.1", byte_order=bo, **kwargs)

402 403 404 405

406 -class XMLGenerator(object):

407 - def __init__(self, compressor=None):

408 if compressor == "zlib": 409 try: 410 import zlib 411 except ImportError: 412 compress = False 413 elif compressor is None: 414 pass 415 else: 416 raise ValueError, "Invalid compressor name `%s'" % compressor 417 418 self.compressor = compressor

419

420 - def __call__(self, vtkobj):

421 child = self.rec(vtkobj) 422 vtkf = make_vtkfile(child.tag, self.compressor) 423 vtkf.add_child(child) 424 return XMLRoot(vtkf)

425

426 - def rec(self, vtkobj):

427 return vtkobj.invoke_visitor(self)

428 429 430 431 432

433 -class InlineXMLGenerator(XMLGenerator):

434 - def gen_unstructured_grid(self, ugrid):

435 el = XMLElement("UnstructuredGrid") 436 piece = XMLElement("Piece", 437 NumberOfPoints=ugrid.point_count, NumberOfCells=ugrid.cell_count) 438 el.add_child(piece) 439 440 pointdata = XMLElement("PointData") 441 piece.add_child(pointdata) 442 for data_array in ugrid.pointdata: 443 pointdata.add_child(self.rec(data_array)) 444 445 points = XMLElement("Points") 446 piece.add_child(points) 447 points.add_child(self.rec(ugrid.points)) 448 449 cells = XMLElement("Cells") 450 piece.add_child(cells) 451 cells.add_child(self.rec(ugrid.cell_connectivity)) 452 cells.add_child(self.rec(ugrid.cell_offsets)) 453 cells.add_child(self.rec(ugrid.cell_types)) 454 455 return el

456

457 - def gen_data_array(self, data):

458 el = XMLElement("DataArray", type=data.type, Name=data.name, 459 NumberOfComponents=data.components, format="binary") 460 data.encode(self.compressor, el) 461 el.add_child("\n") 462 return el

463 464 465 466

467 -class AppendedDataXMLGenerator(InlineXMLGenerator):

468 - def __init__(self, compressor=None):

469 InlineXMLGenerator.__init__(self, compressor) 470 471 self.base64_len = 0 472 self.app_data = XMLElement("AppendedData", encoding="base64") 473 self.app_data.add_child("_")

474

475 - def __call__(self, vtkobj):

476 xmlroot = XMLGenerator.__call__(self, vtkobj) 477 self.app_data.add_child("\n") 478 xmlroot.children[0].add_child(self.app_data) 479 return xmlroot

480

481 - def gen_data_array(self, data):

482 el = XMLElement("DataArray", type=data.type, Name=data.name, 483 NumberOfComponents=data.components, format="appended", 484 offset=self.base64_len) 485 486 self.base64_len += data.encode(self.compressor, self.app_data) 487 488 return el

489 490 491 492

493 -class ParallelXMLGenerator(XMLGenerator):

494 - def __init__(self, pathnames):

495 XMLGenerator.__init__(self, compressor=None) 496 497 self.pathnames = pathnames

498

499 - def gen_unstructured_grid(self, ugrid):

500 el = XMLElement("PUnstructuredGrid") 501 502 pointdata = XMLElement("PPointData") 503 el.add_child(pointdata) 504 for data_array in ugrid.pointdata: 505 pointdata.add_child(self.rec(data_array)) 506 507 points = XMLElement("PPoints") 508 el.add_child(points) 509 points.add_child(self.rec(ugrid.points)) 510 511 cells = XMLElement("PCells") 512 el.add_child(cells) 513 cells.add_child(self.rec(ugrid.cell_connectivity)) 514 cells.add_child(self.rec(ugrid.cell_offsets)) 515 cells.add_child(self.rec(ugrid.cell_types)) 516 517 for pn in self.pathnames: 518 el.add_child(XMLElement("Piece", Source=pn)) 519 520 return el

521

522 - def gen_data_array(self, data):

523 el = XMLElement("PDataArray", type=data.type, Name=data.name, 524 NumberOfComponents=data.components) 525 return el

526

Source Code for Module hedge.vtk