Deltares / xugrid

Xarray and unstructured grids
https://deltares.github.io/xugrid/
MIT License
58 stars 8 forks source link

Add function to compute length of intersection of lines with 2D topology #159

Open Huite opened 11 months ago

Huite commented 11 months ago

For example, to compute conductances in a MODFLOW model, it's very useful to get the total length of a river within a cell:

def length_of_intersection(
    gdf: "geopandas.GeoDataframe",  # type: ignore # noqa
    like: Union["xugrid.Ugrid2d", "xugrid.UgridDataArray", "xugrid.UgridDataset"],
):
    """
    Compute (total) lenght of line intersection with the faces of a Ugrid2d mesh.

    Parameters
    ----------
    gdf: geopandas.GeoDataFrame
        Lines to be burned into the grid.
    like: UgridDataArray, UgridDataset, or Ugrid2d
        Grid to burn the vector data into.
    Returns
    -------
    length: UgridDataArray
    """
    import geopandas as gpd
    import pandas as pd

    if not isinstance(gdf, gpd.GeoDataFrame):
        raise TypeError(f"gdf must be GeoDataFrame, received: {type(like).__name__}")
    if isinstance(like, (xugrid.UgridDataArray, xugrid.UgridDataset)):
        like = like.ugrid.grid
    if not isinstance(like, xugrid.Ugrid2d):
        raise TypeError(
            "Like must be Ugrid2d, UgridDataArray, or UgridDataset;"
            f"received: {type(like).__name__}"
        )
    geometry_id = shapely.get_type_id(gdf.geometry)
    allowed_types = (LINESTRING, LINEARRING)
    if not np.isin(geometry_id, allowed_types).all():
        received = ", ".join(
            [GEOM_NAMES[geom_id] for geom_id in np.unique(geometry_id)]
        )
        raise TypeError(
            "GeoDataFrame contains unsupported geometry types. Can only "
            "compute intersection length for LineString and LinearRing "
            f"geometries. Received: {received}"
        )

    lines = gdf.geometry
    xy, index = shapely.get_coordinates(lines, return_index=True)
    # From the coordinates and the index, create the (n_edge, 2, 2) shape array
    # containing the edge_coordinates.
    linear_index = np.arange(index.size)
    segments = np.column_stack([linear_index[:-1], linear_index[1:]])
    # Only connections with vertices with the same index are valid.
    valid = np.diff(index) == 0
    segments = segments[valid]
    edges = xy[segments]
    # Now query the grid for these edges.
    _, face_index, intersections = like.intersect_edges(edges)
    # Compute the length
    length = np.linalg.norm(intersections[:, 1] - intersections[:, 0], axis=1)
    # Find the associated values.
    output = np.full(like.n_face, np.nan)
    length_series = pd.DataFrame({"face_index": face_index, "length": length}).groupby("face_index").sum()
    output[length_series.index] = length_series["length"]
    return xugrid.UgridDataArray(
        obj=xr.DataArray(output, dims=[like.face_dimension], name="length"),
        grid=like,
    )

However, this function has a lot of overlap with _burn_lines, ideally we can do the least amount of duplication. We could also add the length of the intersection as a coordinate, but it's likely a bad idea (it will contain a lot of NaNs).

Huite commented 11 months ago

From the vector conversion examples:

lines = gpd.GeoDataFrame(geometry=provinces.exterior)
length = xu.length_of_intersection(lines, uda)

fig, ax = plt.subplots()
length.ugrid.plot(ax=ax)
length.ugrid.plot.line(ax=ax, edgecolor="black", linewidth=0.5)
provinces.plot(ax=ax, edgecolor="red", facecolor="none", linewidth=1.5)
ax.set_xlim(xmin, xmax)
ax.set_ylim(ymin, ymax)

image