VTK.py 4.38 KB
Newer Older
1
2
import zlib
import struct
3
import os
4
import numpy as np
5
import re
6
7
8
from base64 import b64decode,b64encode
from xml.etree import ElementTree as ET

9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
def _typename(a) :
    if a.dtype == np.int32 :
        return b"Int32"
    elif a.dtype == np.int64 :
        return b"Int64"
    elif a.dtype == np.float64 :
        return b"Float64"
    elif a.dtype == np.float32 :
        return b"Float32"
    else :
        raise(NameError("unown array type : " + str(a.dtype)))

def _write_array_ascii(f,a,attr) :
    f.write(b'<DataArray %s type="%s" format="ascii">\n'%(attr,_typename(a)))
    f.write((" ".join(map(str,a.flat))).encode("utf-8"))
    f.write(b'\n</DataArray>\n')


def _write_array(f,appended,a,ds,attr) :
    f.write(b'<DataArray %s type="%s" format="appended" offset="%i"/>\n'%(attr,_typename(a),len(appended)-1))
29
    data = zlib.compress(a,2)
30
31
#    return appended + struct.pack("iiii",1,a.size*ds,0,len(data)) + data
    return appended + b64encode(struct.pack("iiii",1,a.size*ds,0,len(data))) + b64encode(data)
32

33
34
35
def write(basename,i,t,elements,x,data,field_data) :
    appended = b"_"
    filename = "%s_%05d.vtu"%(basename,i)
36
37
    f = open(filename,"wb")
    f.write(b'<?xml version="1.0"?>\n')
38
39
40
41
    f.write(b'<VTKFile type="UnstructuredGrid" format="0.1" compressor="vtkZLibDataCompressor">\n')
    f.write(b'<UnstructuredGrid>\n')
    nel = elements.shape[0]
    nptbyel = elements.shape[1]
42
43
44
    npt = x.shape[0]
    f.write(b'<Piece NumberOfPoints="%i" NumberOfCells="%i">\n'%(npt,nel))
    f.write(b'<Points>\n')
45
    appended = _write_array(f,appended,x,8,b'NumberOfComponents="3"')
46
    f.write(b'</Points>\n')
47
48
49
50
51
52
53
54
55
56
57
58
    f.write(b'<Cells>\n')
    appended = _write_array(f,appended,elements,4,b'Name="connectivity"')
    offsets = np.array(range(nptbyel,(nel+1)*nptbyel,nptbyel),np.int32)
    appended = _write_array(f,appended,offsets,4,b'Name="offsets"')
    types = np.ones(nel,np.int32)*(5 if elements.shape[1] == 3 else 13)
    appended = _write_array(f,appended,types,4,b'Name="types"')
    f.write(b'</Cells>\n')
    f.write(b'<PointData>\n')
    for name,v in data :
        vc = np.ascontiguousarray(v)
        appended = _write_array(f,appended,vc,8,b'Name="%s" NumberOfComponents="%i"' % (name.encode("utf8"),vc.shape[-1]))
    f.write(b'</PointData>\n')
59
    f.write(b'</Piece>\n')
60
61
62
63
64
65
66
67
68
    f.write(b'<FieldData>\n');
    for name,v in field_data :
        vc = np.ascontiguousarray(v)
        appended = _write_array(f,appended,vc,4,b'Name="%s" NumberOfTuples="%i" NumberOfComponents="%i"' % (name,vc.shape[0],vc.shape[1]))
    f.write(b'</FieldData>\n');
    f.write(b'</UnstructuredGrid>\n')
    f.write(b'<AppendedData encoding="base64">\n')
    f.write(appended)
    f.write(b'</AppendedData>\n')
69
70
    f.write(b'</VTKFile>\n')
    f.close()
71
72
73
74
75
76
77
78
79
80
81
82
83
84


def _get_array_info(a) :
    ncomp = int(a["NumberOfComponents"]) if "NumberOfComponents" in a else None
    ntuples = int(a["NumberOfTuples"]) if "NumberOfTuples" in a else None
    return (a["Name"],a["type"],int(a["offset"]),(ntuples,ncomp))

def _get_binary_array(appended,pos,dtype,shape) :
    data = appended[pos:pos+24]
    _,n,_,s = struct.unpack("iiii",b64decode(data))
    data = appended[pos+24:pos+24-(-s//3)*4]
    data = b64decode(data)
    data = zlib.decompress(data)
    return np.frombuffer(data,dtype).reshape(shape)
85
86

def read(fname) :
87
88
89
    f = open(fname,"rb")
    root = ET.parse(f).getroot()
    grid = root.find("UnstructuredGrid")
90
    piece = grid.find("Piece")
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119

    npoints = int(piece.attrib["NumberOfPoints"])
    ncells = int(piece.attrib["NumberOfCells"])

    appended = root.find("AppendedData").text
    shift = appended.find("_")+1

    dim = 2
    pointsoffset = int(piece.find("Points/DataArray").attrib["offset"])
    x = _get_binary_array(appended,shift+pointsoffset,np.float64,(npoints,3))

    for f in piece.find("Cells") :
        if f.attrib["Name"] == "connectivity" :
            connectivityoffset = int(f.attrib["offset"])
    el = _get_binary_array(appended,shift+connectivityoffset,np.int32,(ncells,dim+1))

    fdata = []
    for f in grid.find("FieldData") :
        name,ntype,offset,shape = _get_array_info(f.attrib)
        v = _get_binary_array(appended,shift+offset,np.int32,shape)
        fdata.append((name,v))

    data = []
    for f in piece.find("PointData") :
        name,ntype,offset,shape = _get_array_info(f.attrib)
        v = _get_binary_array(appended,shift+offset,np.float64,(npoints,shape[1]))
        data.append((name,v))

    return x,el,data,fdata