Validating Hydro-Conditioning Against NHD Flowlines

import logging

import geopandas as gpd
import numpy as np
import rasterio
from rasterio.features import rasterize
from shapely.geometry import mapping

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


def validate_conditioning(
    facc_path: str,
    nhd_path: str,
    threshold_cells: int,
    gauges: gpd.GeoDataFrame | None = None,
    area_field: str = "drain_area_km2",
    buffer_cells: int = 2,
) -> dict:
    """
    Score a conditioned surface against NHD and, where available, against
    published gauge drainage areas.

    Parameters
    ----------
    facc_path       : Flow accumulation raster from the conditioned DEM.
    nhd_path        : Reference flowlines.
    threshold_cells : Calibrated channel-initiation threshold.
    gauges          : Optional gauge points carrying a published area field.
    area_field      : Column on `gauges` holding the published area in km².
    buffer_cells    : Spatial tolerance for the overlap comparison.

    Returns
    -------
    Dict with overlap metrics and, when gauges are supplied, per-gauge
    drainage-area residuals.
    """
    with rasterio.open(facc_path) as src:
        acc = src.read(1)
        transform, crs, shape = src.transform, src.crs, (src.height, src.width)
        cell_size = abs(src.transform.a)
        cell_area_km2 = (cell_size ** 2) / 1e6

    extracted = acc >= threshold_cells

    # --- Rasterize NHD with a tolerance buffer. The buffer absorbs the real
    # geometric difference between a generalised vector line and a network
    # that follows cell centres — typically one to two cells. ---
    nhd = gpd.read_file(nhd_path).to_crs(crs)
    buffered = nhd.geometry.buffer(buffer_cells * cell_size)
    reference = rasterize(
        ((mapping(g), 1) for g in buffered if g is not None and not g.is_empty),
        out_shape=shape, transform=transform, fill=0, dtype=np.uint8,
    ).astype(bool)

    tp = int((extracted & reference).sum())
    fp = int((extracted & ~reference).sum())
    fn = int((~extracted & reference).sum())
    precision = tp / (tp + fp) if tp + fp else 0.0
    recall = tp / (tp + fn) if tp + fn else 0.0
    f1 = 2 * precision * recall / (precision + recall) if precision + recall else 0.0

    log.info("Network overlap: precision %.3f, recall %.3f, F1 %.3f "
             "(buffer %d cells)", precision, recall, f1, buffer_cells)

    result = {"precision": precision, "recall": recall, "f1": f1}

    # --- The independent check. Delineated area at each gauge against the
    # published area, which was derived from neither the DEM nor the NHD
    # geometry used above. ---
    if gauges is not None and len(gauges):
        residuals = []
        gauges = gauges.to_crs(crs)
        for _, row in gauges.iterrows():
            r, c = rasterio.transform.rowcol(transform, row.geometry.x, row.geometry.y)
            if not (0 <= r < shape[0] and 0 <= c < shape[1]):
                log.warning("Gauge %s falls outside the raster", row.get("site_no", "?"))
                continue
            delineated_km2 = float(acc[r, c]) * cell_area_km2
            published_km2 = float(row[area_field])
            if published_km2 <= 0:
                continue
            pct = 100.0 * (delineated_km2 - published_km2) / published_km2
            residuals.append({
                "site": row.get("site_no", "?"),
                "delineated_km2": delineated_km2,
                "published_km2": published_km2,
                "residual_pct": pct,
            })
            log.info("gauge %s: delineated %.1f km², published %.1f km² (%+.1f %%)",
                     row.get("site_no", "?"), delineated_km2, published_km2, pct)

        if residuals:
            pcts = np.array([r["residual_pct"] for r in residuals])
            median, spread = float(np.median(pcts)), float(np.percentile(np.abs(pcts), 90))
            result.update({
                "gauges": residuals,
                "median_residual_pct": median,
                "p90_abs_residual_pct": spread,
            })
            # A consistent one-sided offset is a systematic problem, not a set
            # of individual snapping errors.
            if abs(median) > 5.0:
                log.error("Median residual %+.1f %% across %d gauges — this is a "
                          "conditioning or extent problem, not per-gauge snapping",
                          median, len(residuals))
            else:
                log.info("Median residual %+.1f %%, p90 |residual| %.1f %% over %d gauges",
                         median, spread, len(residuals))

    return result


# --- Example usage ---
# gauges = gpd.read_file("usgs_gauges.gpkg")
# validate_conditioning(
#     facc_path="basin_conditioned_acc.tif",
#     nhd_path="nhd_flowlines.gpkg",
#     threshold_cells=500,
#     gauges=gauges, area_field="drain_area_km2",
# )