polyline mapslice for vectors (thalweg)

[veenstrajelmer]

Two initiatives to plot vectors in a cross section plot were undertaken:

• thalweg_mapslice by MG
• EPD-support-tool contains polyline_mapslice_vectors(), which is an extended version of dfmt.polyline_mapslice(). This code is also documented in polyline_mapslice for vectors #675
• it was also requested in xugrid: Investigate a cross section plot xugrid#59 (comment)

Good to align the initiatives in dfm_tools, or preferably in xugrid.

[mgeraeds]

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.