Deconvolving Observed Hydrographs to Derive a Unit Hydrograph

import logging

import numpy as np
import pandas as pd
from scipy.optimize import nnls

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


def separate_baseflow(q: pd.Series, alpha: float = 0.925, passes: int = 3) -> pd.Series:
    """
    Recursive digital filter baseflow separation (Lyne-Hollick).

    Returns direct runoff. The filter is run forward, backward and forward
    again, which is the usual three-pass convention and reduces the
    dependence on where the series happens to start.
    """
    q_arr = q.to_numpy(dtype="float64")
    direct = q_arr.copy()
    for p in range(passes):
        seq = direct if p % 2 == 0 else direct[::-1]
        out = np.zeros_like(seq)
        prev_direct, prev_q = 0.0, seq[0]
        for i, val in enumerate(seq):
            d = alpha * prev_direct + 0.5 * (1 + alpha) * (val - prev_q)
            out[i] = min(max(d, 0.0), val)   # direct runoff is bounded by total
            prev_direct, prev_q = out[i], val
        direct = out if p % 2 == 0 else out[::-1]

    result = pd.Series(direct, index=q.index, name="direct_runoff")
    frac = float(result.sum() / max(q.sum(), 1e-9))
    log.info("Baseflow separation: %.1f %% of total volume is direct runoff",
             100 * frac)
    if frac > 0.95:
        log.warning("Almost no baseflow separated — check alpha, or the event "
                    "may start on a rising limb")
    return result


def derive_unit_hydrograph(
    excess_mm: pd.Series,
    direct_runoff_cms: pd.Series,
    area_km2: float,
    n_ordinates: int | None = None,
) -> pd.Series:
    """
    Recover unit hydrograph ordinates by non-negative deconvolution.

    Parameters
    ----------
    excess_mm         : Excess rainfall per timestep.
    direct_runoff_cms : Direct runoff (baseflow already removed).
    area_km2          : Catchment area, for the volume check.
    n_ordinates       : Length of the unit hydrograph. Defaults to the
                        difference in length between runoff and excess plus one,
                        which is the maximum the data can support.
    """
    p = excess_mm.to_numpy(dtype="float64")
    q = direct_runoff_cms.to_numpy(dtype="float64")
    p = p[p > 1e-9] if (p > 1e-9).any() else p

    m = len(q)
    n = n_ordinates or (m - len(p) + 1)
    if n < 3:
        raise ValueError("the runoff series is too short relative to the "
                         "rainfall to support a unit hydrograph")

    # --- Build the convolution matrix. Row i, column j holds the excess
    # increment that, delayed by j steps, contributes to runoff at step i. ---
    P = np.zeros((m, n))
    for i in range(m):
        for j in range(n):
            k = i - j
            if 0 <= k < len(p):
                P[i, j] = p[k]

    # --- Non-negative least squares. An ordinary lstsq here returns
    # alternating-sign ordinates that fit marginally better and mean nothing. ---
    u, residual = nnls(P, q)
    log.info("Deconvolved %d ordinate(s) from %d timesteps, residual norm %.3f",
             n, m, float(residual))

    n_zero = int((u < 1e-9).sum())
    if n_zero > n // 2:
        log.warning("%d of %d ordinates are zero — the constraint is doing most "
                    "of the work, which usually means a multi-peaked event",
                    n_zero, n)

    # --- Volume check. One unit (1 mm) of excess over the catchment must equal
    # the area under the unit hydrograph. This is the check that catches a
    # units error, and it is the reason to pass area in. ---
    dt_s = (direct_runoff_cms.index[1] - direct_runoff_cms.index[0]).total_seconds()
    uh_volume_m3 = float(u.sum()) * dt_s
    expected_m3 = area_km2 * 1e6 * 0.001          # 1 mm over the catchment
    ratio = uh_volume_m3 / expected_m3
    log.info("Unit hydrograph volume %.0f m³ vs %.0f m³ for 1 mm (ratio %.3f)",
             uh_volume_m3, expected_m3, ratio)
    if abs(ratio - 1.0) > 0.05:
        log.warning("Rescaling ordinates by %.3f to enforce unit volume", 1 / ratio)
        u = u / ratio

    index = [direct_runoff_cms.index[0] + pd.Timedelta(seconds=dt_s * i)
             for i in range(n)]
    uh = pd.Series(u, index=index, name="uh_cms_per_mm")
    log.info("Peak %.3f m³/s per mm at ordinate %d of %d",
             float(uh.max()), int(np.argmax(u)) + 1, n)
    return uh


def average_unit_hydrographs(uhs: list[pd.Series]) -> pd.Series:
    """
    Average several derived unit hydrographs by aligning their peaks.

    Averaging ordinate by ordinate smears the peak whenever the individual
    peaks fall at different times, producing a result flatter than any input.
    """
    if not uhs:
        raise ValueError("no unit hydrographs supplied")
    peaks = [int(np.argmax(u.to_numpy())) for u in uhs]
    target_peak = int(round(float(np.mean(peaks))))
    length = max(len(u) for u in uhs) + max(peaks) - min(peaks)

    stack = np.full((len(uhs), length), np.nan)
    for i, (u, pk) in enumerate(zip(uhs, peaks)):
        offset = target_peak - pk
        stack[i, offset:offset + len(u)] = u.to_numpy()
    mean = np.nanmean(stack, axis=0)
    mean = np.nan_to_num(mean, nan=0.0)

    # Rescale so the average still integrates to one unit.
    scale = sum(float(u.sum()) for u in uhs) / len(uhs) / max(mean.sum(), 1e-12)
    mean = mean * scale

    log.info("Averaged %d unit hydrographs; individual peaks at ordinates %s, "
             "aligned on %d", len(uhs), peaks, target_peak)
    log.info("Averaged peak %.3f, mean of individual peaks %.3f",
             float(mean.max()), float(np.mean([u.max() for u in uhs])))
    return pd.Series(mean, name="uh_averaged")


# --- Example usage ---
# direct = separate_baseflow(observed_q)
# uh = derive_unit_hydrograph(excess, direct, area_km2=214.0)
# final = average_unit_hydrographs([uh1, uh2, uh3, uh4, uh5])