-
Notifications
You must be signed in to change notification settings - Fork 0
/
pvuLoder.py
88 lines (67 loc) · 3.5 KB
/
pvuLoder.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
import os
import numpy as np
import xmltodict, meshio
class pvuFile(object):
def __init__(self,file,**kwargs):
self.file=file
self.dataDir=kwargs.get('dataDir',os.path.split(file)[0])
with open(file) as data:
self.pXML = xmltodict.parse(data.read())
# store fields for convenience
self.fields=self.pXML['VTKFile']['PUnstructuredGrid']['PPointData']['PDataArray']
self.connectivity = None
self.coordinates = None
self.node_data = None
def load(self):
conlist=[] # list of 2D connectivity arrays
coordlist=[] # global, concatenated coordinate array
nodeDictList=[] # list of node_data dicts, same length as conlist
con_offset=-1
pieces = self.pXML['VTKFile']['PUnstructuredGrid']['Piece']
if not isinstance(pieces,list):
pieces = [pieces]
for mesh_id,src in enumerate(pieces):
print(f"processing {mesh_id} of {pieces} vtu files")
mesh_name="connect{meshnum}".format(meshnum=mesh_id+1) # connect1, connect2, etc.
srcFi=os.path.join(self.dataDir,src['@Source']) # full path to .vtu file
[con,coord,node_d]=self.loadPiece(srcFi,mesh_name,con_offset+1)
con_offset=con.max()
conlist.append(con.astype("i8"))
coordlist.extend(coord.astype("f8"))
nodeDictList.append(node_d)
print("concatenating data")
self.connectivity=conlist
self.coordinates=np.array(coordlist)
self.node_data=nodeDictList
def loadPiece(self,srcFi,mesh_name,connectivity_offset=0):
meshPiece=meshio.read(srcFi) # read it in with meshio
coords=meshPiece.points # coords and node_data are already global
cell_type = list(meshPiece.cells_dict.keys())[0]
connectivity=meshPiece.cells_dict[cell_type] # 2D connectivity array
# parse node data
node_data=self.parseNodeData(meshPiece.point_data,connectivity,mesh_name)
# offset the connectivity matrix to global value
connectivity=np.array(connectivity)+connectivity_offset
return [connectivity,coords,node_data]
def parseNodeData(self,point_data,connectivity,mesh_name):
# for each field, evaluate field data by index, reshape to match connectivity
con1d=connectivity.ravel()
conn_shp=connectivity.shape
comp_hash={0:'cx',1:'cy',2:'cz'}
def rshpData(data1d):
return np.reshape(data1d[con1d],conn_shp)
node_data={}
for fld in self.fields:
nm=fld['@Name']
if nm in point_data.keys():
if '@NumberOfComponents' in fld.keys() and int(fld['@NumberOfComponents'])>1:
# we have a vector, deal with components
for component in range(int(fld['@NumberOfComponents'])):
comp_name=nm+'_'+comp_hash[component] # e.g., velocity_cx
m_F=(mesh_name,comp_name) # e.g., ('connect1','velocity_cx')
node_data[m_F]=rshpData(point_data[nm][:,component])
else:
# just a scalar!
m_F=(mesh_name,nm) # e.g., ('connect1','T')
node_data[m_F]=rshpData(point_data[nm])
return node_data