pyfrm2xdmf: add first and second order pyramids, prisms and tetrahedrons

1 files changed, 123 insertions, 54 deletions

diff --git a/pyfrm2xdmf b/pyfrm2xdmf
index 172b76c..93dd9d3 100755
--- a/pyfrm2xdmf
+++ b/pyfrm2xdmf
@@ -14,6 +14,8 @@ def meshFile(param):
     base, ext = os.path.splitext(param)
     if ext.lower() != '.pyfrm':
         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
 
 def writeHeader(xdmfFile):
@@ -26,58 +28,58 @@ def writeHeader(xdmfFile):
 
 def writeTopology(xdmfFile, tType, nCells, connFile):
     "write Topology element"
-    xdmfFile.write('   <Topology TopologyType="{}" NumberOfElements="{}">\n'.format(tType, nCells))
-    xdmfFile.write('    <xi:include href="{}"/>\n'.format(connFile))
-    xdmfFile.write('   </Topology>\n')
+    xdmfFile.write('    <Topology TopologyType="{}" NumberOfElements="{}">\n'.format(tType, nCells))
+    xdmfFile.write('     <xi:include href="{}"/>\n'.format(connFile))
+    xdmfFile.write('    </Topology>\n')
     return
 
 def writeGeometry(xdmfFile, nDims, nCells, nodeArray, pyfrm, dataset):
     "write Geometry element"
     nVerts = len(nodeArray)
     if nDims == 2:
-        xdmfFile.write('   <Geometry GeometryType="X_Y">\n')   # co-ordinates in separate arrays
+        xdmfFile.write('    <Geometry GeometryType="X_Y">\n')   # co-ordinates in separate arrays
     else:
-        xdmfFile.write('   <Geometry GeometryType="X_Y_Z">\n') # co-ordinates in separate arrays
+        xdmfFile.write('    <Geometry GeometryType="X_Y_Z">\n') # co-ordinates in separate arrays
     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))))
+        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>\n')
-    xdmfFile.write('   </Geometry>\n')
+        xdmfFile.write('     </DataItem>\n')
+    xdmfFile.write('    </Geometry>\n')
     return
 
 def writeHyperSlab(xdmfFile, coord, nDims, nCells, nodeArray, pyfrm, dataset):
     "write HyperSlab element"
     nVerts = len(nodeArray)
     for vert in nodeArray:
-        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
-        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
+        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
+        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
+        xdmfFile.write('          </DataItem>\n')
+        xdmfFile.write('          <DataItem\n')
+        xdmfFile.write('          Name="Points" \n')
+        xdmfFile.write('          Dimensions="{} {} {}"\n'.format(nVerts, nCells, nDims))
+        xdmfFile.write('          Format="HDF">\n')
+        xdmfFile.write('          {}:/{}\n'.format(pyfrm, dataset))
         xdmfFile.write('         </DataItem>\n')
-        xdmfFile.write('         <DataItem\n')
-        xdmfFile.write('         Name="Points" \n')
-        xdmfFile.write('         Dimensions="{} {} {}"\n'.format(nVerts, nCells, nDims))
-        xdmfFile.write('         Format="HDF">\n')
-        xdmfFile.write('         {}:/{}\n'.format(pyfrm, dataset))
-        xdmfFile.write('        </DataItem>\n')
-        xdmfFile.write('      </DataItem>\n')
+        xdmfFile.write('       </DataItem>\n')
     return
 
 def writeAttribute(xdmfFile, tag):
     "write Attribute element"
-    xdmfFile.write('   <Attribute Name="Partition" Center="Grid">\n')
-    xdmfFile.write('    <DataItem\n')
-    xdmfFile.write('     Dimensions="1"\n')
-    xdmfFile.write('     Format="XML">\n')
-    xdmfFile.write('     {}\n'.format(tag)) # tag with partition number
-    xdmfFile.write('    </DataItem>\n')
-    xdmfFile.write('   </Attribute>\n')
+    xdmfFile.write('    <Attribute Name="Partition" Center="Grid">\n')
+    xdmfFile.write('     <DataItem\n')
+    xdmfFile.write('      Dimensions="1"\n')
+    xdmfFile.write('      Format="XML">\n')
+    xdmfFile.write('      {}\n'.format(tag)) # tag with partition number
+    xdmfFile.write('     </DataItem>\n')
+    xdmfFile.write('    </Attribute>\n')
     return
 
 def writeConnectivities(connFile, nCells, nVerts):
@@ -120,7 +122,7 @@ def readH5lsOutput(os_output):
             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 cell is quadrangular
+            if re.search('quad', line.decode()):  # check whether cells are quadrilaterals
                 if nnodes == 4:
                     quad4[partno] = ncells
                 elif nnodes == 9:
@@ -133,7 +135,7 @@ def readH5lsOutput(os_output):
                     quad36[partno] = ncells
                 else:
                     print("unknown cell order")
-            elif re.search('tri', line.decode()): # check whether cell is triangular
+            elif re.search('tri', line.decode()): # check whether cells are triangles
                 if nnodes == 3:
                     tri3[partno] = ncells
                 elif nnodes == 6:
@@ -146,7 +148,7 @@ def readH5lsOutput(os_output):
                     tri21[partno] = ncells
                 else:
                     print("unknown cell order")
-            elif re.search('hex', line.decode()): # check whether cell is hexahedral
+            elif re.search('hex', line.decode()): # check whether cells are hexahedrons
                 if nnodes == 8:
                     hex8[partno] = ncells
                 elif nnodes == 27:
@@ -159,6 +161,51 @@ def readH5lsOutput(os_output):
                     hex216[partno] = ncells
                 else:
                     print("unknown cell order")
+            elif re.search('pyr', line.decode()):  # check whether cells are pyramids
+                if nnodes == 5:
+                    pyr5[partno] = ncells
+                elif nnodes == 14:
+                    pyr14[partno] = ncells
+                elif nnodes == 30:
+                    pyr30[partno] = ncells
+                elif nnodes == 55:
+                    pyr55[partno] = ncells
+                elif nnodes == 91:
+                    pyr91[partno] = ncells
+                elif nnodes == 140:
+                    pyr140[partno] = ncells
+                else:
+                    print("unknown cell order")
+            elif re.search('pri', line.decode()): # check whether cells are prisms
+                if nnodes == 6:
+                    pri6[partno] = ncells
+                elif nnodes == 18:
+                    pri18[partno] = ncells
+                elif nnodes == 40:
+                    pri40[partno] = ncells
+                elif nnodes == 75:
+                    pri75[partno] = ncells
+                elif nnodes == 126:
+                    pri126[partno] = ncells
+                elif nnodes == 196:
+                    pri196[partno] = ncells
+                else:
+                    print("unknown cell order")
+            elif re.search('tet', line.decode()): # check whether cells are tetrahedrons
+                if nnodes == 4:
+                    tet4[partno] = ncells
+                elif nnodes == 10:
+                    tet10[partno] = ncells
+                elif nnodes == 20:
+                    tet20[partno] = ncells
+                elif nnodes == 35:
+                    tet35[partno] = ncells
+                elif nnodes == 56:
+                    tet56[partno] = ncells
+                elif nnodes == 84:
+                    tet84[partno] = ncells
+                else:
+                    print("unknown cell order")
             else:
                 print("unknown cell type")
                 break
@@ -196,28 +243,55 @@ 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'}
 
-# node identification (see pyfr/readers/nodemaps.py): reduce high-order cells to first order
+# node identification (see pyfr/readers/nodemaps.py): reduces high-order cells to first order
+# example, for second-order pyramid:
+# >>> 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
 
-# sort datasets
-hex8Keys  = list(nhex8.keys())    # keys are partition numbers
+# 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())
-quad4Keys = list(nquad4.keys())
 tri3Keys  = list(ntri3.keys())
-allKeys   = quad4Keys + tri3Keys + hex8Keys + hex27Keys   # concatenate keys
-numTypes  = Counter(allKeys)        # number of types present in each partition
-partKeys  = list(numTypes.keys())   # partition 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)
+# number of types present in each partition
+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]   # list of cell dictionaries
+              nquad4, nquad9, nquad16, nquad25, nquad36]   
 numCellTypes = len(partitions)
 
 if partKeys:
@@ -226,31 +300,26 @@ if partKeys:
     writeHeader(g)
     
     for part in partKeys:
-        if numTypes[part] > 1: # check whether partition contain multiple cell types
-            g.write('  <Grid Name="Partition{}" GridType="Collection">\n'.format(part))
-        else:
-            g.write('  <Grid Name="Partition{}" GridType="Uniform">\n'.format(part))
-    
+        g.write('  <Grid Name="Partition{}" GridType="Collection">\n'.format(part))
         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))
-                if numTypes[part] > 1:
-                    g.write('  <Grid Name="cellType{}" GridType="Uniform">\n'.format(cellType))
+                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)
-                if numTypes[part] > 1:
-                    g.write('  </Grid>\n')
+                writeAttribute(g, part)
+                g.write('   </Grid>\n')
                 # connectivities file
                 writeConnectivities(xfname, partitions[cellType][part], 
                                     len(nodeIDs[cellType]))
-    
-        writeAttribute(g, part)
         g.write('  </Grid>\n')
     
     writeFooter(g)
     g.close()
+else:
+    print("mesh does not contain any supported cell types")