D3D to VTK

[ssolson]

Below is a custom script created by Chris and Jessica which converts D3D output to VTK format which can be used by ParaView for model analysis.

#!/usr/bin/env python3

## Convert dflowFM netCDF mapp to VTK format
## runs pretty slow, would be much faster in C++


import numpy as np
import os, sys
import xarray as xr
import matplotlib.dates as mdates
import vtk

def main():

    ##
    ## Bathymetry file
    ##
  
    filename='./FlowFM_EmBath.nc' 
    dflowMap = xr.open_dataset(filename)

    nodeX = dflowMap.variables['mesh2d_node_x'][:]
    nodeY = dflowMap.variables['mesh2d_node_y'][:]
    nodeZ = dflowMap.variables['mesh2d_node_z'][:]
    xyzArray=np.array(list(zip(nodeX.values,nodeY.values,nodeZ.values)))

    ## get connectivity
    faceConn = list(dflowMap.variables['mesh2d_face_nodes'][:].values -1) ## change count to start at 0
    faceConn = cleanConn(faceConn)

    ## start making faces
    write2DVTU('bathy.vtu',xyzArray,faceConn)

    ##
    ## DataFile
    ##

    filename='./FlowFM_map.nc'   ## Map output
    #filename='./FlowFM_0004_map.nc'   ## Test Map output
    dflowMap = xr.open_dataset(filename)

    nTimes = dflowMap.coords.dims['time']
    dateList = dflowMap['time'].values.flatten()
    timeList = [ mdates.date2num(date)-mdates.date2num(dateList[0]) for date in dateList ]

    ## time of interest
    ## eventually, we can animate all times
    itime = -1


    nodeX = dflowMap.variables['mesh2d_node_x'].values[:]
    nodeY = dflowMap.variables['mesh2d_node_y'].values[:]
    ## not always defined
    #nodeZ = dflowMap.variables['mesh2d_node_z'][:]

    nNodes = len(nodeX)


    ## I'm not sure why this variables is available sometimes and not others
    try:
        sigma = dflowMap.variables['mesh2d_layer_sigma'].values[:]
        nLayers = len(sigma)
    except:
        sigma = dflowMap['mesh2d_nLayers'].values[:]
        nLayers = len(sigma)
        ## if no sigma variable, assume even layers
        sigma = np.linspace(0,1,nLayers+1)[:-1]-1.0 + 1.0/nLayers/2.0
        #ds = np.linspace(0,1,nLayers+1)[1:]  ## calculated later

    ## these are at faces
    ## again, variable names are different in different files
    try:
        bedLevel    = dflowMap.variables['mesh2d_bldepth'].values  # constant in time
    except:
        bedLevel    = dflowMap.mesh2d_flowelem_bl.values

    waterHeight = dflowMap.variables['mesh2d_waterdepth'].values

    ## get connectivity
    edgeConn = list(dflowMap.variables['mesh2d_edge_nodes'][:].values -1) ## change count to start at 0
    faceConn = list(dflowMap.variables['mesh2d_face_nodes'][:].values -1) ## change count to start at 0
    nEdges = len(edgeConn)
    nFaces = len(faceConn)
    edgeConn = cleanConn(edgeConn) ## remove NaNs, and shorten triangles to 3 entries
    faceConn = cleanConn(faceConn)

    ## open the file
    ugrid = open3DVTU(nodeX,nodeY,faceConn,sigma,bedLevel,waterHeight[itime])  ## water height determines mesh

    velocity=np.stack((dflowMap.variables['mesh2d_ucx'].values[itime],
                        dflowMap.variables['mesh2d_ucy'].values[itime],
                        dflowMap.variables['mesh2d_ucz'].values[itime]),
                        axis=-1)
    ## add velocity data
    ugrid = addCellVector(ugrid,'Velocity',velocity,nFaces,nLayers)

    surfaceHeight = dflowMap.variables['mesh2d_s1'].values

    ## add depth data
    ugrid = addFaceScalar(ugrid,'Height',surfaceHeight[itime],nFaces,nLayers)
    ugrid = addFaceScalar(ugrid,'Depth',waterHeight[itime],nFaces,nLayers)
    ugrid = addFaceScalar(ugrid,'BedLevel',bedLevel,nFaces,nLayers)

    ## write out the VTK file
    write3DVTU(ugrid,'data.vtu')



def cleanConn(connList):
    for cell_idx in range(len(connList)):
        cell = connList[cell_idx]
        ## remove NaNs (which mean it's a triangle, not a quad)
        while np.isnan(connList[cell_idx][-1]):
            connList[cell_idx] = connList[cell_idx][:-1]
        connList[cell_idx] = list(connList[cell_idx].astype(int))
    return connList
 
def open3DVTU(nodeX,nodeY,faceConn,sigma,bedLevel,waterHeight):

    ugrid = vtk.vtkUnstructuredGrid()
    points = vtk.vtkPoints()
    cell_array = vtk.vtkCellArray()

    nNodes  = len(nodeX)
    nLayers = len(sigma)
    ds = np.zeros(nLayers)
    ds[0]= (sigma[0]+1.0)*2.0  ## sigma is negative
    for iLayer in range(1,nLayers):
        ds[iLayer]= (sigma[iLayer]+1.0)*2.0 - ds[iLayer-1]

    ## calculate node z
    ## scatter z to nodes and average
    nodeZ0 = np.zeros(nNodes,dtype=float)  ## bottom z
    nodeZ1 = np.zeros(nNodes,dtype=float)  ## top z
    nface  = np.zeros(nNodes,dtype=float)
    for iface in range(len(faceConn)):
         for inode in faceConn[iface]:
             nodeZ0[inode] += bedLevel[iface]
             nodeZ1[inode] += bedLevel[iface] + waterHeight[iface]
             nface[inode] += 1
    ## do average
    for inode in range(nNodes):
        nodeZ0[inode] = nodeZ0[inode]/nface[inode]
        nodeZ1[inode] = nodeZ1[inode]/nface[inode]

    #calculate points and store them

    ## layer 0
    for idx in range(len(nodeX)):
        points.InsertNextPoint([nodeX[idx],nodeY[idx],nodeZ0[idx]])

    for iLayer in range(1,nLayers):
        Z = ds[iLayer-1]*(nodeZ1[:] - nodeZ0[:]) + nodeZ0[:]
        for idx in range(len(nodeX)):
            points.InsertNextPoint([nodeX[idx],nodeY[idx],Z[idx]])

    ## layer N
    for idx in range(len(nodeX)):
        points.InsertNextPoint([nodeX[idx],nodeY[idx],nodeZ1[idx]])

    ugrid.SetPoints(points)

    ## make cells
    vtk_elemType=[]
    for iLayer in range(nLayers):
        for face_idx in range(len(faceConn)):
            nodes = faceConn[face_idx]
            nPoly=len(nodes)
            if nPoly == 4:
                poly = vtk.vtkHexahedron()
                vtk_elemType.append(vtk.VTK_HEXAHEDRON)
            elif nPoly == 3:
                poly = vtk.vtkWedge()
                vtk_elemType.append(vtk.VTK_WEDGE)

            for iVert in range(nPoly):
                poly.GetPointIds().SetId(iVert,      nodes[iVert]+nNodes*iLayer)
                poly.GetPointIds().SetId(iVert+nPoly,nodes[iVert]+nNodes*(iLayer+1))
            cell_array.InsertNextCell(poly)

    ugrid.SetCells(vtk_elemType,cell_array)

    return ugrid

def addCellVector(ugrid,arrayName,cellData,nFaces,nLayers):

    data_array = vtk.vtkFloatArray()
    data_array.SetName(arrayName)
    data_array.SetNumberOfComponents(3)
    for iLayer in range(nLayers):
        for iFace in range(nFaces):
            dataEntry = cellData[iFace,iLayer,:]
            data_array.InsertNextTuple(dataEntry)

    ugrid.GetCellData().AddArray(data_array)

    return ugrid

def addFaceScalar(ugrid,arrayName,cellData,nFaces,nLayers):

    data_array = vtk.vtkFloatArray()
    data_array.SetName(arrayName)
    data_array.SetNumberOfComponents(1)
    for iLayer in range(nLayers):
        for iFace in range(nFaces):
            dataEntry = cellData[iFace]  ## just repeat this for all layers
            data_array.InsertNextValue(dataEntry)

    ugrid.GetCellData().AddArray(data_array)

    return ugrid

def write3DVTU(ugrid,filename):

    writer = vtk.vtkXMLUnstructuredGridWriter()
    writer.SetFileName(filename)

    writer.SetInputData(ugrid)
    writer.Write()

def write2DVTU(filename,vertList,connList):
    ugrid = vtk.vtkUnstructuredGrid()
    points = vtk.vtkPoints()
    cell_array = vtk.vtkCellArray()

    #points.InsertPoints(vertList)
    for p in vertList:
        points.InsertNextPoint(p)
    
    for cell_idx in range(len(connList)):
        cell = connList[cell_idx]

        nPoly=len(cell)
        poly = vtk.vtkPolygon()
        poly.GetPointIds().SetNumberOfIds(nPoly)

        for iVert in range(nPoly):
            poly.GetPointIds().SetId(iVert,cell[iVert])
        cell_array.InsertNextCell(poly)

    ugrid.SetPoints(points)
    ugrid.SetCells(vtk.VTK_POLYGON, cell_array)

    writer = vtk.vtkXMLUnstructuredGridWriter()
    writer.SetFileName(filename)

    writer.SetInputData(ugrid)
    writer.Write()

if __name__ == "__main__":
    main()