pyfrm2xdmf: include third, fourth, fifth and sixth order elements

1 files changed, 154 insertions, 105 deletions

diff --git a/pyfrm2xdmf b/pyfrm2xdmf
index 93dd9d3..1a47fd6 100755
--- a/pyfrm2xdmf
+++ b/pyfrm2xdmf
@@ -13,7 +13,8 @@ def meshFile(param):
     global base
     base, ext = os.path.splitext(param)
     if ext.lower() != '.pyfrm':
-        raise argparse.ArgumentTypeError('Mesh file must have a .pyfrm extension')
+        raise argparse.ArgumentTypeError('Mesh file must have a .pyfrm '
+                                          'extension')
     if not os.path.exists(param):
         raise argparse.ArgumentTypeError("{} not found".format(param))
     return param
@@ -22,13 +23,15 @@ def writeHeader(xdmfFile):
     "write XDMF header"
     xdmfFile.write('<?xml version="1.0" ?>\n')
     xdmfFile.write('<!DOCTYPE Xdmf SYSTEM "Xdmf.dtd" []>\n')
-    xdmfFile.write('<Xdmf xmlns:xi="http://www.w3.org/2003/XInclude" Version="2.2">\n')
+    xdmfFile.write('<Xdmf xmlns:xi="http://www.w3.org/2003/XInclude"')
+    xdmfFile.write(' Version="2.2">\n')
     xdmfFile.write(' <Domain>\n')
     return
 
 def writeTopology(xdmfFile, tType, nCells, connFile):
     "write Topology element"
-    xdmfFile.write('    <Topology TopologyType="{}" NumberOfElements="{}">\n'.format(tType, nCells))
+    xdmfFile.write('    <Topology TopologyType="{}"'.format(tType))
+    xdmfFile.write(' NumberOfElements="{}">\n'.format(nCells))
     xdmfFile.write('     <xi:include href="{}"/>\n'.format(connFile))
     xdmfFile.write('    </Topology>\n')
     return
@@ -36,14 +39,22 @@ def writeTopology(xdmfFile, tType, nCells, connFile):
 def writeGeometry(xdmfFile, nDims, nCells, nodeArray, pyfrm, dataset):
     "write Geometry element"
     nVerts = len(nodeArray)
+    # co-ordinates in separate arrays
     if nDims == 2:
-        xdmfFile.write('    <Geometry GeometryType="X_Y">\n')   # co-ordinates in separate arrays
+        xdmfFile.write('    <Geometry GeometryType="X_Y">\n')
     else:
-        xdmfFile.write('    <Geometry GeometryType="X_Y_Z">\n') # co-ordinates in separate arrays
+        xdmfFile.write('    <Geometry GeometryType="X_Y_Z">\n')
+
     for coord in range(nDims):
-        xdmfFile.write('     <DataItem ItemType="Function" Dimensions="{}"\n'.format(nVerts*nCells)) # 1D-array
-        xdmfFile.write('       Function="JOIN({})">\n'.format(' ; '.join("$" + str(k) for k in range(nVerts))))
-        writeHyperSlab(xdmfFile, coord, nDims, nCells, nodeArray, pyfrm, dataset)
+        xdmfFile.write('     <DataItem ItemType="Function"')
+        # 1D-array
+        xdmfFile.write(' Dimensions="{}"\n'.format(nVerts*nCells))
+        xdmfFile.write('       Function="JOIN({})">\n'.format(
+                           ' ; '.join( "$" + str(k) for k in range(nVerts))))
+        writeHyperSlab(
+                xdmfFile,
+                coord, nDims, nCells,
+                nodeArray, pyfrm, dataset)
         xdmfFile.write('     </DataItem>\n')
     xdmfFile.write('    </Geometry>\n')
     return
@@ -55,16 +66,21 @@ def writeHyperSlab(xdmfFile, coord, nDims, nCells, nodeArray, pyfrm, dataset):
         xdmfFile.write('       <DataItem ItemType="HyperSlab"\n')
         xdmfFile.write('         Dimensions="{} 1 1"\n'.format(nCells))
         xdmfFile.write('         Type="HyperSlab">\n')
-        xdmfFile.write('         <DataItem\n') # start, stride and count of hyperslab region
+        # start, stride and count of hyperslab region
+        xdmfFile.write('         <DataItem\n')
         xdmfFile.write('          Dimensions="3 3"\n')
         xdmfFile.write('          Format="XML">\n')
-        xdmfFile.write('          {:<3} 0   {}\n'.format(vert, coord)) # select vertex and co-ordinate (format is vertex, cell, co-ordinate)
-        xdmfFile.write('          1   1   1\n')             # select every cell, for this vertex and co-ordinate
-        xdmfFile.write('          1   {} 1\n'.format(nCells)) # loop over cells
+        # select vertex and co-ordinate (format is vertex, cell, co-ordinate)
+        xdmfFile.write('          {:<3} 0   {}\n'.format(vert, coord))
+        # select every cell, for this vertex and co-ordinate
+        xdmfFile.write('          1   1   1\n')
+        # loop over cells
+        xdmfFile.write('          1   {} 1\n'.format(nCells))
         xdmfFile.write('          </DataItem>\n')
         xdmfFile.write('          <DataItem\n')
         xdmfFile.write('          Name="Points" \n')
-        xdmfFile.write('          Dimensions="{} {} {}"\n'.format(nVerts, nCells, nDims))
+        xdmfFile.write('          Dimensions=')
+        xdmfFile.write('"{} {} {}"\n'.format(nVerts, nCells, nDims))
         xdmfFile.write('          Format="HDF">\n')
         xdmfFile.write('          {}:/{}\n'.format(pyfrm, dataset))
         xdmfFile.write('         </DataItem>\n')
@@ -77,7 +93,8 @@ def writeAttribute(xdmfFile, tag):
     xdmfFile.write('     <DataItem\n')
     xdmfFile.write('      Dimensions="1"\n')
     xdmfFile.write('      Format="XML">\n')
-    xdmfFile.write('      {}\n'.format(tag)) # tag with partition number
+    # tag with partition number
+    xdmfFile.write('      {}\n'.format(tag))
     xdmfFile.write('     </DataItem>\n')
     xdmfFile.write('    </Attribute>\n')
     return
@@ -88,13 +105,13 @@ def writeConnectivities(connFile, nCells, nVerts):
     cf.write('<DataItem DataType="Int"\n')
     cf.write('  Dimensions="{} {}"\n'.format(nCells, nVerts))
     cf.write('  Format="XML">\n')
-    
+
     for i in range (0, nCells):
         cf.write(' ')
         for j in range(nVerts):
             cf.write(' ' + repr(j*nCells+i).ljust(1))
         cf.write('\n')
-    
+
     cf.write('</DataItem>\n')
     cf.close()
     print('connectivities written to ' + connFile)
@@ -108,21 +125,29 @@ def writeFooter(xdmfFile):
 
 def readH5lsOutput(os_output):
     "read output from 'h5ls' command"
-    tet4,  tet10, tet20,  tet35,  tet56,  tet84  = ({} for i in range(6)) # tet cells
-    pri6,  pri18, pri40,  pri75,  pri126, pri196 = ({} for i in range(6)) # pri cells
-    pyr5,  pyr14, pyr30,  pyr55,  pyr91,  pyr140 = ({} for i in range(6)) # pyr cells
-    hex8,  hex27, hex64,  hex125, hex216  = ({} for i in range(5))        # hex cells
-    tri3,  tri6,  tri10,  tri15,  tri21   = ({} for i in range(5))        # tri cells
-    quad4, quad9, quad16, quad25, quad36  = ({} for i in range(5))        # quad cells
+    # tetrahedral cells (first order to sixth order)
+    tet4,  tet10, tet20,  tet35,  tet56,  tet84  = ({} for i in range(6))
+    # prism cells
+    pri6,  pri18, pri40,  pri75,  pri126, pri196 = ({} for i in range(6))
+    # pyramid cells
+    pyr5,  pyr14, pyr30,  pyr55,  pyr91,  pyr140 = ({} for i in range(6))
+    # hexahedral cells (first order to fifth order)
+    hex8,  hex27, hex64,  hex125, hex216  = ({} for i in range(5))
+    # triangular cells
+    tri3,  tri6,  tri10,  tri15,  tri21   = ({} for i in range(5))
+    # quadrilateral cells
+    quad4, quad9, quad16, quad25, quad36  = ({} for i in range(5))
 
     for line in os_output.splitlines():
-        spt = re.search('spt', line.decode())     # restrict to 'spt' arrays
+        # restrict to 'spt' arrays
+        spt = re.search('spt', line.decode())
         if spt:
             chunk = line.decode().split()
             partno = int(re.search('\d+', chunk[0]).group())
             nnodes = int(re.search('(\d+),', line.decode()).group(1))
             ncells = int(re.search(' (\d+),', line.decode()).group(1))
-            if re.search('quad', line.decode()):  # check whether cells are quadrilaterals
+            # check whether cells are quadrilaterals
+            if re.search('quad', line.decode()):
                 if nnodes == 4:
                     quad4[partno] = ncells
                 elif nnodes == 9:
@@ -135,7 +160,8 @@ def readH5lsOutput(os_output):
                     quad36[partno] = ncells
                 else:
                     print("unknown cell order")
-            elif re.search('tri', line.decode()): # check whether cells are triangles
+            # check whether cells are triangles
+            elif re.search('tri', line.decode()):
                 if nnodes == 3:
                     tri3[partno] = ncells
                 elif nnodes == 6:
@@ -148,7 +174,8 @@ def readH5lsOutput(os_output):
                     tri21[partno] = ncells
                 else:
                     print("unknown cell order")
-            elif re.search('hex', line.decode()): # check whether cells are hexahedrons
+            # check whether cells are hexahedrons
+            elif re.search('hex', line.decode()):
                 if nnodes == 8:
                     hex8[partno] = ncells
                 elif nnodes == 27:
@@ -161,7 +188,8 @@ def readH5lsOutput(os_output):
                     hex216[partno] = ncells
                 else:
                     print("unknown cell order")
-            elif re.search('pyr', line.decode()):  # check whether cells are pyramids
+            # check whether cells are pyramids
+            elif re.search('pyr', line.decode()):
                 if nnodes == 5:
                     pyr5[partno] = ncells
                 elif nnodes == 14:
@@ -176,7 +204,8 @@ def readH5lsOutput(os_output):
                     pyr140[partno] = ncells
                 else:
                     print("unknown cell order")
-            elif re.search('pri', line.decode()): # check whether cells are prisms
+            # check whether cells are prisms
+            elif re.search('pri', line.decode()):
                 if nnodes == 6:
                     pri6[partno] = ncells
                 elif nnodes == 18:
@@ -191,7 +220,8 @@ def readH5lsOutput(os_output):
                     pri196[partno] = ncells
                 else:
                     print("unknown cell order")
-            elif re.search('tet', line.decode()): # check whether cells are tetrahedrons
+            # check whether cells are tetrahedrons
+            elif re.search('tet', line.decode()):
                 if nnodes == 4:
                     tet4[partno] = ncells
                 elif nnodes == 10:
@@ -210,27 +240,26 @@ def readH5lsOutput(os_output):
                 print("unknown cell type")
                 break
 
-    return (tet4,  tet10, tet20,  tet35,  tet56,  tet84,
-            pri6,  pri18, pri40,  pri75,  pri126, pri196,
-            pyr5,  pyr14, pyr30,  pyr55,  pyr91,  pyr140,
-            hex8,  hex27, hex64,  hex125, hex216,
-            tri3,  tri6,  tri10,  tri15,  tri21,
-            quad4, quad9, quad16, quad25, quad36)
+    # list of cell dictionaries
+    cellDictList = [tet4,  tet10, tet20,  tet35,  tet56,  tet84,
+                    pri6,  pri18, pri40,  pri75,  pri126, pri196,
+                    pyr5,  pyr14, pyr30,  pyr55,  pyr91,  pyr140,
+                    hex8,  hex27, hex64,  hex125, hex216,
+                    tri3,  tri6,  tri10,  tri15,  tri21,
+                    quad4, quad9, quad16, quad25, quad36]
+
+    return cellDictList
 
 # read command line arguments
-parser = argparse.ArgumentParser(description="extract connectivities from mesh file")
+parser = argparse.ArgumentParser(
+            description="extract connectivities from mesh file"
+                        ": write xdmf file")
 parser.add_argument("mesh", help="mesh file (.pyfrm)", type=meshFile)
 args = parser.parse_args()
 
 # use 'h5ls' command to provide array dimensions
 h5lsOutput = check_output(["h5ls", args.mesh])
-
-(ntet4,  ntet10, ntet20,  ntet35,  ntet56,  ntet84,
-npri6,  npri18, npri40,  npri75,  npri126, npri196,
-npyr5,  npyr14, npyr30,  npyr55,  npyr91,  npyr140,
-nhex8,  nhex27, nhex64,  nhex125, nhex216,
-ntri3,  ntri6,  ntri10,  ntri15,  ntri21,
-nquad4, nquad9, nquad16, nquad25, nquad36) = readH5lsOutput(h5lsOutput)
+partitions = readH5lsOutput(h5lsOutput)
 
 # cell types
 firstOrderCellType = ['tet', 'tet', 'tet', 'tet', 'tet', 'tet',
@@ -239,87 +268,107 @@ firstOrderCellType = ['tet', 'tet', 'tet', 'tet', 'tet', 'tet',
                       'hex', 'hex', 'hex', 'hex', 'hex',
                       'tri', 'tri', 'tri', 'tri', 'tri',
                       'quad', 'quad', 'quad', 'quad', 'quad']
-ndims              = {'quad': 2, 'tri': 2, 'hex': 3, 'pyr': 3, 'pri': 3, 'tet':3}
-xdmfTopologyType   = {'quad': 'Quadrilateral', 'tri': 'Triangle', 'hex': 'Hexahedron', 
-                      'pyr' : 'Pyramid',       'pri': 'Wedge',    'tet': 'Tetrahedron'}
+xdmfTopologyType   = {'quad': 'Quadrilateral', 'tri': 'Triangle',
+                      'hex' : 'Hexahedron',    'pyr': 'Pyramid',
+                      'pri' : 'Wedge',         'tet': 'Tetrahedron'}
+ndims              = {'quad': 2, 'tri': 2,
+                      'hex' : 3, 'pyr': 3, 'pri': 3, 'tet': 3}
 
-# node identification (see pyfr/readers/nodemaps.py): reduces high-order cells to first order
-# example, for second-order pyramid:
+# node identification: reduces high-order cells to first order
+# (see pyfr/readers/nodemaps.py)
+# Example, for second-order pyramids having 14 solution points and 5 vertices:
 # >>> from pyfr.readers.nodemaps import GmshNodeMaps
 # >>> [GmshNodeMaps.to_pyfr['pyr', 14][i] for i in range(5)]
 nodeIDs = {}
-nodeIDs[29] = [0, 2, 8, 6]                 # quad9
-nodeIDs[28] = [0, 1, 3, 2]                 # quad4
-nodeIDs[24] = [0, 2, 5]                    # tri6
-nodeIDs[23] = [0, 1, 2]                    # tri3
-nodeIDs[19] = [0, 2, 8, 6, 18, 20, 26, 24] # hex27
-nodeIDs[18] = [0, 1, 3, 2, 4, 5, 7, 6]     # hex8
-nodeIDs[13] = [0, 2, 8, 6, 13]             # pyr14
-nodeIDs[12] = [0, 1, 3, 2, 4]              # pyr5
-nodeIDs[7]  = [0, 2, 5, 12, 14, 17]        # pri18
-nodeIDs[6]  = [0, 1, 2, 3, 4, 5]           # pri6
-nodeIDs[1]  = [0, 2, 5, 9]                 # tet10
-nodeIDs[0]  = [0, 1, 2, 3]                 # tet4
+nodeIDs[0]  = [0, 1, 2, 3]                       # tet4
+nodeIDs[1]  = [0, 2, 5, 9]                       # tet10
+nodeIDs[2]  = [0, 3, 9, 19]                      # tet20
+nodeIDs[3]  = [0, 4, 14, 34]                     # tet35
+nodeIDs[4]  = [0, 5, 20, 55]                     # tet56
+nodeIDs[5]  = [0, 6, 27, 83]                     # tet84
+nodeIDs[6]  = [0, 1, 2, 3, 4, 5]                 # pri6
+nodeIDs[7]  = [0, 2, 5, 12, 14, 17]              # pri18
+nodeIDs[8]  = [0, 3, 9, 30, 33, 39]              # pri40
+nodeIDs[9]  = [0, 4, 14, 60, 64, 74]             # pri75
+nodeIDs[10] = [0, 5, 20, 105, 110, 125]          # pri126
+nodeIDs[11] = [0, 6, 27, 168, 174, 195]          # pri196
+nodeIDs[12] = [0, 1, 3, 2, 4]                    # pyr5
+nodeIDs[13] = [0, 2, 8, 6, 13]                   # pyr14
+nodeIDs[14] = [0, 3, 15, 12, 29]                 # pyr30
+nodeIDs[15] = [0, 4, 24, 20, 54]                 # pyr55
+nodeIDs[16] = [0, 5, 35, 30, 90]                 # pyr91
+nodeIDs[17] = [0, 6, 48, 42, 139]                # pyr140
+nodeIDs[18] = [0, 1, 3, 2, 4, 5, 7, 6]           # hex8
+nodeIDs[19] = [0, 2, 8, 6, 18, 20, 26, 24]       # hex27
+nodeIDs[20] = [0, 3, 15, 12, 48, 51, 63, 60]     # hex64
+nodeIDs[21] = [0, 4, 24, 20, 100, 104, 124, 120] # hex125
+nodeIDs[22] = [0, 5, 35, 30, 180, 185, 215, 210] # hex216
+nodeIDs[23] = [0, 1, 2]                          # tri3
+nodeIDs[24] = [0, 2, 5]                          # tri6
+nodeIDs[25] = [0, 3, 9]                          # tri10
+nodeIDs[26] = [0, 4, 14]                         # tri15
+nodeIDs[27] = [0, 5, 20]                         # tri21
+nodeIDs[28] = [0, 1, 3, 2]                       # quad4
+nodeIDs[29] = [0, 2, 8, 6]                       # quad9
+nodeIDs[30] = [0, 3, 15, 12]                     # quad16
+nodeIDs[31] = [0, 4, 24, 20]                     # quad25
+nodeIDs[32] = [0, 5, 35, 30]                     # quad36
+
+# number of supported cell types
+numCellTypes = len(partitions)
 
 # keys are partition numbers
-tet4Keys  = list(ntet4.keys())    
-tet10Keys = list(ntet10.keys())
-pri6Keys  = list(npri6.keys())    
-pri18Keys = list(npri18.keys())
-pyr5Keys  = list(npyr5.keys())    
-pyr14Keys = list(npyr14.keys())
-hex8Keys  = list(nhex8.keys())    
-hex27Keys = list(nhex27.keys())
-tri3Keys  = list(ntri3.keys())
-tri6Keys  = list(ntri6.keys())
-quad4Keys = list(nquad4.keys())
-quad9Keys = list(nquad9.keys())
-# concatenate keys
-allKeys   = (tet4Keys + tet10Keys +
-             pri6Keys + pri18Keys +
-             pyr5Keys + pyr14Keys +
-             hex8Keys + hex27Keys +
-             tri3Keys + tri6Keys +
-             quad4Keys + quad9Keys)
+allKeys = []
+for cd in partitions:
+    allKeys += list(cd.keys())
+
 # number of types present in each partition
-numTypes  = Counter(allKeys)        
+numTypes  = Counter(allKeys)
+
 # list of partition keys
-partKeys  = list(numTypes.keys())   
-# list of cell dictionaries
-partitions = [ntet4,  ntet10, ntet20,  ntet35,  ntet56,  ntet84,
-              npri6,  npri18, npri40,  npri75,  npri126, npri196,
-              npyr5,  npyr14, npyr30,  npyr55,  npyr91,  npyr140,
-              nhex8,  nhex27, nhex64,  nhex125, nhex216,
-              ntri3,  ntri6,  ntri10,  ntri15,  ntri21,
-              nquad4, nquad9, nquad16, nquad25, nquad36]   
-numCellTypes = len(partitions)
+partKeys  = list(numTypes.keys())
 
 if partKeys:
     # write files
     g = open(os.path.join(base + '.xdmf'), 'w')
     writeHeader(g)
-    
+
     for part in partKeys:
-        g.write('  <Grid Name="Partition{}" GridType="Collection">\n'.format(part))
+        g.write(
+            '  <Grid Name="Partition{}" GridType="Collection">\n'.format(part))
+        # step through all supported cell types
         for cellType in range(numCellTypes):
-            if part in partitions[cellType]: # check whether these cells exist in this partition
-                xfname = os.path.join('con_' + firstOrderCellType[cellType] + '_p' + str(part) + '.xml')
-                dname  = os.path.join('spt_' + firstOrderCellType[cellType] + '_p' + str(part))
-                g.write('   <Grid Name="{}s[{}]" GridType="Uniform">\n'.format(xdmfTopologyType[firstOrderCellType[cellType]], part))
-                writeTopology(g, xdmfTopologyType[firstOrderCellType[cellType]],
-                              partitions[cellType][part], xfname)
-                writeGeometry(g, ndims[firstOrderCellType[cellType]], 
-                               partitions[cellType][part], 
-                               nodeIDs[cellType], 
-                               args.mesh, dname)
+            # check whether these cells exist in this partition
+            if part in partitions[cellType]:
+                idname = firstOrderCellType[cellType] + '_p' + str(part)
+                xfname = os.path.join('con_' + idname + '.xml')
+                dname  = os.path.join('spt_' + idname)
+                g.write(
+                    '   <Grid Name="{}s[{}]" GridType="Uniform">\n'.format(
+                        xdmfTopologyType[firstOrderCellType[cellType]],
+                        part))
+                writeTopology(
+                        g,
+                        xdmfTopologyType[firstOrderCellType[cellType]],
+                        partitions[cellType][part],
+                        xfname)
+                writeGeometry(
+                        g,
+                        ndims[firstOrderCellType[cellType]],
+                        partitions[cellType][part],
+                        nodeIDs[cellType],
+                        args.mesh,
+                        dname)
                 writeAttribute(g, part)
                 g.write('   </Grid>\n')
                 # connectivities file
-                writeConnectivities(xfname, partitions[cellType][part], 
-                                    len(nodeIDs[cellType]))
+                writeConnectivities(
+                        xfname,
+                        partitions[cellType][part],
+                        len(nodeIDs[cellType]))
         g.write('  </Grid>\n')
-    
+
     writeFooter(g)
     g.close()
 else:
-    print("mesh does not contain any supported cell types")
+    print("no supported cell types")