Open veenstrajelmer opened 7 months ago
Hi Jelmer,
I started a first initiative for my own post-processing for this. My extended version of the dfmt.polyline_mapslice
is as follows:
def thalweg_mapslice(uds:xu.UgridDataset, cross_section:np.array, reference_location=None) -> (xu.UgridDataset, xu.UgridDataset):
"""Function to slice through mapdata based on a cross-section, and transform the relevant dataset variables to an along-channel and cross-channel component.
:param uds xu.UgridDataset: Dataset to-be-sliced. Should only contain a single timestep.
:param cross_section np.array: Cross-section polyline
:returns uds_sliced xu.UgridDataset: Dataset in cross-sectional (transformed) reference frame, main dimension is nCrossedFaces
:returns xr_crs_ugrid xu.UgridDataset: Cross-sectional dataframe that is plottable (mesh2d_nFaces contain x,z-information)
"""
# > Do imports
from dfm_tools.xugrid_helpers import get_vertical_dimensions
from dfm_tools.xarray_helpers import Dataset_varswithdim
from dfm_tools.get_nc import intersect_edges_withsort
from scipy.spatial import cKDTree
# > Compute intersection coordinates of crossings between edges and faces and their respective indices
dimn_layer, dimn_interfaces = get_vertical_dimensions(uds)
gridname = uds.grid.name
edges = np.stack([cross_section[:-1], cross_section[1:]], axis=1)
edge_index, face_index, intersections = intersect_edges_withsort(uds=uds, edges=edges)
if len(edge_index) == 0:
raise ValueError('Selected cross-section does not cross mapdata')
if uds.grids[0].is_geographic:
calc_dist = dfmt.calc_dist_haversine
else:
calc_dist = dfmt.calc_dist_pythagoras
# > Compute pyt/haversine start/stop distances for all intersections
edge_len = calc_dist(edges[:,0,0], edges[:,0,1], edges[:,1,0], edges[:,1,1])
edge_len_cum = np.cumsum(edge_len)
edge_len_cum0 = np.concatenate([[0], edge_len_cum[:-1]])
crs_dist_starts = calc_dist(edges[edge_index,0,0], intersections[:,0,0], edges[edge_index,0,1], intersections[:,0,1]) + edge_len_cum0[edge_index]
crs_dist_stops = calc_dist(edges[edge_index,0,0], intersections[:,1,0], edges[edge_index,0,1], intersections[:,1,1]) + edge_len_cum0[edge_index]
# > Drop all variables that do not contain a face dimension, then select only all sliced faceidx
xu_facedim = uds.grid.face_dimension
face_index_xr = xr.DataArray(face_index,dims=(f'{gridname}_nCrossedFaces'))
uds = Dataset_varswithdim(uds, dimname=xu_facedim) #TODO: is there an xugrid alternative?
uds_sel = uds.sel({xu_facedim:face_index_xr})
# > Reshape the intersections to be able to calculate width of the cross-section through which the transport is calculated
uds_sel[f'{gridname}_dxc'] = xr.DataArray((crs_dist_stops-crs_dist_starts), dims={f'{gridname}_nCrossedFaces'}, name=f'{gridname}_dyc')
# > Calculate the angle of the cross-section relative to the x-axis
# > Do this based on the average velocity vector over depth (PCA-like)
uds_sel[f'{gridname}_alpha'] = np.rad2deg(np.arctan(uds_sel[f'{gridname}_ucy'].mean(dim='mesh2d_nLayers') / uds_sel[f'{gridname}_ucx'].mean(dim='mesh2d_nLayers')))
uds_sel[f'{gridname}_alpha'].assign_attrs({'standard_name': 'cross-section angle thalweg', 'unit': 'degrees'})
# > Calculate the normal vector relative to cross-section
# > BE AWARE: this function assumes alpha is in the direction of the flow.
uds_sel[f'{gridname}_ucs'], uds_sel[f'{gridname}_ucn'] = fixed2bearing(uds_sel['mesh2d_ucx'], uds_sel['mesh2d_ucy'], uds_sel['mesh2d_alpha'])
# > Generate the cumulative distance from the left-most point of the cross-section
uds_sel[f'{gridname}_distance'] = crs_dist_starts + uds_sel[f'{gridname}_dxc']/2
# > Calculate the bed level if not present in the dataset
if not f'{gridname}_flowelem_bl' in uds_sel:
uds_sel[f'{gridname}_flowelem_bl'] = uds_sel[f'{gridname}_flowelem_zcc'].min(dim=dimn_layer)
# > See if reference location is given and calculate a starting point distance
if reference_location:
if isinstance(reference_location, tuple):
# > Stack dataset faces
faces = np.stack([uds_sel[f"{gridname}_face_x"], uds_sel[f'{gridname}_face_y']], axis=1)
# > Make cKDTree based on the faces
tree = cKDTree(faces)
# > Find index of nearest face
dist, closest_idx = tree.query(reference_location)
# > If the closest index is not 0, that means that the reference location falls within the cross-section
# > We have to then subtract the reference distance to get a new distance axis
if closest_idx != 0:
# > Select closest reference distance
ref_dist = uds_sel[f"{gridname}_distance"][closest_idx].values
# > Update distance arrays by subtracting reference distance to previously calculated distances
uds_sel[f"{gridname}_distance"] = uds_sel[f"{gridname}_distance"] - ref_dist
crs_dist_starts = crs_dist_starts - ref_dist
crs_dist_stops = crs_dist_stops - ref_dist
else:
# > Update distance arrays by adding reference distance to previously calculated distances
uds_sel[f"{gridname}_distance"] = uds_sel[f"{gridname}_distance"] + dist
crs_dist_starts = crs_dist_starts + dist
crs_dist_stops = crs_dist_stops + dist
# > Take zvals_interface
if dimn_layer in uds_sel.dims: #3D model
nlay = uds.dims[dimn_layer]
zvals_interface_filled = uds_sel[f'{gridname}_flowelem_zw'].bfill(dim=dimn_interfaces) # Fill nan values (below bed) with equal values
zvals_interface = zvals_interface_filled.to_numpy().T #transpose to make in line with 2D sigma dataset
else: #2D model, no layers
nlay = 1
data_frommap_wl3_sel = uds_sel[f'{gridname}_s1'].to_numpy() #TODO: add escape for missing wl/bl vars
data_frommap_bl_sel = uds_sel[f'{gridname}_flowelem_bl'].to_numpy()
zvals_interface = np.linspace(data_frommap_bl_sel, data_frommap_wl3_sel, nlay+1)
# > Derive crs_verts
crs_dist_starts_matrix = np.repeat(crs_dist_starts[np.newaxis], nlay, axis=0)
crs_dist_stops_matrix = np.repeat(crs_dist_stops[np.newaxis], nlay, axis=0)
crs_verts_x_all = np.array([[crs_dist_starts_matrix.ravel(), crs_dist_stops_matrix.ravel(), crs_dist_stops_matrix.ravel(), crs_dist_starts_matrix.ravel()]]).T
crs_verts_z_all = np.ma.array([zvals_interface[1:,:].ravel(), zvals_interface[1:,:].ravel(), zvals_interface[:-1,:].ravel(), zvals_interface[:-1,:].ravel()]).T[:,:,np.newaxis]
crs_verts = np.ma.concatenate([crs_verts_x_all, crs_verts_z_all], axis=2)
# > Define grid
shape_crs_grid = crs_verts[:,:,0].shape
shape_crs_flat = crs_verts[:,:,0].ravel().shape
xr_crs_grid = xu.Ugrid2d(node_x=crs_verts[:,:,0].ravel(),
node_y=crs_verts[:,:,1].ravel(),
fill_value=-1,
face_node_connectivity=np.arange(shape_crs_flat[0]).reshape(shape_crs_grid),
)
# > Define dataset
if dimn_layer in uds_sel.dims:
crs_plotdata_clean = uds_sel.stack({xr_crs_grid.face_dimension:[dimn_layer,f'{gridname}_nCrossedFaces']},create_index=False)
else: #2D: still make sure xr_crs_grid.face_dimension is created, using stack since .rename() gives "UserWarning: rename 'ncrossed_faces' to 'mesh2d_nFaces' does not create an index anymore."
crs_plotdata_clean = uds_sel.stack({xr_crs_grid.face_dimension:[f'{gridname}_nCrossedFaces']},create_index=False)
# > Combine into xugrid
xr_crs_ugrid = xu.UgridDataset(crs_plotdata_clean, grids=[xr_crs_grid])
# > Drop variables that don't make sense anymore
xr_crs_ugrid = xr_crs_ugrid.drop(f'{gridname}_flowelem_bl')
return xr_crs_ugrid
Additionally, the fixed2bearing
function that I wrote that is referenced:
def fixed2bearing(east_velocity, north_velocity, principal_direction):
"""
Function to calculate the velocity components in a reference frame parallel to the principal direction of the
velocity cluster.
:param east_velocity: velocity in Eastern direction (float)
:param north_velocity: velocity in Northern direction (float)
:param principal_direction: principal direction in degrees in fixed reference frame North-East (float)
:return: coordinates in a reference frame relative to the given principal direction (tup)
"""
bearing_rad = np.radians(principal_direction)
x_velocity = np.cos(bearing_rad) * east_velocity + np.sin(bearing_rad) * north_velocity
y_velocity = -1 * np.sin(bearing_rad) * east_velocity + np.cos(bearing_rad) * north_velocity
return x_velocity, y_velocity
It works for what I wanted for now, but could definitely be made more modular. For example, one thing that is not ideal is that I'd still like to use uds_sel
after calculation of the along- and cross-channel velocity, because the xr_crs_ugrid
variable has an altered grid (for plotting). I need some of the grid-related variables (like the connectivity properties) for later use. Maybe something can be implemented in xugrid
for this?
Also, I called this thalweg_mapslice
because I assumed that the slice is in the direction of the flow, but it essentially just divides it up into an along-slice and cross-slice component.
Another way to really calculate along-channel and cross-channel components is by doing a principal component analysis (PCA) which can decompose the velocity based on the velocity signal, assuming that the signal is in the direction of the deepest point in the channel, but there are some things to take into account when doing this. This also doesn't work when there's not a clear thalweg or main direction.
Two initiatives to plot vectors in a cross section plot were undertaken:
thalweg_mapslice
by MGpolyline_mapslice_vectors()
, which is an extended version ofdfmt.polyline_mapslice()
. This code is also documented in https://github.com/Deltares/dfm_tools/issues/675Good to align the initiatives in dfm_tools, or preferably in xugrid.