Batch Validating Delineated Areas Against USGS Gauge Metadata

import logging

import geopandas as gpd
import numpy as np
import pandas as pd
import rasterio
import dataretrieval.nwis as nwis

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

SQMI_TO_SQKM = 2.58999


def gauges_in_domain(bounds_4326, min_area_km2: float = 5.0) -> gpd.GeoDataFrame:
    """
    Every streamflow gauge inside a bounding box, with its published area.

    Filters out very small basins, where a few boundary cells dominate the
    residual and the comparison says more about the DEM than the pipeline.
    """
    minx, miny, maxx, maxy = bounds_4326
    sites, _ = nwis.get_info(bBox=f"{minx:.6f},{miny:.6f},{maxx:.6f},{maxy:.6f}",
                             siteType="ST", hasDataTypeCd="dv")
    if sites.empty:
        raise ValueError("no stream gauges in this bounding box")

    df = sites.dropna(subset=["drain_area_va"]).copy()
    df["published_km2"] = df["drain_area_va"].astype(float) * SQMI_TO_SQKM

    # Contributing drainage area, where published, excludes closed basins
    # upstream. Where it differs from the total, it is the number a DEM
    # delineation with depressions removed should be compared against.
    if "contrib_drain_area_va" in df.columns:
        contrib = pd.to_numeric(df["contrib_drain_area_va"], errors="coerce")
        n_contrib = int(contrib.notna().sum())
        if n_contrib:
            log.info("%d site(s) publish a separate contributing area — using "
                     "it where present", n_contrib)
            df["published_km2"] = np.where(
                contrib.notna(), contrib * SQMI_TO_SQKM, df["published_km2"])

    df = df[df["published_km2"] >= min_area_km2]
    gdf = gpd.GeoDataFrame(
        df, geometry=gpd.points_from_xy(df["dec_long_va"], df["dec_lat_va"]),
        crs=4326,
    )
    log.info("%d gauges with a published area ≥ %.1f km² in the domain",
             len(gdf), min_area_km2)
    return gdf[["site_no", "station_nm", "published_km2", "geometry"]]


def batch_validate(
    facc_path: str,
    gauges: gpd.GeoDataFrame,
    snap_radius_m: float = 150.0,
) -> pd.DataFrame:
    """
    Compare the delineated area at every gauge against its published area.

    Snapping is done by taking the maximum accumulation within the search
    radius, which is the cheap equivalent of a full snap and adequate for a
    batch check.
    """
    with rasterio.open(facc_path) as src:
        acc = src.read(1).astype("float64")
        transform, crs = src.transform, src.crs
        cell_km2 = (abs(src.transform.a) ** 2) / 1e6
        h, w = acc.shape
        radius_cells = int(round(snap_radius_m / abs(src.transform.a)))

    pts = gauges.to_crs(crs)
    rows = []
    for _, g in pts.iterrows():
        r, c = rasterio.transform.rowcol(transform, g.geometry.x, g.geometry.y)
        if not (0 <= r < h and 0 <= c < w):
            log.warning("%s falls outside the raster", g.site_no)
            continue
        r0, r1 = max(0, r - radius_cells), min(h, r + radius_cells + 1)
        c0, c1 = max(0, c - radius_cells), min(w, c + radius_cells + 1)
        window = acc[r0:r1, c0:c1]
        if not np.isfinite(window).any():
            continue
        delineated = float(np.nanmax(window)) * cell_km2
        published = float(g.published_km2)
        rows.append({
            "site_no": g.site_no,
            "station_nm": g.station_nm,
            "published_km2": published,
            "delineated_km2": delineated,
            "residual_pct": 100.0 * (delineated - published) / published,
        })

    df = pd.DataFrame(rows)
    if df.empty:
        raise ValueError("no gauges could be evaluated")

    median = float(df["residual_pct"].median())
    mad = float(np.median(np.abs(df["residual_pct"] - median)))
    # Flag outliers relative to the median, not to zero: if the whole layer is
    # shifted, an outlier is a gauge that departs from the SHIFT.
    df["is_outlier"] = np.abs(df["residual_pct"] - median) > max(5.0, 4.0 * mad)

    log.info("%d gauges evaluated: median residual %+.2f %%, MAD %.2f %%",
             len(df), median, mad)
    if abs(median) > 3.0:
        log.error("A median residual of %+.2f %% across %d gauges is a "
                  "SYSTEMATIC problem — check the DEM extent, the conditioning "
                  "and the projection before looking at individual gauges",
                  median, len(df))
    n_out = int(df["is_outlier"].sum())
    if n_out:
        log.warning("%d gauge(s) depart from the population — likely snapped to "
                    "the wrong branch:", n_out)
        for _, r in df[df["is_outlier"]].iterrows():
            log.warning("  %s  %.1f vs %.1f km² (%+.1f %%)  %s",
                        r.site_no, r.delineated_km2, r.published_km2,
                        r.residual_pct, r.station_nm[:40])

    # A residual that grows with basin area is a scale error, not a
    # delineation one — test it explicitly rather than reading it off a plot.
    if len(df) >= 8:
        corr = float(np.corrcoef(np.log10(df["published_km2"]),
                                 df["residual_pct"])[0, 1])
        if abs(corr) > 0.6:
            log.error("Residual correlates with basin size (r = %.2f) — this is "
                      "a projection or unit error, not a delineation error", corr)
        else:
            log.info("No size dependence in the residuals (r = %.2f)", corr)
    return df.sort_values("residual_pct")


# --- Example usage ---
# gauges = gauges_in_domain((-84.5, 35.2, -83.1, 36.4))
# report = batch_validate("basin_acc.tif", gauges)
# report.to_csv("delineation_validation.csv", index=False)