Retrieving USGS Streamflow with dataretrieval-python

import logging

import numpy as np
import pandas as pd
import dataretrieval.nwis as nwis

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

CFS_TO_CMS = 0.0283168
DISCHARGE = "00060"     # discharge, cubic feet per second
DAILY_MEAN = "00003"    # statistic code: daily mean


def fetch_daily_discharge(
    site: str,
    start: str,
    end: str,
    approved_only: bool = True,
) -> pd.Series:
    """
    Daily mean discharge for one site, in m³/s, screened on approval flags.

    Parameters
    ----------
    site          : Site number as a STRING — an int drops the leading zero.
    start, end    : ISO dates.
    approved_only : Drop provisional values, which can change on review.
    """
    if not isinstance(site, str):
        raise TypeError("site must be a string; an int drops the leading zero "
                        "that most eastern sites carry")

    df, meta = nwis.get_dv(sites=site, start=start, end=end,
                           parameterCd=DISCHARGE, statCd=DAILY_MEAN)
    if df.empty:
        raise ValueError(f"no daily discharge for {site} in {start}..{end}")

    value_col = next(c for c in df.columns if c.endswith(f"{DISCHARGE}_{DAILY_MEAN}"))
    flag_col = value_col + "_cd"

    n_all = len(df)
    if approved_only and flag_col in df.columns:
        # A = approved, P = provisional. Provisional high flows in particular
        # move on review, because the rating is extrapolated there.
        approved = df[flag_col].astype(str).str.contains("A")
        n_prov = int((~approved).sum())
        if n_prov:
            log.warning("Dropping %d provisional value(s) of %d (%.1f %%)",
                        n_prov, n_all, 100.0 * n_prov / n_all)
        df = df[approved]

    series = df[value_col].astype(float) * CFS_TO_CMS
    series.index = pd.to_datetime(df.index).tz_localize(None)
    series.name = "discharge_cms"

    # --- Report the gaps explicitly. A metric computed over a series with a
    # four-month hole is not wrong, but it is not what the caller thinks. ---
    full = pd.date_range(series.index.min(), series.index.max(), freq="D")
    missing = len(full) - len(series)
    log.info("%s: %d days from %s to %s, %d missing (%.2f %%)",
             site, len(series), series.index.min().date(),
             series.index.max().date(), missing, 100.0 * missing / len(full))
    if missing:
        gaps = full.difference(series.index)
        runs = _gap_runs(gaps)
        longest = max((r[1] for r in runs), default=0)
        log.warning("%d gap(s), longest %d consecutive day(s)", len(runs), longest)
    return series


def fetch_annual_peaks(site: str) -> pd.DataFrame:
    """
    The reviewed annual peak series, with qualification codes preserved.

    Returns a frame indexed by water year with columns peak_cms and codes.
    """
    df, _ = nwis.get_discharge_peaks(sites=site)
    if df.empty:
        raise ValueError(f"no annual peak series published for {site}")

    out = pd.DataFrame({
        "peak_cms": df["peak_va"].astype(float) * CFS_TO_CMS,
        "codes": df.get("peak_cd", pd.Series(index=df.index, dtype=object)),
    })
    dates = pd.to_datetime(df["peak_dt"], errors="coerce")
    # Water year runs October to September, so a December peak belongs to the
    # following water year. Using calendar years splits winter events.
    out["water_year"] = dates.dt.year + (dates.dt.month >= 10).astype(int)
    out = out.dropna(subset=["water_year"]).set_index("water_year")

    # Qualification codes 5, 6 and C flag peaks affected by regulation or
    # by an urban or agricultural change — not natural floods.
    regulated = out["codes"].astype(str).str.contains("[56C]", regex=True, na=False)
    if regulated.any():
        log.warning("%d of %d peaks are flagged as regulated or affected by "
                    "change — exclude them or fit the periods separately",
                    int(regulated.sum()), len(out))
    out["regulated"] = regulated

    log.info("%s: %d annual peaks, %d–%d, max %.1f m³/s in %d",
             site, len(out), int(out.index.min()), int(out.index.max()),
             float(out["peak_cms"].max()), int(out["peak_cms"].idxmax()))
    return out


def _gap_runs(gaps: pd.DatetimeIndex):
    """Collapse a set of missing dates into (start, length) runs."""
    if len(gaps) == 0:
        return []
    breaks = np.where(np.diff(gaps.values).astype("timedelta64[D]").astype(int) > 1)[0]
    starts = np.concatenate([[0], breaks + 1])
    ends = np.concatenate([breaks, [len(gaps) - 1]])
    return [(gaps[s], int(e - s + 1)) for s, e in zip(starts, ends)]


# --- Example usage ---
# q = fetch_daily_discharge("03339000", "1990-10-01", "2024-09-30")
# peaks = fetch_annual_peaks("03339000")
# natural = peaks[~peaks.regulated]["peak_cms"]