Fix missing tiles in corners for fast gridding by using off-boundary pixels

geowrangler/gridding_utils/polygon_fill.py

@@ -11,16 +11,18 @@
 import geopandas as gpd
 import polars as pl
 
-# %% ../../notebooks/15_polygon_fill.ipynb 11
+# %% ../../notebooks/15_polygon_fill.ipynb 12
 def voxel_traversal_2d(
     start_vertex: Tuple[int, int],
     end_vertex: Tuple[int, int],
     debug: bool = False,  # if true, prints diagnostic info for the algorithm
-) -> List[Tuple[int, int]]:
+) -> Dict[str, List[Tuple[int, int]]]:
     """
     Returns all pixels between two points as inspired by Amanatides & Woo's “A Fast Voxel Traversal Algorithm For Ray Tracing”
 
     Implementation adapted from https://www.redblobgames.com/grids/line-drawing/ in the supercover lines section
+
+    This also returns the off-diagonal pixels that can be useful for correcting errors at the corners of polygons during polygon fill
     """
 
     # Setup initial conditions
@@ -30,26 +32,32 @@ def voxel_traversal_2d(
     direction_x = 1 if x2 > x1 else -1
     direction_y = 1 if y2 > y1 else -1
 
+    result = {"line_pixels": [], "off_diagonal_pixels": []}
+
     # Single point
     if (x1 == x2) and (y1 == y2):
         pixels = [(x1, y1)]
-        return pixels
+        result["line_pixels"].extend(pixels)
+        return result
 
     # Vertical line
     elif x1 == x2:
         pixels = [(x1, y) for y in range(y1, y2 + direction_y, direction_y)]
-        return pixels
+        result["line_pixels"].extend(pixels)
+        return result
 
     # Horizontal line
     elif y1 == y2:
         pixels = [(x, y1) for x in range(x1, x2 + direction_x, direction_x)]
-        return pixels
+        result["line_pixels"].extend(pixels)
+        return result
 
     dy = y2 - y1
     dx = x2 - x1
 
     pixel_x, pixel_y = x1, y1
     pixels = [(pixel_x, pixel_y)]
+    off_diagonal_pixels = []
 
     is_finished = False
 
@@ -73,6 +81,10 @@ def voxel_traversal_2d(
 
         decision = (1 + 2 * ix) * ny - (1 + 2 * iy) * nx
         if decision == 0:
+
+            off_diagonal_pixels.append((pixel_x + direction_x, pixel_y))
+            off_diagonal_pixels.append((pixel_x, pixel_y + direction_y))
+
             # diagonal step
             pixel_x += direction_x
             pixel_y += direction_y
@@ -106,9 +118,10 @@ def voxel_traversal_2d(
         if is_x_finished and is_y_finished:
             break
 
-    return pixels
+    result = {"line_pixels": pixels, "off_diagonal_pixels": off_diagonal_pixels}
+    return result
 
-# %% ../../notebooks/15_polygon_fill.ipynb 15
+# %% ../../notebooks/15_polygon_fill.ipynb 16
 def interpolate_x(
     start_vertex: Tuple[int, int],
     end_vertex: Tuple[int, int],
@@ -125,7 +138,7 @@ def interpolate_x(
     interpolated_x = x1 + (y - y1) * inverse_slope
     return interpolated_x
 
-# %% ../../notebooks/15_polygon_fill.ipynb 16
+# %% ../../notebooks/15_polygon_fill.ipynb 17
 def scanline_fill(
     vertices: List[
         Tuple[int, int]
@@ -182,38 +195,50 @@ def scanline_fill(
 
     return filled_pixels
 
-# %% ../../notebooks/15_polygon_fill.ipynb 20
+# %% ../../notebooks/15_polygon_fill.ipynb 21
 def voxel_traversal_scanline_fill(
     vertices_df: Union[
         pd.DataFrame, pl.DataFrame
     ],  # dataframe with x_col and y_col for the polygon vertices
     x_col: str = "x",
     y_col: str = "y",
     debug: bool = False,  # if true, prints diagnostic info for both voxel traversal and scanline fill algorithms
-) -> Set[Tuple[int, int]]:
+) -> Dict[str, Set[Tuple[int, int]]]:
     """
     Returns pixels that intersect a polygon.
 
     This uses voxel traversal to fill the boundary, and scanline fill for the interior. All coordinates are assumed to be integers.
+
+    This also returns the off-boundary pixels that can be useful for correcting errors at the corners of polygons during polygon fill
     """
 
     vertices = list(zip(vertices_df[x_col].to_list(), vertices_df[y_col].to_list()))
     offset_vertices = vertices[1:] + vertices[:1]
 
     polygon_pixels = set()
+    off_boundary_pixels = set()
 
     for start_vertex, end_vertex in zip(vertices, offset_vertices):
-        polygon_pixels.update(voxel_traversal_2d(start_vertex, end_vertex, debug))
+        voxel_traversal_results = voxel_traversal_2d(start_vertex, end_vertex, debug)
+        polygon_pixels.update(voxel_traversal_results["line_pixels"])
+        off_boundary_pixels.update(voxel_traversal_results["off_diagonal_pixels"])
 
     polygon_pixels.update(scanline_fill(vertices, debug))
 
-    return polygon_pixels
+    # removing off boundary tiles that are actually in the interior
+    off_boundary_pixels = off_boundary_pixels - polygon_pixels
+
+    result = {
+        "polygon_pixels": polygon_pixels,
+        "off_boundary_pixels": off_boundary_pixels,
+    }
+    return result
 
-# %% ../../notebooks/15_polygon_fill.ipynb 27
+# %% ../../notebooks/15_polygon_fill.ipynb 28
 SUBPOLYGON_ID_COL = "__subpolygon_id__"
 PIXEL_DTYPE = pl.Int32
 
-# %% ../../notebooks/15_polygon_fill.ipynb 28
+# %% ../../notebooks/15_polygon_fill.ipynb 29
 def polygons_to_vertices(
     polys_gdf: gpd.GeoDataFrame,
     unique_id_col: Optional[
@@ -251,13 +276,13 @@ def polygons_to_vertices(
 
     return vertices_df
 
-# %% ../../notebooks/15_polygon_fill.ipynb 32
+# %% ../../notebooks/15_polygon_fill.ipynb 33
 def fast_polygon_fill(
     vertices_df: pl.DataFrame,  # integer vertices of all polygons in the AOI
     unique_id_col: Optional[
         str
     ] = None,  # the ids under this column will be preserved in the output tiles
-) -> pl.DataFrame:
+) -> Dict[str, pl.DataFrame]:
 
     if unique_id_col is not None:
         id_cols = [SUBPOLYGON_ID_COL, unique_id_col]
@@ -276,6 +301,7 @@ def fast_polygon_fill(
     polygon_ids = vertices_df.select(id_cols).unique(maintain_order=True).rows()
 
     tiles_in_geom = set()
+    tiles_off_boundary = set()
     for polygon_id in polygon_ids:
         subpolygon_id, unique_id = polygon_id
         filter_expr = (pl.col(SUBPOLYGON_ID_COL) == subpolygon_id) & (
@@ -284,14 +310,21 @@ def fast_polygon_fill(
         poly_vertices = vertices_df.filter(filter_expr)
 
         poly_vertices = poly_vertices.unique(maintain_order=True)
-        _tiles_in_geom = voxel_traversal_scanline_fill(
+        voxel_traversal_results = voxel_traversal_scanline_fill(
             poly_vertices, x_col="x", y_col="y"
         )
+        _tiles_in_geom = voxel_traversal_results["polygon_pixels"]
+        _tiles_off_boundary = voxel_traversal_results["off_boundary_pixels"]
 
         if has_unique_id_col:
             _tiles_in_geom = [(x, y, unique_id) for (x, y) in _tiles_in_geom]
+            _tiles_off_boundary = [(x, y, unique_id) for (x, y) in _tiles_off_boundary]
 
         tiles_in_geom.update(_tiles_in_geom)
+        tiles_off_boundary.update(_tiles_off_boundary)
+
+    # removing off boundary tiles that are actually in the interior
+    tiles_off_boundary = tiles_off_boundary - tiles_in_geom
 
     schema = {"x": PIXEL_DTYPE, "y": PIXEL_DTYPE}
     if has_unique_id_col:
@@ -303,4 +336,12 @@ def fast_polygon_fill(
         schema=schema,
     )
 
-    return tiles_in_geom
+    tiles_off_boundary = pl.from_records(
+        data=list(tiles_off_boundary),
+        orient="row",
+        schema=schema,
+    )
+
+    result = {"tiles_in_geom": tiles_in_geom, "tiles_off_boundary": tiles_off_boundary}
+
+    return result

geowrangler/grids.py

@@ -231,10 +231,25 @@ def generate_grid(
         vertices, boundary, northing_col="y", easting_col="x"
     )
 
-    tiles_in_geom = polygon_fill.fast_polygon_fill(vertices, unique_id_col)
-
+    polygon_fill_result = polygon_fill.fast_polygon_fill(vertices, unique_id_col)
+    tiles_in_geom = polygon_fill_result["tiles_in_geom"]
     bboxes = self._xy_to_bbox(tiles_in_geom, boundary, "x", "y")
 
+    # this is error correction on the polygon boundary (not the square boundary)
+    tiles_off_boundary = polygon_fill_result["tiles_off_boundary"]
+    if not tiles_off_boundary.is_empty():
+        off_boundary_bboxes = self._xy_to_bbox(tiles_off_boundary, boundary, "x", "y")
+        all_polygon_boundary = reprojected_gdf.boundary.union_all(method="unary")
+        intersects_boundary_bool = off_boundary_bboxes.intersects(all_polygon_boundary)
+
+        addtl_tiles_in_geom = tiles_off_boundary.filter(
+            pl.Series(intersects_boundary_bool)
+        )
+        addtl_boundary_bboxes = off_boundary_bboxes[intersects_boundary_bool].copy()
+
+        tiles_in_geom = pl.concat([tiles_in_geom, addtl_tiles_in_geom])
+        bboxes = pd.concat([bboxes, addtl_boundary_bboxes], ignore_index=True)
+
     column_order = ["x", "y"]
     if unique_id_col is not None:
         column_order += [unique_id_col]
@@ -664,7 +679,20 @@ def generate_grid(
     vertices = polygon_fill.polygons_to_vertices(aoi_gdf, unique_id_col)
     vertices = self._latlng_to_xy(vertices, lat_col="y", lng_col="x")
 
-    tiles_in_geom = polygon_fill.fast_polygon_fill(vertices, unique_id_col)
+    polygon_fill_result = polygon_fill.fast_polygon_fill(vertices, unique_id_col)
+    tiles_in_geom = polygon_fill_result["tiles_in_geom"]
+
+    # this is error correction on the polygon boundary (not the square boundary)
+    tiles_off_boundary = polygon_fill_result["tiles_off_boundary"]
+    if not tiles_off_boundary.is_empty():
+        off_boundary_bboxes = self._xy_to_bbox(tiles_off_boundary, "x", "y")
+        all_polygon_boundary = aoi_gdf.boundary.union_all(method="unary")
+        intersects_boundary_bool = off_boundary_bboxes.intersects(all_polygon_boundary)
+        addtl_tiles_in_geom = tiles_off_boundary.filter(
+            pl.Series(intersects_boundary_bool)
+        )
+
+        tiles_in_geom = pl.concat([tiles_in_geom, addtl_tiles_in_geom])
 
     quadkey_expr = self._xyz_to_quadkey(
         pl.col("x"),
@@ -675,6 +703,10 @@ def generate_grid(
     if self.return_geometry:
         bboxes = self._xy_to_bbox(tiles_in_geom, "x", "y")
 
+        if not tiles_off_boundary.is_empty():
+            addtl_boundary_bboxes = off_boundary_bboxes[intersects_boundary_bool].copy()
+            bboxes = pd.concat([bboxes, addtl_boundary_bboxes], ignore_index=True)
+
     if not self.add_xyz_cols:
         tiles_in_geom = tiles_in_geom.drop(["x", "y"])
     else:

notebooks/00_grids.ipynb

@@ -311,7 +311,9 @@
     "This uses these optimizations to speed up grid generation:\n",
     "\n",
     "1. Vectorized Translation Functions: Functions that translate between lat,lon and x,y are written in polars.\n",
-    "2. Voxel Traversal and Scanline Fill Algorithms: Faster alternative to finding all pixels within a polygon without using polygon intersection operations. These are implemented in `polygon_fill.fast_polygon_fill()`"
+    "2. Voxel Traversal and Scanline Fill Algorithms: Faster alternative to finding all pixels within a polygon without using polygon intersection operations. These are implemented in `polygon_fill.fast_polygon_fill()`\n",
+    "\n",
+    "This also does error correction on the polygon boundary using off-boundary pixels. Read more in the [polygon fill module reference](https://geowrangler.thinkingmachin.es/polygon_fill.html)"
    ]
   },
   {
@@ -362,10 +364,23 @@
     "        vertices = self._remove_out_of_bounds_polygons(vertices, boundary)\n",
     "    vertices = self._northingeasting_to_xy(vertices, boundary, northing_col=\"y\", easting_col=\"x\")\n",
     "    \n",
-    "    tiles_in_geom = polygon_fill.fast_polygon_fill(vertices, unique_id_col)\n",
-    "\n",
+    "    polygon_fill_result = polygon_fill.fast_polygon_fill(vertices, unique_id_col)\n",
+    "    tiles_in_geom = polygon_fill_result[\"tiles_in_geom\"]\n",
     "    bboxes = self._xy_to_bbox(tiles_in_geom, boundary, \"x\", \"y\")\n",
     "\n",
+    "    # this is error correction on the polygon boundary (not the square boundary) \n",
+    "    tiles_off_boundary = polygon_fill_result[\"tiles_off_boundary\"]\n",
+    "    if not tiles_off_boundary.is_empty():\n",
+    "        off_boundary_bboxes = self._xy_to_bbox(tiles_off_boundary, boundary, \"x\", \"y\") \n",
+    "        all_polygon_boundary = reprojected_gdf.boundary.union_all(method=\"unary\")\n",
+    "        intersects_boundary_bool = off_boundary_bboxes.intersects(all_polygon_boundary)\n",
+    "    \n",
+    "        addtl_tiles_in_geom = tiles_off_boundary.filter(pl.Series(intersects_boundary_bool))\n",
+    "        addtl_boundary_bboxes = off_boundary_bboxes[intersects_boundary_bool].copy()\n",
+    "    \n",
+    "        tiles_in_geom = pl.concat([tiles_in_geom, addtl_tiles_in_geom])\n",
+    "        bboxes = pd.concat([bboxes, addtl_boundary_bboxes], ignore_index=True)\n",
+    "\n",
     "    column_order = [\"x\",\"y\"]\n",
     "    if unique_id_col is not None:\n",
     "        column_order += [unique_id_col]\n",
@@ -854,7 +869,9 @@
     "This uses these optimizations to speed up grid generation:\n",
     "\n",
     "1. Vectorized Translation Functions: Functions that translate between lat,lon and x,y are written in polars.\n",
-    "2. Voxel Traversal and Scanline Fill Algorithms: Faster alternative to finding all pixels within a polygon without using polygon intersection operations. These are implemented in `polygon_fill.fast_polygon_fill()`"
+    "2. Voxel Traversal and Scanline Fill Algorithms: Faster alternative to finding all pixels within a polygon without using polygon intersection operations. These are implemented in `polygon_fill.fast_polygon_fill()`\n",
+    "\n",
+    "This also does error correction on the polygon boundary using off-boundary pixels. Read more in the [polygon fill module reference](https://geowrangler.thinkingmachin.es/polygon_fill.html)"
    ]
   },
   {
@@ -902,7 +919,18 @@
     "    vertices = polygon_fill.polygons_to_vertices(aoi_gdf, unique_id_col)\n",
     "    vertices = self._latlng_to_xy(vertices, lat_col=\"y\", lng_col=\"x\")\n",
     "\n",
-    "    tiles_in_geom = polygon_fill.fast_polygon_fill(vertices, unique_id_col)\n",
+    "    polygon_fill_result = polygon_fill.fast_polygon_fill(vertices, unique_id_col)\n",
+    "    tiles_in_geom = polygon_fill_result[\"tiles_in_geom\"]\n",
+    "\n",
+    "    # this is error correction on the polygon boundary (not the square boundary)\n",
+    "    tiles_off_boundary = polygon_fill_result[\"tiles_off_boundary\"]\n",
+    "    if not tiles_off_boundary.is_empty():\n",
+    "        off_boundary_bboxes = self._xy_to_bbox(tiles_off_boundary, \"x\", \"y\")\n",
+    "        all_polygon_boundary = aoi_gdf.boundary.union_all(method=\"unary\")\n",
+    "        intersects_boundary_bool = off_boundary_bboxes.intersects(all_polygon_boundary)\n",
+    "        addtl_tiles_in_geom = tiles_off_boundary.filter(pl.Series(intersects_boundary_bool))\n",
+    "    \n",
+    "        tiles_in_geom = pl.concat([tiles_in_geom, addtl_tiles_in_geom])\n",
     "\n",
     "    quadkey_expr = self._xyz_to_quadkey(\n",
     "        pl.col(\"x\"),\n",
@@ -912,6 +940,10 @@
     "\n",
     "    if self.return_geometry:\n",
     "        bboxes = self._xy_to_bbox(tiles_in_geom, \"x\", \"y\")\n",
+    "\n",
+    "        if not tiles_off_boundary.is_empty():\n",
+    "            addtl_boundary_bboxes = off_boundary_bboxes[intersects_boundary_bool].copy()\n",
+    "            bboxes = pd.concat([bboxes, addtl_boundary_bboxes], ignore_index=True)\n",
     "    \n",
     "    if not self.add_xyz_cols:\n",
     "        tiles_in_geom = tiles_in_geom.drop([\"x\",\"y\"])\n",