More efficient burn_polygons. Expand user guide, especially on all_to…

…uched option.

docs/changelog.rst

@@ -9,6 +9,15 @@ The format is based on `Keep a Changelog`_, and this project adheres to
 Unreleased
 ----------
 
+Added
+~~~~~
+
+- Added :func:`xugrid.burn_vector_geometries` to burn vector geometries in the
+  form of geopandas GeoDataFrames into a Ugrid2d topology.
+- Added :func:`xugrid.polygonize` to create vector polygons for all connected
+  regions of a Ugrid2d topology sharing a common value. The result is a
+  geopandas GeoDataFrame.
+
 [0.6.2] 2023-07-26
 ------------------
 

examples/vector_conversion.py

@@ -13,6 +13,7 @@
 # %%
 
 import geopandas as gpd
+import matplotlib.pyplot as plt
 
 import xugrid as xu
 
@@ -32,7 +33,7 @@
 # ``.to_geodataframe``, a shapely Polygon is created for every face (cell).
 
 gdf = uda.ugrid.to_geodataframe()
-gdf
+print(gdf)
 
 # %%
 # We see that a GeoDataFrame with 5248 rows is created: one row for each face.
@@ -73,7 +74,7 @@
 #
 # In this example, we mark the faces that are covered by a certain province.
 
-provinces = xu.data.provinces_nl.to_crs(28992)
+provinces = xu.data.provinces_nl().to_crs(28992)
 provinces["value"] = range(len(provinces))
 burned = xu.burn_vector_geometry(provinces, uda, column="value")
 burned.ugrid.plot()
@@ -83,7 +84,49 @@
 
 utrecht = provinces[provinces["name"] == "Utrecht"]
 burned = xu.burn_vector_geometry(utrecht, uda)
-burned.ugrid.plot()
+xmin, ymin, xmax, ymax = utrecht.buffer(10_000).total_bounds
+
+fig, ax = plt.subplots()
+burned.ugrid.plot(ax=ax)
+burned.ugrid.plot.line(ax=ax, edgecolor="black", linewidth=0.5)
+utrecht.plot(ax=ax, edgecolor="red", facecolor="none", linewidth=1.5)
+ax.set_xlim(xmin, xmax)
+ax.set_ylim(ymin, ymax)
+
+# %%
+# By default, ``burn_vector_geometry`` will only include grid faces whose
+# centroid are located in a polygon. We can also mark all intersected faces
+# by setting ``all_touched=True``:
+
+burned = xu.burn_vector_geometry(utrecht, uda, all_touched=True)
+
+fig, ax = plt.subplots()
+burned.ugrid.plot(ax=ax)
+burned.ugrid.plot.line(ax=ax, edgecolor="black", linewidth=0.5)
+utrecht.plot(ax=ax, edgecolor="red", facecolor="none", linewidth=1.5)
+ax.set_xlim(xmin, xmax)
+ax.set_ylim(ymin, ymax)
+
+# %%
+# Note that ``all_touched=True`` is less suitable when differently valued
+# polygons are present that share borders. While the centroid of a face is
+# contained by only a single polygon, the area of the polygon may be located
+# in more than one polygon. In this case, the results of each polygon will
+# overwrite each other (in arbitrary order).
+
+by_centroid = xu.burn_vector_geometry(provinces, uda, column="value")
+by_touch = xu.burn_vector_geometry(provinces, uda, column="value", all_touched=True)
+
+fig, axes = plt.subplots(ncols=2, figsize=(10, 5))
+by_centroid.ugrid.plot(ax=axes[0], add_colorbar=False)
+by_touch.ugrid.plot(ax=axes[1], add_colorbar=False)
+
+for ax, title in zip(axes, ("centroid", "all touched")):
+    burned.ugrid.plot.line(ax=ax, edgecolor="black", linewidth=0.5)
+    utrecht.plot(ax=ax, edgecolor="red", facecolor="none", linewidth=1.5)
+    ax.set_xlim(xmin, xmax)
+    ax.set_ylim(ymin, ymax)
+    ax.set_title(title)
 
 # %%
 # This function can also be used to burn points or lines into the faces of an
@@ -94,7 +137,13 @@
 
 lines = gpd.GeoDataFrame(geometry=provinces.exterior)
 burned = xu.burn_vector_geometry(lines, uda)
-burned.ugrid.plot()
+
+fig, ax = plt.subplots()
+burned.ugrid.plot(ax=ax)
+burned.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)
 
 # %%
 # Polygonizing
@@ -111,12 +160,12 @@
 # %%
 # We see that the results consists of two large polygons, in which the
 # triangles of the triangular grid have been merged to form a single polygon,
-# and many smaller polygons, some of which correspond one-to-one to the
+# and many smaller polygons, some of which correspond one to one to the
 # triangles of the grid.
 #
 # .. note::
 #   The produced polygon edges will follow exactly the face boundaries. When
 #   the data consists of many unique values (e.g. unbinned elevation data), the
-#   result will essentially be one polygon per face. In such cases, it is much
-#   more efficient to use ``xugrid.UgridDataArray.to_geodataframe``, which
-#   directly converts every face to a polygon.
+#   result will essentially be one polygon per face. In such cases, it is more
+#   efficient to use ``xugrid.UgridDataArray.to_geodataframe``, which directly
+#   converts every face to a polygon.

xugrid/ugrid/burn.py

@@ -1,29 +1,102 @@
 from typing import List, Union
 
+import numba as nb
 import numpy as np
 import xarray as xr
+from numba_celltree.constants import Point, Triangle
+from numba_celltree.geometry_utils import (
+    as_point,
+    as_triangle,
+    cross_product,
+    to_vector,
+)
 
 import xugrid
-from xugrid.constants import BoolArray, FloatArray, MissingOptionalModule
+from xugrid.constants import FloatArray, IntArray, MissingOptionalModule
 
 try:
     import shapely
 except ImportError:
     shapely = MissingOptionalModule("shapely")
 
 
+@nb.njit(inline="always")
+def point_in_triangle(p: Point, triangle: Triangle) -> bool:
+    """Unrolled half-plane check."""
+    # TODO: move this in numba_celltree instead?
+    a = cross_product(to_vector(triangle.a, triangle.b), to_vector(triangle.a, p)) > 0
+    b = cross_product(to_vector(triangle.b, triangle.c), to_vector(triangle.b, p)) > 0
+    c = cross_product(to_vector(triangle.c, triangle.a), to_vector(triangle.c, p)) > 0
+    return (a == b) and (b == c)
+
+
+@nb.njit(inline="always", parallel=True, cache=True)
+def points_in_triangles(
+    points: FloatArray,
+    face_indices: IntArray,
+    faces: IntArray,
+    vertices: FloatArray,
+):
+    # TODO: move this in numba_celltree instead?
+    n_points = len(points)
+    inside = np.empty(n_points, dtype=np.bool_)
+    for i in nb.prange(n_points):
+        face_index = face_indices[i]
+        face = faces[face_index]
+        triangle = as_triangle(vertices, face)
+        point = as_point(points[i])
+        inside[i] = point_in_triangle(point, triangle)
+    return inside
+
+
 def _locate_polygon(
-    xy: FloatArray,
+    grid: "xu.Ugrid2d",  # type: ignore # noqa
     exterior: FloatArray,
     interiors: List[FloatArray],
-) -> BoolArray:
+    all_touched: bool,
+) -> IntArray:
+    """
+    This algorithm burns polygon vector geometries in a 2d topology by:
+
+    * Extracting the exterior and interiors (holes) coordinates from the
+      polygons.
+    * Breaking every polygon down into a triangles using an "earcut" algorithm.
+    * Searching the grid for these triangles.
+
+    Due to the use of the separating axes theorem, _locate_faces finds all
+    intersecting triangles, including those who only touch the edge.
+    To enable all_touched=False, we have to search the centroids of candidates
+    in the intersecting triangles.
+
+    Parameters
+    ----------
+    grid: Ugrid2d
+    exterior: FloatArray
+        Exterior of the polygon.
+    interiors: List[FloatArray]
+        Interior holes of the polygon.
+    all_touched: bool
+        Whether to include include cells whose centroid falls inside, of every
+        cell that is touched.
+    """
+
     import mapbox_earcut
 
     rings = np.cumsum([len(exterior)] + [len(interior) for interior in interiors])
     vertices = np.vstack([exterior] + interiors).astype(np.float64)
     triangles = mapbox_earcut.triangulate_float64(vertices, rings).reshape((-1, 3))
-    grid = xugrid.Ugrid2d(*vertices.T, -1, triangles)
-    return grid.locate_points(xy) != -1
+    triangle_indices, grid_indices = grid.celltree._locate_faces(vertices, triangles)
+    if all_touched:
+        return grid_indices
+    else:
+        centroids = grid.centroids[grid_indices]
+        inside = points_in_triangles(
+            points=centroids,
+            face_indices=triangle_indices,
+            faces=triangles,
+            vertices=vertices,
+        )
+        return grid_indices[inside]
 
 
 def _burn_polygons(
@@ -33,61 +106,20 @@ def _burn_polygons(
     all_touched: bool,
     output: FloatArray,
 ) -> None:
-    """
-    This algorithm burns polygon vector geometries in a 2d topology by:
-
-    * Extracting the exterior and interiors (holes) coordinates from the
-      polygons.
-    * Breaking every polygon down into a triangles using an "earcut" algorithm.
-    * Building a Ugrid2d from the triangles.
-
-    When all_touched=False:
-
-    * The grid created from the polygon is searched for every centroid of the
-      ``like`` grid.
-
-    When all_touched=True:
-
-    * The grid created from the polygon is searched for every node of the
-      ``like`` grid.
-    * If any of the associated nodes of a face is found in the polygon, the
-      value is burned into the face.
-
-    """
     exterior_coordinates = [
         shapely.get_coordinates(exterior) for exterior in polygons.exterior
     ]
     interior_coordinates = [
         [shapely.get_coordinates(p_interior) for p_interior in p_interiors]
         for p_interiors in polygons.interiors
     ]
+    to_burn = np.empty(like.n_face, dtype=bool)
 
-    if all_touched:
-        # Pre-allocate work arrays so we don't have re-allocate for every
-        # polygon.
-        to_burn2d = np.empty_like(like.face_node_connectivity, dtype=bool)
-        to_burn = np.empty(like.n_face, dtype=bool)
-        # These data are static:
-        mask = like.face_node_connectivity == like.fill_value
-        xy = like.node_coordinates
-        for exterior, interiors, value in zip(
-            exterior_coordinates, interior_coordinates, values
-        ):
-            location = _locate_polygon(xy, exterior, interiors)
-            # Equal to: to_burn2d = location[like.face_node_connectivity]
-            location.take(like.face_node_connectivity, out=to_burn2d)
-            to_burn2d[mask] = False
-            # Equal to: to_burn = to_burn2d.any(axis=1)
-            np.bitwise_or.reduce(to_burn2d, axis=1, out=to_burn)
-            output[to_burn] = value
-
-    else:
-        xy = like.centroids
-        for exterior, interiors, value in zip(
-            exterior_coordinates, interior_coordinates, values
-        ):
-            to_burn = _locate_polygon(xy, exterior, interiors)
-            output[to_burn] = value
+    for exterior, interiors, value in zip(
+        exterior_coordinates, interior_coordinates, values
+    ):
+        to_burn = _locate_polygon(like, exterior, interiors, all_touched)
+        output[to_burn] = value
 
     return