polyline_mapslice for vectors

[SCLaan]

The current function polyline_mapslice is used to take a vertical slice through the *_map.nc along a user-defined transect. However, for plotting velocity vectors, but possibly also transport variables, a projection to the transect is required. The below code includes some additional lines (the last ~25) to project the x- and y-component of velocity to a perpendicular and parallel component.

def polyline_mapslice_vectors(uds:xu.UgridDataset, line_array:np.array, calcdist_fromlatlon:bool = None) -> xu.UgridDataset:
    """
    Slice trough mapdata, combine: intersect_edges_withsort, calculation of distances and conversion to ugrid dataset.

    Parameters
    ----------
    uds : xu.UgridDataset
        DESCRIPTION.
    line_array : np.array
        DESCRIPTION.
    calcdist_fromlatlon : bool, optional
        DESCRIPTION. The default is None.

    Raises
    ------
    Exception
        DESCRIPTION.

    Returns
    -------
    xr_crs_ugrid : xu.UgridDataset
        DESCRIPTION.

    """

    #compute intersection coordinates of crossings between edges and faces and their respective indices
    edges = np.stack([line_array[:-1],line_array[1:]],axis=1)
    edge_index, face_index, intersections = dfmt.intersect_edges_withsort(uds=uds, edges=edges)
    if len(edge_index) == 0:
        raise ValueError('polyline does not cross mapdata')

    #auto determine if cartesian/sperical distance should be computed
    if calcdist_fromlatlon is None:
        if hasattr(uds,'projected_coordinate_system'):
            calcdist_fromlatlon = False
        elif hasattr(uds,'wgs84'):
            calcdist_fromlatlon = True
        else:
            raise KeyError('To auto determine calcdist_fromlatlon, a variable "projected_coordinate_system" or "wgs84" is required, please provide calcdist_fromlatlon=True/False yourself.')
    if calcdist_fromlatlon:
        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[:,1,0], edges[:,0,1], 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]

    #derive vertices from cross section (distance from first point)
    xr_crs_ugrid = dfmt.get_xzcoords_onintersection(uds=uds, face_index=face_index, crs_dist_starts=crs_dist_starts, crs_dist_stops=crs_dist_stops)

    ##### CODE BELOW IS IN ADDITION TO ORIGINAL DFM_TOOLS FUNCTION #####
    # Get the distance on the transect
    v_grid = xr_crs_ugrid.ugrid.grid.to_dataset()
    node_1 = v_grid['mesh2d_node_x'][v_grid['mesh2d_face_nodes'].isel(mesh2d_nMax_face_nodes=0)]

    # Calculate angles with the transect
    dist_xy        = [edges[:,0,0]-edges[:,1,0],edges[:,0,1]-edges[:,1,1]]
    angle_segments = np.arctan2(dist_xy[1],dist_xy[0]) * -1

    # Make a matrix with the corresponding angle for each cell on the transect
    angles_all_points = np.zeros(xr_crs_ugrid['mesh2d_nFaces'].shape)
    for i in range(len(angle_segments)):
        mask = (node_1 >= edge_len_cum0[i])
        angles_all_points[mask] = angle_segments[i]

    # Calculate the velocity parallel to the transect (perpendicular not used)
    v_par = xr_crs_ugrid['mesh2d_ucx'] * np.cos(angles_all_points) * -1 + xr_crs_ugrid['mesh2d_ucy'] * np.sin(angles_all_points)
    v_per = xr_crs_ugrid['mesh2d_ucx'] * np.sin(angles_all_points) + xr_crs_ugrid['mesh2d_ucy'] * np.cos(angles_all_points)

    # Overwrite x- and y-components with parallel and upward components (or: add new variables in final function)
    xr_crs_ugrid['mesh2d_ucx']   = v_par
    xr_crs_ugrid['mesh2d_ucy']   = xr_crs_ugrid['mesh2d_ucz']
    # Use parallel velocity as magnitude, because z-velocities are very small (maybe change this in final function)
    xr_crs_ugrid['mesh2d_ucmag'] = np.abs(v_par)
    #xr_crs_ugrid['mesh2d_ucmag'] = np.sqrt(xr_crs_ugrid['mesh2d_ucx']**2+xr_crs_ugrid['mesh2d_ucy']**2)

    return xr_crs_ugrid