hydrun_functions.py

"""Python translation of the HydRun MATLAB functions.

All time series are NumPy arrays with columns ``[MATLAB serial date, value]``.
Keeping serial dates makes the original thresholds and ``BB_data.mat`` directly
compatible. Public functions also have MATLAB-style aliases at the end.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Callable, Sequence

import numpy as np
from scipy.optimize import least_squares


@dataclass(frozen=True)
class HydroTimeInstants:
    start: float
    end: float
    peak: float
    centroid: float


@dataclass(frozen=True)
class RainTimeInstants:
    start: float
    end: float
    centroid: float


TIME_CHARACTERISTIC_NAMES = (
    "T_w", "T_LR", "T_r", "T_LP", "T_LPC", "T_LC", "T_b", "T_c"
)


def _series(data, allow_empty=False):
    array = np.asarray(data, dtype=float)
    if allow_empty and array.size == 0:
        return np.empty((0, 2))
    if array.ndim != 2 or array.shape[1] != 2:
        raise ValueError("A time series must have shape (n, 2).")
    return array


def _mode_step(times):
    if len(times) < 2:
        return np.nan
    values, counts = np.unique(np.round(np.diff(times), 12), return_counts=True)
    return float(values[np.argmax(counts)])


def smooth_curve(values, passes):
    """Apply HydRun's sequential three-point moving-average smoother."""
    result = np.asarray(values, dtype=float).copy().reshape(-1)
    for _ in range(max(0, int(passes))):
        # MATLAB updates the preceding point in-place, so this is sequential.
        for i in range(1, len(result) - 1):
            result[i] = (result[i - 1] + 2 * result[i] + result[i + 1]) / 4
    return result


def find_turning_points(line):
    """Return zero-based indices and labels: peak=1, valley=0."""
    line = np.asarray(line, dtype=float).reshape(-1)
    sign_difference = np.diff(np.sign(np.diff(line)))
    indices = np.flatnonzero(sign_difference != 0) + 1
    labels = (sign_difference[indices - 1] < 0).astype(int)
    return np.column_stack((indices, labels)).astype(int)


def separate_baseflow(hydrograph, filter_coefficient, passes):
    """Separate streamflow into stormflow and baseflow using a digital filter."""
    hydrograph = _series(hydrograph)
    flow = hydrograph[:, 1]
    previous_pass = flow.copy()
    baseflow = np.empty_like(flow)
    for pass_number in range(int(passes)):
        indices = (list(range(len(flow))) if pass_number % 2 == 0
                   else list(range(len(flow) - 1, -1, -1)))
        baseflow[indices[0]] = previous_pass[indices[0]]
        for current, previous in zip(indices[1:], indices[:-1]):
            candidate = (filter_coefficient * baseflow[previous] +
                         (1 - filter_coefficient) *
                         (previous_pass[current] + previous_pass[previous]) / 2)
            baseflow[current] = min(candidate, previous_pass[current])
        previous_pass = baseflow.copy()
    stormflow = flow - baseflow
    return (np.column_stack((hydrograph[:, 0], stormflow)),
            np.column_stack((hydrograph[:, 0], baseflow)))


def extract_runoff(stormflow, min_difference, return_ratio, beginning_slope,
                   end_slope, smoothing_passes, minimum_duration=0,
                   dynamic_slope=0.001):
    """Extract runoff events from a baseflow-free hydrograph."""
    stormflow = _series(stormflow)
    smooth = smooth_curve(stormflow[:, 1], smoothing_passes)
    turning = find_turning_points(smooth)
    if not len(turning):
        return [], 0
    turning = np.column_stack((turning, smooth[turning[:, 0]]))
    mean_flow = np.mean(smooth)
    if turning[0, 1] == 1 and smooth[0] < mean_flow / 10:
        turning = np.vstack(([0, 0, smooth[0]], turning))
    if turning[-1, 1] == 1 and smooth[-1] < mean_flow / 10:
        turning = np.vstack((turning, [len(smooth) - 1, 0, smooth[-1]]))
    while len(turning) and turning[0, 1] == 1:
        turning = turning[1:]
    while len(turning) and turning[-1, 1] == 1:
        turning = turning[:-1]
    if len(turning) < 3:
        return [], 0

    differences = np.r_[np.diff(turning[:, 2]), 0.0]
    starts, ends = [], []
    i, complete = 0, True
    return_constant = min_difference / 3
    while i < len(turning) - 2:
        j = 1
        difference = differences[i] + differences[i + j]
        while difference > max(
            return_ratio * np.max(np.abs(differences[i:i + j + 1])),
            return_constant,
        ):
            if i + j < len(turning) - 2:
                j += 1
                difference += differences[i + j]
            else:
                complete = False
                break
        starts.append(i)
        ends.append(i + j)
        i += j + 1
    if not complete and starts:
        starts.pop()
        ends.pop()

    events = []
    for start, end in zip(starts, ends):
        first = int(turning[start, 0])
        last = int(turning[end + 1, 0])
        event = np.column_stack((stormflow[first:last + 1], smooth[first:last + 1]))
        if (np.max(event[:, 1]) - event[0, 1] <= min_difference or
                np.max(event[:, 1]) - event[-1, 1] <= min_difference):
            continue
        # Reset for each event (the MATLAB code accidentally mutates dyslp globally).
        slope_limit = dynamic_slope * np.ptp(event[:, 1])
        begin_limit = min(beginning_slope, slope_limit)
        finish_limit = min(end_slope, slope_limit)
        delta = np.diff(event[:, 2])
        while len(delta) and delta[0] < begin_limit:
            event, delta = event[1:], delta[1:]
        while len(delta) and delta[-1] > -finish_limit:
            event, delta = event[:-1], delta[:-1]
        delta = np.diff(event[:, 1])
        while len(delta) and delta[0] <= begin_limit:
            event, delta = event[1:], delta[1:]
        while len(delta) and delta[-1] >= -finish_limit:
            event, delta = event[:-1], delta[:-1]
        if len(delta) > minimum_duration:
            events.append(event[:, :2].copy())
    return events, len(events)


def extract_precipitation_events(precipitation, max_allowable_gap,
                                 threshold=0.0):
    """Extract rainfall events; gap is in hours and threshold is in mm."""
    original = _series(precipitation)
    clean = original.copy()
    clean[np.isnan(clean[:, 1]), 1] = 0
    gap_steps = round((max_allowable_gap / 2) / (_mode_step(clean[:, 0]) * 24))
    gap_steps = max(0, int(gap_steps))
    active = np.convolve(clean[:, 1], np.ones(2 * gap_steps + 1), mode="same") != 0
    padded = np.r_[False, active, False]
    starts = np.flatnonzero(np.diff(padded.astype(int)) == 1)
    ends = np.flatnonzero(np.diff(padded.astype(int)) == -1)
    events, missing_proportions = [], []
    for start, end in zip(starts, ends):
        core_start, core_end = start + gap_steps, end - gap_steps
        if core_end <= core_start:
            continue
        event = clean[core_start:core_end].copy()
        if np.sum(event[:, 1]) > threshold:
            events.append(event)
            missing_proportions.append(np.mean(np.isnan(original[start:end, 1])))
    return events, len(events), np.asarray(missing_proportions)


def compute_hydrograph_time_instants(runoff_event):
    event = _series(runoff_event)
    peak = int(np.argmax(event[:, 1]))
    return HydroTimeInstants(
        float(event[0, 0]), float(event[-1, 0]), float(event[peak, 0]),
        float(np.sum(event[:, 0] * event[:, 1]) / np.sum(event[:, 1])),
    )


def compute_hyetograph_time_instants(rainfall_event):
    event = _series(rainfall_event, allow_empty=True)
    if not len(event) or np.all(event[:, 1] == 0):
        return RainTimeInstants(np.nan, np.nan, np.nan)
    interval = _mode_step(event[:, 0])
    nonzero = np.flatnonzero(event[:, 1] != 0)
    centroid = np.nansum(event[:, 0] * event[:, 1]) / np.nansum(event[:, 1])
    return RainTimeInstants(
        float(event[nonzero[0], 0] - interval),
        float(event[nonzero[-1], 0]),
        float(centroid - interval / 2),
    )


def match_rainfall_runoff(rain_events, runoff_events, search_hours):
    """Pair rainfall events overlapping each runoff event's search window."""
    rain_starts = np.array([event[0, 0] for event in rain_events])
    rain_ends = np.array([event[-1, 0] for event in rain_events])
    pairs, relations, previous = [], [], None
    for runoff_number, runoff in enumerate(runoff_events, 1):
        hydro = compute_hydrograph_time_instants(runoff)
        window_start = hydro.start - search_hours / 24
        if previous is not None:
            previous_centroid = compute_hydrograph_time_instants(previous).centroid
            window_start = max(window_start, previous_centroid)
        turning = find_turning_points(smooth_curve(runoff[:, 1], 2))
        peak_indices = turning[turning[:, 1] == 1, 0]
        window_end = runoff[peak_indices[-1], 0] if len(peak_indices) else hydro.peak
        matched = np.flatnonzero((rain_starts <= window_end) &
                                 (rain_ends >= window_start))
        rainfall = (np.vstack([rain_events[index] for index in matched])
                    if len(matched) else np.empty((0, 2)))
        pairs.append((rainfall, runoff))
        relations.append({
            "runoff_index": runoff_number,
            "num_rain_events": len(matched),
            "rain_event_indices": (matched + 1).tolist(),
        })
        previous = runoff
    return pairs, relations


def compute_time_characteristics(rainfall_event, runoff_event, kind=None):
    """Compute the eight rainfall-runoff time characteristics in hours."""
    rain = compute_hyetograph_time_instants(rainfall_event)
    runoff = compute_hydrograph_time_instants(runoff_event)
    values = np.array([
        rain.end - rain.start,
        runoff.start - rain.start,
        runoff.peak - runoff.start,
        runoff.peak - rain.start,
        runoff.peak - rain.centroid,
        runoff.centroid - rain.centroid,
        runoff.end - runoff.start,
        runoff.end - rain.end,
    ]) * 24
    if kind is None:
        return values
    lookup = {"w": 0, "LR": 1, "r": 2, "LP": 3, "LPC": 4,
              "LC": 5, "b": 6, "c": 7}
    return float(values[lookup[kind]]) if kind in lookup else np.nan


def compute_runoff_ratio(rainfall_event, runoff_event, drainage_area):
    """Return ratio, runoff volume (m3), and precipitation volume (m3)."""
    rainfall = _series(rainfall_event, allow_empty=True)
    runoff = _series(runoff_event)
    interval_seconds = (runoff[1, 0] - runoff[0, 0]) * 86400
    runoff_volume = float(np.sum(runoff[:, 1]) * interval_seconds)
    precip_volume = (float(np.nansum(rainfall[:, 1]) * drainage_area * 1000)
                     if len(rainfall) else 0.0)
    ratio = runoff_volume / precip_volume if precip_volume else np.inf
    return ratio, runoff_volume, precip_volume


def compute_initial_abstraction(rainfall_event, runoff_event):
    rainfall = _series(rainfall_event, allow_empty=True)
    runoff = _series(runoff_event)
    if not len(rainfall):
        return 0.0
    return float(np.nansum(rainfall[rainfall[:, 0] < runoff[0, 0], 1]))


def compute_antecedent_precipitation(precipitation, runoff_event,
                                     antecedent_days, search_hours=0):
    precipitation, runoff = _series(precipitation), _series(runoff_event)
    values = precipitation[
        (precipitation[:, 0] > runoff[0, 0] - antecedent_days) &
        (precipitation[:, 0] < runoff[0, 0] - search_hours / 24), 1
    ]
    missing = int(np.isnan(values).sum())
    proportion = missing / len(values) if len(values) else np.nan
    return float(np.nansum(values)), missing, proportion


def _fit_exponential(x, y):
    fit = least_squares(lambda p: p[0] * np.exp(p[1] * x) - y, [0.0, 0.0])
    return fit.x


def compute_recession_constant(runoff_event, peak_choice="last"):
    event = _series(runoff_event)
    if peak_choice == "last":
        turning = find_turning_points(smooth_curve(event[:, 1], 4))
        peaks = turning[turning[:, 1] == 1, 0]
        start = int(peaks[-1]) if len(peaks) else int(np.argmax(event[:, 1]))
        recession = event[start:].copy()
        recession = recession[np.argmax(recession[:, 1]):]
    elif peak_choice == "highest":
        recession = event[np.argmax(event[:, 1]):].copy()
    else:
        raise ValueError("peak_choice must be 'last' or 'highest'.")
    while len(recession) and recession[-1, 1] == 0:
        recession = recession[:-1]
    if len(recession) < 2:
        raise ValueError("The recession limb has fewer than two nonzero points.")
    x = (recession[:, 0] - recession[0, 0]) * 24
    coefficients = _fit_exponential(x, recession[:, 1])
    simulated = coefficients[0] * np.exp(coefficients[1] * x)
    span = np.ptp(recession[:, 1])
    nrmse = (np.sqrt(np.mean((recession[:, 1] - simulated) ** 2)) / span
             if span else np.nan)
    constant = -1 / coefficients[1] if coefficients[1] else np.inf
    return (float(constant), float(nrmse), recession,
            np.column_stack((recession[:, 0], simulated)))


def generalize_recession(recession_limbs):
    """Fit a master recession curve to normalized recession limbs."""
    normalized = []
    for limb in recession_limbs:
        limb = _series(limb)
        normalized.append(np.column_stack((
            (limb[:, 0] - limb[0, 0]) * 24,
            limb[:, 1] / np.max(limb[:, 1]),
        )))
    combined = np.vstack(normalized)
    coefficients = _fit_exponential(combined[:, 0], combined[:, 1])
    fitted = coefficients[0] * np.exp(coefficients[1] * combined[:, 0])
    residual = np.sum((combined[:, 1] - fitted) ** 2)
    total = np.sum((combined[:, 1] - np.mean(combined[:, 1])) ** 2)
    r_squared = 1 - residual / total
    maximum_time = max(limb[-1, 0] for limb in normalized)
    x = np.linspace(0, maximum_time, 500)
    master_curve = np.column_stack((
        x, coefficients[0] * np.exp(coefficients[1] * x)
    ))
    return (-1 / coefficients[1], r_squared, coefficients,
            master_curve, normalized)


def batch_process(function: Callable, event_parameter_indices, *arguments):
    """Apply a function to event lists; parameter indices are zero-based."""
    indices = np.atleast_1d(event_parameter_indices).astype(int)
    event_count = len(arguments[indices[0]])
    rows = []
    for event_index in range(event_count):
        call_arguments = list(arguments)
        for parameter_index in indices:
            call_arguments[parameter_index] = arguments[parameter_index][event_index]
        result = function(*call_arguments)
        rows.append(result if isinstance(result, tuple) else (result,))
    columns = tuple([row[i] for row in rows] for i in range(len(rows[0])))
    return columns[0] if len(columns) == 1 else columns


def plot_runoff_events(runoff_events, hydrograph, show_ids=False, ax=None):
    import matplotlib.pyplot as plt
    hydrograph = _series(hydrograph)
    ax = ax or plt.subplots()[1]
    ax.plot(hydrograph[:, 0], hydrograph[:, 1], label="streamflow", linewidth=1)
    for number, event in enumerate(runoff_events, 1):
        selected = hydrograph[(hydrograph[:, 0] >= event[0, 0]) &
                              (hydrograph[:, 0] <= event[-1, 0])]
        ax.plot(selected[:, 0], selected[:, 1], "g", linewidth=1.5,
                label="runoff event" if number == 1 else None)
        ax.plot(selected[[0, -1], 0], selected[[0, -1], 1], "r.", markersize=10)
        if show_ids:
            peak = selected[np.argmax(selected[:, 1])]
            ax.text(peak[0], peak[1], str(number))
    ax.set(xlim=(hydrograph[0, 0], hydrograph[-1, 0]),
           ylim=(0, 1.05 * np.nanmax(hydrograph[:, 1])))
    ax.legend()
    return ax


def plot_master_recession_curve(master_curve, normalized_limbs, ax=None):
    import matplotlib.pyplot as plt
    ax = ax or plt.subplots()[1]
    for i, limb in enumerate(normalized_limbs):
        ax.plot(limb[:, 0], limb[:, 1], color="0.7",
                label="Individual recession" if i == 0 else None)
    ax.plot(master_curve[:, 0], master_curve[:, 1], "b", linewidth=1.2,
            label="MRC")
    ax.set(xlim=(master_curve[0, 0], master_curve[-1, 0]), ylim=(0, 1),
           xlabel="Hours", ylabel="Normalized discharge")
    ax.legend()
    return ax


def plot_event(rainfall_event, runoff_event, axes=None):
    """Plot one rainfall-runoff pair and mark its characteristic times."""
    import matplotlib.pyplot as plt
    rainfall = _series(rainfall_event, allow_empty=True)
    runoff = _series(runoff_event, allow_empty=True)
    if not len(rainfall) and not len(runoff):
        raise ValueError("At least one event must be non-empty.")
    if not len(rainfall):
        rainfall = np.column_stack((runoff[:, 0], np.full(len(runoff), np.nan)))
    if not len(runoff):
        runoff = np.column_stack((rainfall[:, 0], np.full(len(rainfall), np.nan)))
    if axes is None:
        _, axes = plt.subplots(2, 1, sharex=True)
    rain_time = compute_hyetograph_time_instants(rainfall)
    flow_time = compute_hydrograph_time_instants(runoff)
    axes[0].bar(rainfall[:, 0], rainfall[:, 1], width=_mode_step(rainfall[:, 0]),
                color="navy", edgecolor="dodgerblue")
    for time in (rain_time.start, rain_time.end):
        axes[0].axvline(time, linestyle="--", color="teal")
    axes[0].axvline(rain_time.centroid, linestyle="--", color="m")
    axes[0].set_ylabel("Precipitation (mm)")
    axes[1].plot(runoff[:, 0], runoff[:, 1], linewidth=2)
    for time in (flow_time.start, flow_time.end):
        axes[1].axvline(time, linestyle="--", color="teal")
    axes[1].axvline(flow_time.peak, linestyle="--", color="m")
    axes[1].set(xlabel="MATLAB serial date", ylabel="Discharge (m3/s)")
    return axes


# MATLAB-compatible names for users following the original documentation.
smoothcurve = smooth_curve
findTP = find_turning_points
separatebaseflow = separate_baseflow
extractrunoff = extract_runoff
extractprecipevent = extract_precipitation_events
computehydroITC = compute_hydrograph_time_instants
computehyetoITC = compute_hyetograph_time_instants
matchrainfallrunoff = match_rainfall_runoff
computeTC = compute_time_characteristics
computeRR = compute_runoff_ratio
computeIA = compute_initial_abstraction
computeAP = compute_antecedent_precipitation
computeRC = compute_recession_constant
generalizerecession = generalize_recession
batchprocessing = batch_process
plotrunoffevent = plot_runoff_events
plotMRC = plot_master_recession_curve
plotevent = plot_event