MHKiT-Software / MHKiT-Python

MHKiT-Python provides the marine renewable energy (MRE) community tools for data processing, visualization, quality control, resource assessment, and device performance.
https://mhkit-software.github.io/MHKiT/
BSD 3-Clause "New" or "Revised" License
47 stars 45 forks source link

D3D to VTK #300

Open ssolson opened 3 months ago

ssolson commented 3 months ago

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()
ssolson commented 1 month ago

Still working in a standalone environment. I have converted the above into a functional approach compatible with MHKiT and works with our current Tanana river example. I still need to integrate this approach with the current MHKiT d3d module.

from netCDF2VTK import *
filename = "./tanana50_map.nc"
dflow_map = read_d3d_file(filename)

# Read unstructured grid
u_grid = get_d3d_geometry(dflow_map)

# Write unstructured grid to VTK file
write_3d_vtu(u_grid, "tanana.vtu")
import numpy as np
import xarray as xr
import vtk

def get_water_height(dflow_map):
    if "mesh2d_waterdepth" in dflow_map.variables:
         water_height = dflow_map.variables["mesh2d_waterdepth"].values
    else:
        water_height = dflow_map.variables["waterdepth"].values
    return water_height

def get_bed_level(dflow_map):
    ## Bed level are at faces & constant in time
    if "mesh2d_bldepth" in dflow_map.variables:
        bed_level = dflow_map.variables["mesh2d_bldepth"].values
    elif "mesh2d_flowelem_bl" in dflow_map:
        bed_level = dflow_map.mesh2d_flowelem_bl.values
    else:
        bed_level = dflow_map.variables['FlowElem_bl'].values
    return bed_level

def get_layers(dflow_map):
    if "mesh2d_layer_sigma" in dflow_map.variables:
        sigma = dflow_map.variables["mesh2d_layer_sigma"].values[:]
        n_layers = len(sigma)
    elif "mesh2d_nLayers" in dflow_map.dims:
        sigma = dflow_map["mesh2d_nLayers"].values[:]
        n_layers = len(sigma)
        # TODO: THIS IS OVERWRITING THE SIGMA VARIABLE & impacting result
        # If no sigma variable, assume even layers
        sigma = np.linspace(0, 1, n_layers + 1)[:-1] - 1.0 + 1.0 / n_layers / 2.0
        # ds = np.linspace(0, 1, n_layers + 1)[1:]  # Calculated later    
    else:
        sigma = dflow_map["laydim"].values[:]
        n_layers = len(sigma)
    return sigma, n_layers

def read_d3d_file(filename):
    dflow_map = xr.open_dataset(filename)
    return dflow_map

def clean_connections(connections):
    """
    Remove NaNs, and shorten triangles to 3 entries
    """
    for cell_idx, cell in enumerate(connections):
        while np.isnan(cell[-1]):
            cell = cell[:-1]
        connections[cell_idx] = list(cell.astype(int))
    return connections

def get_face_connectivity(dflow_map):
    ## get connectivity & change count to start at 0
    if "mesh2d_face_nodes" in dflow_map.variables:
        face_connectivity = list(dflow_map.variables["mesh2d_face_nodes"][:].values - 1)
    else:
        face_connectivity = list(dflow_map.variables["NetElemNode"][:].values - 1)
    n_faces = len(face_connectivity)

    # Remove NaNs, and shorten triangles to 3 entries
    face_connectivity = clean_connections(face_connectivity)
    return face_connectivity, n_faces

def get_velocity(dflow_map, i_time):
    if "mesh2d_ucx" in dflow_map.variables:

        velocity = np.stack(
            (
                dflow_map.variables["mesh2d_ucx"].values[i_time],
                dflow_map.variables["mesh2d_ucy"].values[i_time],
                dflow_map.variables["mesh2d_ucz"].values[i_time],
            ),
            axis=-1,
        )
    else:
        velocity = np.stack(
            (
                dflow_map.variables["ucx"].values[i_time],
                dflow_map.variables["ucy"].values[i_time],
                dflow_map.variables["ucz"].values[i_time],
            ),
            axis=-1,
        )
    return velocity

def get_d3d_geometry(dflow_map):
    i_time = -1

    if "mesh2d_node_x" in dflow_map.variables:
        node_x = dflow_map.variables["mesh2d_node_x"].values[:]
        node_y = dflow_map.variables["mesh2d_node_y"].values[:]
    else:
        node_x = dflow_map.variables["NetNode_x"].values[:]
        node_y = dflow_map.variables["NetNode_y"].values[:]

    # Get layers
    sigma, n_layers = get_layers(dflow_map)

    # Get bed level
    bed_level = get_bed_level(dflow_map)

    # Water Height
    water_height = get_water_height(dflow_map)

    # Get connectivity
    face_connectivity, n_faces = get_face_connectivity(dflow_map)

    # Open the file
    u_grid = open_3d_vtu(
        node_x, node_y, face_connectivity, sigma, bed_level, water_height[i_time]
    )  ## water height determines mesh

    ## Add velocity data
    velocity = get_velocity(dflow_map, i_time)  
    u_grid = add_cell_vector(u_grid, "Velocity", velocity, n_faces, n_layers)

    ## Add water height data
    if "mesh2d_s1" in dflow_map.variables:
        surface_height = dflow_map.variables["mesh2d_s1"].values
    else:
        surface_height = dflow_map.variables["s1"].values    
    u_grid = add_face_scalar(
        u_grid, "Height", surface_height[i_time], n_faces, n_layers
    )

    # Add water depth data
    u_grid = add_face_scalar(u_grid, "Depth", water_height[i_time], n_faces, n_layers)

    ## Add bed level data
    u_grid = add_face_scalar(u_grid, "BedLevel", bed_level, n_faces, n_layers)

    return u_grid

def open_3d_vtu(node_x, node_y, face_connection, sigma, bed_level, water_height):
    """
    Creates a 3D unstructured grid (VTU) using given node coordinates, face connections,
    sigma layers, bed level, and water height data.

    Parameters:
    node_x: list or numpy.ndarray
        The x-coordinates of the nodes.
    node_y: list or numpy.ndarray
        The y-coordinates of the nodes.
    face_connection: (list of lists
        List of faces, where each face is a list of node indices.
    sigma: list or numpy.ndarray
        The sigma layers representing vertical discretization.
    bed_level: list or numpy.ndarray
        The bed level values corresponding to each face.
    water_height: list or numpy.ndarray
        The water height values corresponding to each face.

    Returns:
    u_grid: vtk.vtkUnstructuredGrid
        A 3D unstructured grid object containing the nodes and cells.
    """

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

    n_nodes = len(node_x)
    n_layers = len(sigma)
    ds = np.zeros(n_layers)
    ds[0] = (sigma[0] + 1.0) * 2.0  ## sigma is negative
    for i_layer in range(1, n_layers):
        ds[i_layer] = (sigma[i_layer] + 1.0) * 2.0 - ds[i_layer - 1]

    ## calculate node z
    ## scatter z to nodes and average
    node_z0 = np.zeros(n_nodes, dtype=float)  ## bottom z
    node_z1 = np.zeros(n_nodes, dtype=float)  ## top z
    n_face = np.zeros(n_nodes, dtype=float)
    for iface, face in enumerate(face_connection):
        for inode in face:
            node_z0[inode] += bed_level[iface]
            node_z1[inode] += bed_level[iface] + water_height[iface]
            n_face[inode] += 1

    ## do average
    for inode in range(n_nodes):
        node_z0[inode] = node_z0[inode] / n_face[inode]
        node_z1[inode] = node_z1[inode] / n_face[inode]

    # calculate points and store them

    ## layer 0
    for idx, x in enumerate(node_x):
        points.InsertNextPoint([x, node_y[idx], node_z0[idx]])

    for i_layer in range(1, n_layers):
        node_z = ds[i_layer - 1] * (node_z1[:] - node_z0[:]) + node_z0[:]
        for idx, x in enumerate(node_x):
            points.InsertNextPoint([x, node_y[idx], node_z[idx]])

    ## layer N
    for idx, x in enumerate(node_x):
        points.InsertNextPoint([x, node_y[idx], node_z1[idx]])

    ugrid.SetPoints(points)

    ## make cells
    vtk_element_type = []
    for i_layer in range(n_layers):
        for _, nodes in enumerate(face_connection):
            n_poly = len(nodes)
            if n_poly == 4:
                poly = vtk.vtkHexahedron()
                vtk_element_type.append(vtk.VTK_HEXAHEDRON)
            elif n_poly == 3:
                poly = vtk.vtkWedge()
                vtk_element_type.append(vtk.VTK_WEDGE)

            for i_vertex in range(n_poly):
                poly.GetPointIds().SetId(i_vertex, nodes[i_vertex] + n_nodes * i_layer)
                poly.GetPointIds().SetId(
                    i_vertex + n_poly, nodes[i_vertex] + n_nodes * (i_layer + 1)
                )
            cell_array.InsertNextCell(poly)

    ugrid.SetCells(vtk_element_type, cell_array)

    return ugrid

def add_cell_vector(u_grid, array_name, cell_data, n_faces, n_layers):
    """
    Adds a vector field to the cell data of a VTK unstructured grid.

    This function creates a new vector array with the specified name and
    adds it to the cell data of the given VTK unstructured grid. The vector
    data is provided for each face and layer in the grid.

    Parameters
    ----------
    u_grid : vtk.vtkUnstructuredGrid
        The unstructured grid to which the vector data will be added.
    array_name : string
        The name of the array to be added to the cell data.
    cell_data : numpy.ndarray
        The vector data to be added. This should be a 3D array of shape
        (n_faces, n_layers, 3), where each entry represents the vector
        components (x, y, z) for a face at a specific layer.
    n_faces : int
        The number of faces in the grid.
    n_layers : int
        The number of layers in the grid.

    Returns
    -------
    u_grid: vtk.vtkUnstructuredGrid
        The unstructured grid with the added vector data in the cell data.
    """
    data_array = vtk.vtkFloatArray()
    data_array.SetName(array_name)
    data_array.SetNumberOfComponents(3)
    for i_layer in range(n_layers):
        for i_face in range(n_faces):
            data_entry = cell_data[i_face, i_layer, :]
            data_array.InsertNextTuple(data_entry)

    u_grid.GetCellData().AddArray(data_array)

    return u_grid

def add_face_scalar(ugrid, array_name, cell_data, n_faces, n_layers):
    """
    Adds a scalar field to the cell data of a VTK unstructured grid,
    repeating the scalar values for each layer.

    Parameters
    ----------
    ugrid : vtk.vtkUnstructuredGrid
        The unstructured grid to which the scalar data will be added.
    array_name : string
        The name of the array to be added to the cell data.
    cell_data : numpy.ndarray
        The scalar data to be added. This should be a 1D array of shape (n_faces),
        where each entry represents the scalar value for a face.
    n_faces : int
        The number of faces in the grid.
    n_layers : int
        The number of layers in the grid.

    Returns
    -------
    ugrid : vtk.vtkUnstructuredGrid
        The unstructured grid with the added scalar data in the cell data.
    """
    data_array = vtk.vtkFloatArray()
    data_array.SetName(array_name)
    data_array.SetNumberOfComponents(1)
    for i_face in range(n_faces):
        data_entry = cell_data[i_face]
        # Repeat for all layers
        for _ in range(n_layers):
            data_array.InsertNextValue(data_entry)

    ugrid.GetCellData().AddArray(data_array)

    return ugrid

def write_3d_vtu(u_grid, filename):
    """
    Writes a VTK unstructured grid to a VTU file.

    This function takes a VTK unstructured grid and writes it to a specified
    file in VTU (VTK Unstructured Grid) format using the VTK XML writer.

    Parameters
    ----------
    u_grid : vtk.vtkUnstructuredGrid
        The unstructured grid to be written to the file.
    filename : string
        The name of the file to write the unstructured grid to. This should
        include the file extension (e.g., 'output.vtu').

    Returns
    -------
    None
    """

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

    writer.SetInputData(u_grid)
    writer.Write()

def write_2d_vtu(filename, vertices, connections):
    """
    Writes a 2D unstructured grid to a VTU file.

    This function takes a list of vertices and a list of face connections,
    constructs a VTK unstructured grid, and writes it to a specified file
    in VTU (VTK Unstructured Grid) format using the VTK XML writer.

    Parameters
    ----------
    filename : string
        The name of the file to write the unstructured grid to. This should
        include the file extension (e.g., 'output.vtu').
    vert_list : list of tuples
        A list of vertex coordinates, where each vertex is represented by a tuple
        of coordinates (x, y, z).
    conn_list : list of lists
        A list of faces, where each face is represented by a list of vertex indices.

    Returns
    -------
    None
    """

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

    for p in vertices:
        points.InsertNextPoint(p)

    for cell in connections:
        n_poly = len(cell)
        poly = vtk.vtkPolygon()
        poly.GetPointIds().SetNumberOfIds(n_poly)

        for i_vert, vert in enumerate(cell):
            poly.GetPointIds().SetId(i_vert, vert)
        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()

image