RichDEM vs WhiteboxTools vs pysheds for D8 Routing

import logging

import numpy as np
import rasterio

log = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")

# Offsets are (row, col) with row increasing southwards, matching raster order.
NEIGHBOUR_OFFSETS = {
    "E": (0, 1), "NE": (-1, 1), "N": (-1, 0), "NW": (-1, -1),
    "W": (0, -1), "SW": (1, -1), "S": (1, 0), "SE": (1, 1),
}
RICHDEM_CODES = {"E": 1, "NE": 2, "N": 3, "NW": 4, "W": 5, "SW": 6, "S": 7, "SE": 8}
ESRI_CODES = {"E": 1, "SE": 2, "S": 4, "SW": 8, "W": 16, "NW": 32, "N": 64, "NE": 128}


def convert_d8_encoding(
    src_path: str,
    dst_path: str,
    src_scheme: str = "richdem",
    dst_scheme: str = "esri",
) -> dict:
    """
    Re-encode a D8 direction raster between conventions.

    Both schemes fit in uint8, so nothing about the file signals which one it
    is. Record the scheme in the raster tags, which this function does.
    """
    schemes = {"richdem": RICHDEM_CODES, "esri": ESRI_CODES}
    if src_scheme not in schemes or dst_scheme not in schemes:
        raise ValueError(f"scheme must be one of {sorted(schemes)}")

    with rasterio.open(src_path) as src:
        arr = src.read(1)
        profile = src.profile.copy()
        tags = src.tags()

    declared = tags.get("D8_ENCODING")
    if declared and declared != src_scheme:
        log.warning("Raster is tagged as %r but was read as %r — trust the tag",
                    declared, src_scheme)

    src_codes, dst_codes = schemes[src_scheme], schemes[dst_scheme]
    out = np.zeros_like(arr)
    for direction, src_val in src_codes.items():
        n = int((arr == src_val).sum())
        out[arr == src_val] = dst_codes[direction]
        log.info("  %-2s: %8d cells, %3d → %3d", direction, n,
                 src_val, dst_codes[direction])

    unmapped = int(((arr != 0) & (out == 0)).sum())
    if unmapped:
        log.error("%d non-zero cell(s) had no mapping — the source is not in "
                  "the %r scheme", unmapped, src_scheme)

    profile.update(dtype="uint8", nodata=0, compress="LZW", tiled=True)
    with rasterio.open(dst_path, "w", **profile) as dst:
        dst.write(out.astype("uint8"), 1)
        # Tag the file so the next reader does not have to guess.
        dst.update_tags(D8_ENCODING=dst_scheme)
    log.info("Wrote %s tagged D8_ENCODING=%s", dst_path, dst_scheme)
    return {"output": dst_path, "unmapped_cells": unmapped}


def compare_implementations(dem_path: str, out_dir: str) -> dict:
    """
    Run all three D8 implementations on one conditioned DEM and compare.

    Converts every direction grid to the ESRI scheme before comparing, so the
    comparison is about the algorithm rather than about the numbering.
    """
    import os
    import time

    os.makedirs(out_dir, exist_ok=True)
    results = {}

    # --- RichDEM: in-memory, returns an rdarray carrying metadata ---
    import richdem as rd
    with rasterio.open(dem_path) as src:
        dem = src.read(1).astype(np.float32)
        nodata = src.nodata if src.nodata is not None else -9999.0
    t0 = time.perf_counter()
    grid = rd.rdarray(dem.copy(), no_data=nodata)
    rd_acc = np.asarray(rd.FlowAccumulation(grid, method="D8"))
    results["richdem"] = {"seconds": time.perf_counter() - t0,
                          "max_accumulation": float(rd_acc.max())}

    # --- WhiteboxTools: streams through files, so peak memory stays low ---
    from whitebox import WhiteboxTools
    wbt = WhiteboxTools()
    wbt.set_working_dir(os.path.abspath(out_dir))
    wbt.verbose = False
    t0 = time.perf_counter()
    wbt.d8_flow_accumulation(os.path.abspath(dem_path), "wbt_acc.tif",
                             out_type="cells")
    with rasterio.open(os.path.join(out_dir, "wbt_acc.tif")) as src:
        wbt_max = float(src.read(1).max())
    results["whitebox"] = {"seconds": time.perf_counter() - t0,
                           "max_accumulation": wbt_max}

    # --- pysheds: pure NumPy, easiest to inspect and modify ---
    from pysheds.grid import Grid
    t0 = time.perf_counter()
    pg = Grid.from_raster(dem_path)
    dem_ps = pg.read_raster(dem_path)
    fdir = pg.flowdir(dem_ps)
    ps_acc = pg.accumulation(fdir)
    results["pysheds"] = {"seconds": time.perf_counter() - t0,
                          "max_accumulation": float(np.asarray(ps_acc).max())}

    for name, r in results.items():
        log.info("%-10s %6.1f s, max accumulation %.0f cells",
                 name, r["seconds"], r["max_accumulation"])

    maxima = [r["max_accumulation"] for r in results.values()]
    spread = (max(maxima) - min(maxima)) / max(maxima)
    if spread > 0.01:
        log.warning("Maximum accumulation differs by %.2f %% across "
                    "implementations — check that all three saw the SAME "
                    "conditioned surface", 100 * spread)
    else:
        log.info("All three agree on maximum accumulation to within %.3f %%",
                 100 * spread)
    return results


# --- Example usage ---
# convert_d8_encoding("fdir_richdem.tif", "fdir_esri.tif", "richdem", "esri")
# compare_implementations("basin_conditioned.tif", "./compare")