xref: /petsc/lib/petsc/bin/petsc_gen_xdmf.py (revision d1f89d58d5508d2aa21db4fceee013bffa40d53d)
1#!/usr/bin/env python3
2# VTK high order: https://blog.kitware.com/modeling-arbitrary-order-lagrange-finite-elements-in-the-visualization-toolkit/
3import h5py
4import numpy as np
5import os, sys
6
7class Xdmf:
8  def __init__(self, filename):
9    self.filename = filename
10    self.cellMap  = {1 : {1 : 'Polyvertex', 2 : 'Polyline'}, 2 : {3 : 'Triangle', 4 : 'Quadrilateral'}, 3 : {4 : 'Tetrahedron', 6: 'Wedge', 8 : 'Hexahedron'}}
11    self.typeMap = {'scalar' : 'Scalar', 'vector' : 'Vector', 'tensor' : 'Tensor6', 'matrix' : 'Matrix'}
12
13  def writeHeader(self, fp, hdfFilename):
14    fp.write('''\
15<?xml version="1.0" ?>
16<!DOCTYPE Xdmf SYSTEM "Xdmf.dtd" [
17<!ENTITY HeavyData "%s">
18]>
19''' % os.path.basename(hdfFilename))
20    fp.write('\n<Xdmf>\n  <Domain Name="domain">\n')
21    return
22
23  def writeCells(self, fp, topologyPath, numCells, numCorners, cellsName = "cells"):
24    fp.write('''\
25    <DataItem Name="%s"
26              ItemType="Uniform"
27              Format="HDF"
28              NumberType="Float" Precision="8"
29              Dimensions="%d %d">
30      &HeavyData;:/%s/cells
31    </DataItem>
32''' % (cellsName, numCells, numCorners, topologyPath))
33    return
34
35  def writeVertices(self, fp, geometryPath, numVertices, spaceDim):
36    fp.write('''\
37    <DataItem Name="vertices"
38              Format="HDF"
39              Dimensions="%d %d">
40      &HeavyData;:/%s/vertices
41    </DataItem>
42    <!-- ============================================================ -->
43''' % (numVertices, spaceDim, geometryPath))
44    return
45
46  def writeLocations(self, fp, numParticles, spaceDim):
47    fp.write('''\
48    <DataItem Name="particle_coordinates"
49              Format="HDF"
50              Dimensions="%d %d">
51      &HeavyData;:/particles/coordinates
52    </DataItem>
53''' % (numParticles, spaceDim))
54    return
55
56  def writeTimeGridHeader(self, fp, time):
57    fp.write('''\
58    <Grid Name="TimeSeries" GridType="Collection" CollectionType="Temporal">
59      <Time TimeType="List">
60        <DataItem Format="XML" NumberType="Float" Dimensions="%d">
61          ''' % (len(time)))
62    fp.write(' '.join([str(float(t)) for t in time]))
63    fp.write('''
64        </DataItem>
65      </Time>
66''')
67    return
68
69  #http://www.xdmf.org/index.php/XDMF_Model_and_Format#Topology
70  def writeHybridSpaceGridHeader(self, fp):
71    fp.write('      <Grid Name="domain" GridType="Collection">\n')
72    return
73
74  def writeSpaceGridHeader(self, fp, numCells, numCorners, cellDim, spaceDim, cellsName = "cells"):
75    fp.write('''\
76      <Grid Name="domain" GridType="Uniform">
77        <Topology
78           TopologyType="%s"
79           NumberOfElements="%d">
80          <DataItem Reference="XML">
81            /Xdmf/Domain/DataItem[@Name="%s"]
82          </DataItem>
83        </Topology>
84        <Geometry GeometryType="%s">
85          <DataItem Reference="XML">
86            /Xdmf/Domain/DataItem[@Name="vertices"]
87          </DataItem>
88        </Geometry>
89''' % (self.cellMap[cellDim][numCorners], numCells, cellsName, "XYZ" if spaceDim > 2 else "XY"))
90    return
91
92  def writeField(self, fp, numSteps, timestep, spaceDim, name, f, domain):
93    vtype = f[1].attrs['vector_field_type']
94    if sys.version_info[0] >= 3:
95      vtype = vtype.decode()
96    xtype = self.typeMap[vtype]
97    if len(f[1].shape) > 2:
98      dof = f[1].shape[1]
99      bs  = f[1].shape[2]
100    elif len(f[1].shape) > 1:
101      if numSteps > 0 and numSteps == f[1].shape[0]:
102        dof = f[1].shape[1]
103        bs  = 1
104      else:
105        dof = f[1].shape[0]
106        bs  = f[1].shape[1]
107    else:
108      dof = f[1].shape[0]
109      bs  = 1
110    if xtype == 'Scalar':
111      hsc = 1
112      hbs = bs
113    else:
114      hsc = bs
115      hbs = bs
116      bs = 1
117    for b in range(bs):
118      fname = '%s.%d' % (f[0], b) if bs > 1 else f[0]
119      if numSteps > 0:
120        fp.write('''\
121            <Attribute
122               Name="%s"
123               Type="%s"
124               Center="%s">
125              <DataItem ItemType="HyperSlab"
126		    Dimensions="1 %d %d"
127		    Type="HyperSlab">
128                <DataItem
129                   Dimensions="3 3"
130                   Format="XML">
131                  %d 0 %d
132                  1 1 1
133                  1 %d %d
134                </DataItem>
135                <DataItem
136                   DataType="Float" Precision="8"
137                   Dimensions="%d %d %d"
138                   Format="HDF">
139                  &HeavyData;:%s
140                </DataItem>
141              </DataItem>
142            </Attribute>
143''' % (fname, xtype, domain, dof, hsc, timestep, b, dof, hsc, numSteps, dof, hbs, name))
144      else:
145        fp.write('''\
146            <Attribute
147               Name="%s"
148               Type="%s"
149               Center="%s">
150              <DataItem ItemType="HyperSlab"
151		    Dimensions="%d %d"
152		    Type="HyperSlab">
153                <DataItem
154                   Dimensions="3 2"
155                   Format="XML">
156                  0 %d
157                  1 1
158                  %d %d
159                </DataItem>
160                <DataItem
161                   DataType="Float" Precision="8"
162                   Dimensions="%d %d"
163                   Format="HDF">
164                  &HeavyData;:%s
165                </DataItem>
166              </DataItem>
167            </Attribute>
168''' % (fname, xtype, domain, dof, hsc, b, dof, hsc, dof, hbs, name))
169    return
170
171  def writeSpaceGridFooter(self, fp):
172    fp.write('      </Grid>\n')
173    return
174
175  def writeParticleGridHeader(self, fp, numParticles, spaceDim):
176    fp.write('''\
177      <Grid Name="particle_domain" GridType="Uniform">
178        <Topology TopologyType="Polyvertex" NodesPerElement="%d" />
179        <Geometry GeometryType="%s">
180          <DataItem Reference="XML">/Xdmf/Domain/DataItem[@Name="particle_coordinates"]</DataItem>
181        </Geometry>
182''' % (numParticles, "XYZ" if spaceDim > 2 else "XY"))
183    return
184
185  def writeParticleField(self, fp, fieldname, numParticles, numComp):
186    fp.write('''\
187    <Attribute Name="particles/%s">
188      <DataItem Name="%s"
189                Format="HDF"
190                Dimensions="%d %d">
191                &HeavyData;:/particle_fields/%s
192      </DataItem>
193    </Attribute>
194''' % (fieldname, fieldname, numParticles, numComp, fieldname))
195    return
196
197  def writeTimeGridFooter(self, fp):
198    fp.write('    </Grid>\n')
199    return
200
201  def writeFooter(self, fp):
202    fp.write('  </Domain>\n</Xdmf>\n')
203    return
204
205  def write(self, hdfFilename, topologyPath, numCells, numCorners, cellDim, htopologyPath, numHCells, numHCorners, geometryPath, numVertices, spaceDim, time, vfields, cfields, numParticles, pfields):
206    useTime = len(time) > 0
207    numSteps = len(time)
208    n = max(numSteps, 1)
209    with open(self.filename, 'w') as fp:
210      self.writeHeader(fp, hdfFilename)
211      # Field information
212      self.writeCells(fp, topologyPath, numCells, numCorners)
213      if numHCells:
214        self.writeCells(fp, htopologyPath, numHCells, numHCorners, "hcells")
215      self.writeVertices(fp, geometryPath, numVertices, spaceDim)
216      if useTime: self.writeTimeGridHeader(fp, time)
217      for t in range(n):
218        if numHCells:
219          self.writeHybridSpaceGridHeader(fp)
220          self.writeSpaceGridHeader(fp, numHCells, numHCorners, cellDim, spaceDim, "hcells")
221          self.writeSpaceGridFooter(fp)
222        self.writeSpaceGridHeader(fp, numCells, numCorners, cellDim, spaceDim)
223        for vf in vfields: self.writeField(fp, numSteps, t, spaceDim, '/vertex_fields/'+vf[0], vf, 'Node')
224        for cf in cfields: self.writeField(fp, numSteps, t, spaceDim, '/cell_fields/'+cf[0], cf, 'Cell')
225        self.writeSpaceGridFooter(fp)
226        if numHCells:
227          self.writeSpaceGridFooter(fp)
228      if useTime: self.writeTimeGridFooter(fp)
229      if numParticles:
230        self.writeLocations(fp, numParticles, spaceDim)
231        if useTime: self.writeTimeGridHeader(fp, time)
232        for t in range(len(time)):
233          self.writeParticleGridHeader(fp, numParticles, spaceDim)
234          for pf in pfields:
235            self.writeParticleField(fp, pf[0], numParticles, int(pf[1].attrs['Nc']))
236          self.writeSpaceGridFooter(fp)
237        if useTime: self.writeTimeGridFooter(fp)
238      self.writeFooter(fp)
239    return
240
241def generateXdmf(hdfFilename, xdmfFilename = None):
242  if xdmfFilename is None:
243    xdmfFilename = os.path.splitext(hdfFilename)[0] + '.xmf'
244  # Read mesh
245  h5          = h5py.File(hdfFilename, 'r')
246  if 'viz' in h5 and 'geometry' in h5['viz']:
247    geomPath  = 'viz/geometry'
248    geom      = h5['viz']['geometry']
249  else:
250    geomPath  = 'geometry'
251    geom      = h5['geometry']
252  if 'viz' in h5 and 'topology' in h5['viz']:
253    topoPath  = 'viz/topology'
254    topo      = h5['viz']['topology']
255  else:
256    topoPath  = 'topology'
257    topo      = h5['topology']
258  if 'viz' in h5 and 'hybrid_topology' in h5['viz']:
259    htopoPath = 'viz/hybrid_topology'
260    htopo     = h5['viz']['hybrid_topology']
261  else:
262    htopoPath = None
263    htopo     = None
264  vertices    = geom['vertices']
265  numVertices = vertices.shape[0]
266  spaceDim    = vertices.shape[1]
267  cells       = topo['cells']
268  numCells    = cells.shape[0]
269  numCorners  = cells.shape[1]
270  cellDim     = topo['cells'].attrs['cell_dim']
271  if htopo:
272    hcells      = htopo['cells']
273    numHCells   = hcells.shape[0]
274    numHCorners = hcells.shape[1]
275  else:
276    numHCells   = 0
277    numHCorners = 0
278  if 'time' in h5:
279    time      = np.array(h5['time']).flatten()
280  else:
281    time      = []
282  vfields     = []
283  cfields     = []
284  pfields     = []
285  pfields     = []
286  if 'vertex_fields' in h5: vfields = h5['vertex_fields'].items()
287  if 'cell_fields' in h5: cfields = h5['cell_fields'].items()
288  numParticles = 0
289  if 'particles' in h5:
290    numParticles = h5['particles']['coordinates'].shape[0]
291  if 'particle_fields' in h5: pfields = h5['particle_fields'].items()
292
293  # Write Xdmf
294  Xdmf(xdmfFilename).write(hdfFilename, topoPath, numCells, numCorners, cellDim, htopoPath, numHCells, numHCorners, geomPath, numVertices, spaceDim, time, vfields, cfields, numParticles, pfields)
295  h5.close()
296  return
297
298if __name__ == '__main__':
299  for f in sys.argv[1:]:
300    generateXdmf(f)
301