Open ssolson opened 8 months 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()
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.