implements h3 compaction

Author: alpha-beta-soup

raster2dggs/common.py

@@ -146,6 +146,7 @@ def get_parent_res(dggs: str, parent_res: Union[None, int], resolution: int) ->
 def address_boundary_issues(
     dggs: str,
     parent_groupby: Callable,
+    compaction: Callable,
     pq_input: tempfile.TemporaryDirectory,
     output: Path,
     resolution: int,
@@ -184,23 +185,25 @@ def address_boundary_issues(
     )
     LOGGER.debug("Aggregating cell values where conflicts exist")
 
-    with TqdmCallback(desc="Repartioning/aggregating"):
-        ddf = (
-            ddf.repartition(  # See "notes" on why divisions expects repetition of the last item https://docs.dask.org/en/stable/generated/dask.dataframe.DataFrame.repartition.html
-                divisions=(uniqueparents + [uniqueparents[-1]])
-            )
-            .map_partitions(
-                parent_groupby, resolution, kwargs["aggfunc"], kwargs["decimals"]
-            )
-            .to_parquet(
-                output,
-                overwrite=kwargs["overwrite"],
-                engine="pyarrow",
-                write_index=True,
-                append=False,
-                name_function=lambda i: f"{uniqueparents[i]}.parquet",
-                compression=kwargs["compression"],
-            )
+    with TqdmCallback(
+        desc=f"Repartioning/aggregating{'/compacting' if compaction else ''}"
+    ):
+        ddf = ddf.repartition(  # See "notes" on why divisions expects repetition of the last item https://docs.dask.org/en/stable/generated/dask.dataframe.DataFrame.repartition.html
+            divisions=(uniqueparents + [uniqueparents[-1]])
+        ).map_partitions(
+            parent_groupby, resolution, kwargs["aggfunc"], kwargs["decimals"]
+        )
+        if compaction:
+            ddf = ddf.map_partitions(compaction, resolution, parent_res)
+
+        ddf.to_parquet(
+            output,
+            overwrite=kwargs["overwrite"],
+            engine="pyarrow",
+            write_index=True,
+            append=False,
+            name_function=lambda i: f"{uniqueparents[i]}.parquet",
+            compression=kwargs["compression"],
         )
 
     LOGGER.debug(
@@ -214,6 +217,7 @@ def initial_index(
     dggs: str,
     dggsfunc: Callable,
     parent_groupby: Callable,
+    compaction: Union[None, Callable],
     raster_input: Union[Path, str],
     output: Path,
     resolution: int,
@@ -324,6 +328,7 @@ def process(window):
             return address_boundary_issues(
                 dggs,
                 parent_groupby,
+                compaction,
                 tmpdir,
                 output,
                 resolution,

raster2dggs/h3.py

@@ -6,6 +6,7 @@
 
 import click
 import click_log
+import h3 as h3py
 import h3pandas  # Necessary import despite lack of explicit use
 import pandas as pd
 import pyarrow as pa
@@ -75,6 +76,52 @@ def _h3_parent_groupby(
         )
 
 
+def h3_cell_to_children_size(cell, desired_resolution: int) -> int:
+    """
+    Use h3 cell conversion to determine total number of children at some offset resolution
+    """
+    # NB we enumerate all children due to the presence of pentagonal cells
+    # The H3 API has cellToChildrenSize, but it is not available in Python API?
+    return len(h3py.cell_to_children(cell, desired_resolution))
+
+
+def _h3_compaction(df: pd.DataFrame, resolution: int, parent_res: int) -> pd.DataFrame:
+    """
+    Returns a compacted version of the input dataframe.
+    Compaction only occurs if all values (i.e. bands) of the input share common values across all sibling cells.
+    Compaction will not be performed beyond parent_res or resolution.
+    It assumes and requires that the input has unique DGGS cell values as the index.
+    """
+    unprocessed_indices = set(df.index)
+    compaction_map = {}
+    for r in range(parent_res, resolution):
+        try:
+            parent_cells = map(lambda x: h3py.cell_to_parent(x, r), unprocessed_indices)
+            grouped = df.loc[list(unprocessed_indices)].groupby(list(parent_cells))
+        except ValueError as e:
+            # Indices that aren't DGGS cells; ignore and break
+            # TODO how is this possible?
+            break
+        for parent, group in grouped:
+            if parent in compaction_map:
+                continue
+            expected_count = h3_cell_to_children_size(parent, resolution)
+            if len(group) == expected_count and all(group.nunique() == 1):
+                compact_row = group.iloc[0]
+                compact_row.name = parent  # Rename the index to the parent cell
+                compaction_map[parent] = compact_row
+                unprocessed_indices -= set(group.index)
+    else:
+        # Didn't break
+        compacted_df = pd.DataFrame(list(compaction_map.values()))
+        remaining_df = df.loc[list(unprocessed_indices)]
+        result_df = pd.concat([compacted_df, remaining_df])
+        result_df = result_df.rename_axis(df.index.name)
+        return result_df
+    # Did break
+    return df
+
+
 @click.command(context_settings={"show_default": True})
 @click_log.simple_verbosity_option(common.LOGGER)
 @click.argument("raster_input", type=click.Path(), nargs=1)
@@ -148,6 +195,12 @@ def _h3_parent_groupby(
     type=click.Path(),
     help="Temporary data is created during the execution of this program. This parameter allows you to control where this data will be written.",
 )
+@click.option(
+    "-co",
+    "--compact",
+    is_flag=True,
+    help="Compact the H3 cells up to the parent resolution. Compaction is not applied for cells without identical values across all bands.",
+)
 @click.version_option(version=__version__)
 def h3(
     raster_input: Union[str, Path],
@@ -162,6 +215,7 @@ def h3(
     overwrite: bool,
     warp_mem_limit: int,
     resampling: str,
+    compact: bool,
     tempdir: Union[str, Path],
 ):
     """
@@ -192,6 +246,7 @@ def h3(
         "h3",
         _h3func,
         _h3_parent_groupby,
+        _h3_compaction if compact else None,
         raster_input,
         output_directory,
         int(resolution),