""" Pinball (quantile) loss -- the actual fix for the failure this project just measured directly: a model trained on plain MSE caught zero real flood events at every lead time (0/203 real events, real held-out test data), not because the architecture was wrong, but because MSE minimizes AVERAGE error, and real discharge is heavily right-skewed (confirmed directly against real data: 99th percentile running 4-12x the median). Predicting close to the mean is a genuinely good strategy for average error and a genuinely bad one for flood risk. Pinball loss for a quantile tau: L_tau(y, yhat) = max(tau * (y - yhat), (tau - 1) * (y - yhat)) This is asymmetric by construction -- for a high tau (e.g. 0.95), underpredicting is penalized far more than overpredicting by the same amount, which is exactly the incentive missing from MSE: a model trained with this loss is pushed toward actually covering the real tail, not toward the safe average. """ from typing import List, Optional import numpy as np try: import torch _HAS_TORCH = True except ImportError: _HAS_TORCH = False def pinball_loss(y_true, y_pred, quantile: float): """ Single-quantile pinball loss, NaN-masked (same principle used throughout this project's other losses -- a missing real target must not poison the loss for every other real value in the batch). Args: y_true, y_pred: same shape, any real tensor/array. quantile: the quantile level this prediction represents, in (0, 1). """ is_np = not _HAS_TORCH or not isinstance(y_true, torch.Tensor) abs_diff = y_true - y_pred loss = np.maximum(quantile * abs_diff, (quantile - 1) * abs_diff) if is_np else \ torch.maximum(quantile * abs_diff, (quantile - 1) * abs_diff) is_finite = np.isfinite(loss) if is_np else torch.isfinite(loss) has_any = is_finite.any() if hasattr(is_finite, "any") else any(is_finite) if not has_any: return loss.sum() * float("nan") if not is_np else np.nan return loss[is_finite].mean() def multi_quantile_loss(y_true, y_pred_quantiles, quantiles: List[float]): """ Sums pinball loss across every predicted quantile level -- what a caller actually uses during training, since the model predicts all quantiles for a target together, not one at a time. Args: y_true: [*batch_dims] -- the real target, same for every quantile. y_pred_quantiles: [*batch_dims, n_quantiles] -- one prediction per quantile level, last axis ordered to match `quantiles`. quantiles: the quantile levels, e.g. [0.5, 0.9, 0.95, 0.99]. """ # Confirmed as a real, persistent bug against real training, not # hypothetical: y_true.sum() * 0.0 was meant as "a zero of the # right shape/dtype/device", but y_true (the real target) is # ~97% NaN by design here -- only a few real gauges ever have real # discharge/water-level data. A plain (non-nan-aware) sum over a # mostly-NaN tensor is ALREADY NaN, and NaN * 0.0 is still NaN, not # zero. total started broken and never recovered, regardless of how # many real pinball terms the loop below correctly computed -- # NaN + anything is NaN. This made multi_quantile_loss return NaN # on every real call in this project, unconditionally, masked until # now by larger, more visible bugs elsewhere in the same runs. total = torch.zeros((), dtype=y_true.dtype, device=y_true.device) if _HAS_TORCH and isinstance(y_true, torch.Tensor) \ else np.zeros((), dtype=y_true.dtype) n_real_terms = 0 for i, q in enumerate(quantiles): term = pinball_loss(y_true, y_pred_quantiles[..., i], q) is_nan = np.isnan(term) if not _HAS_TORCH or not isinstance(term, torch.Tensor) else torch.isnan(term) if not is_nan: total = total + term n_real_terms += 1 if n_real_terms == 0: return total * float("nan") return total / n_real_terms def log_space_multi_quantile_loss( y_true_standardized, y_pred_quantiles_standardized, quantiles: List[float], real_mean: float, real_std: float, log_offset: Optional[float] = None, sample_weight=None, ): """ Same real pinball-loss mechanism as multi_quantile_loss, but computed in log(real discharge + log_offset) space instead of standardized (z-scored) space -- built specifically to address a real, measured problem: real evaluation showed high-flow PBIAS at 25.0% vs. -1.8% at low flow for one real station, confirmed as peak-specific, not a general low-flow issue. In real, linear units, a large absolute error during a real 30,000 L/s peak swamps a proportionally IDENTICAL relative error during a real 500 L/s baseline day, even under pinball loss's linear (not squared) penalty -- log-space compresses that gap so a given real PROPORTIONAL error contributes comparable loss regardless of whether it happens at low or high flow, directly targeting the confirmed peak-specific bias. log_offset MATTERS, CONFIRMED DIRECTLY: the original log1p (offset fixed at 1.0) turned out to weight a low-flow error ~28x more than an identical high-flow error (measured on real station-scale values), and a real training run with that offset produced a real, severe regression -- low-flow NSE collapsed to -116.7, PBIAS to -102.7%, at the same real station whose high-flow bias this was built to fix. A larger, real, data-derived offset softens that ratio (measured directly: 0.25*discharge_std -> ~10x, not eliminating the low-flow priority, but no longer as extreme). log_offset defaults to None, which resolves to 0.25*real_std here -- a real, tested starting point, not the original, now-known- risky offset=1.0 default. HONEST CAVEAT: this addresses one real, understood mechanism (log-space's near-zero sensitivity) -- it does NOT fully explain why kge_loss_from_pooled_pairs, a genuinely different mechanism, produced a similar real low-flow regression independently. Treat this as a real, testable improvement, not a confirmed, complete fix, until re-evaluated the same stratified-low/high-flow way the original regression was found. SCOPED TO THIS FUNCTION ONLY -- deliberately not a project-wide change. confluence_mass_balance_loss and water_balance_loss need real, ADDITIVE, linear-space discharge (Q_downstream ~= sum(Q_upstream)); log(a)+log(b) != log(a+b), so log-transforming discharge throughout this project would break the physical meaning of every real mass-conservation constraint already built and validated. Everything else keeps operating on real, linear discharge exactly as before; only the supervised term changes. Args: y_true_standardized: real target, in the SAME standardized (z-scored) space the model already operates in -- this function does the un-standardization internally, the caller doesn't need to pass real, un-standardized values. y_pred_quantiles_standardized: [..., n_quantiles], same standardized space. quantiles: quantile levels, e.g. [0.5, 0.9, 0.95, 0.99]. real_mean, real_std: this run's real discharge_mean/ discharge_std (from standardize_dynamic_tensors, fit on the real training period only) -- needed to recover real, physical L/s values before the log transform, since log has no valid meaning applied directly to a z-score. log_offset: the real "c" in log(x + c). None resolves to 0.25*real_std. Larger values soften low-flow sensitivity (real, measured trade-off, not a free parameter -- see this docstring's own measured ratios). sample_weight: optional, same meaning as weighted_multi_quantile_loss's own sample_weight -- e.g. train_spatiotemporal_gnn.py's --event-reweight combines with log-space loss by passing the same event-proximity weight tensor through here. None (the default) reproduces the original, unweighted behavior exactly. """ is_torch = _HAS_TORCH and isinstance(y_true_standardized, torch.Tensor) if log_offset is None: log_offset = 0.25 * real_std y_true_real = y_true_standardized * real_std + real_mean y_pred_real = y_pred_quantiles_standardized * real_std + real_mean # Real discharge is never negative; log is undefined at/below # -log_offset. This clamp is NOT the sole force keeping predictions # non-negative the way Manning's n's old hard clamp was (a single, # shared, permanently-reused scalar that could get stuck outside # its range) -- the pinball loss itself, even computed in log- # space, naturally pulls every real prediction toward real, non- # negative observed discharge, and each (node, horizon, quantile) # value here comes from an independent forward pass, so an # occasional real clamp doesn't create the same permanent-zero- # gradient trap. if is_torch: y_true_log = torch.log(torch.clamp(y_true_real, min=0.0) + log_offset) y_pred_log = torch.log(torch.clamp(y_pred_real, min=0.0) + log_offset) else: y_true_log = np.log(np.clip(y_true_real, 0.0, None) + log_offset) y_pred_log = np.log(np.clip(y_pred_real, 0.0, None) + log_offset) if sample_weight is None: return multi_quantile_loss(y_true_log, y_pred_log, quantiles) return weighted_multi_quantile_loss(y_true_log, y_pred_log, quantiles, sample_weight) def per_station_relative_multi_quantile_loss( y_true_standardized, y_pred_quantiles_standardized, quantiles: List[float], real_mean: float, real_std: float, per_node_std, sample_weight=None, ): """ Same real pinball-loss mechanism as multi_quantile_loss, but normalizes each real node's real discharge by THAT NODE's OWN real historical std before computing pinball loss, instead of the single, GLOBAL discharge_std every other version of this loss uses. Built after real, direct evidence (check_discharge_quantile_distribution.py) showed this project's 6 real gauges have WILDLY different real scale and skew -- Q50/Q01 ranged from 1.4x at one real station (nearly flat at low flow) to literally infinite (real zero) at another, and Q99/Q50 ranged from 4.3x to 12.0x. Under the existing GLOBAL standardization, a real, proportionally large error at a small station (e.g. real Q50~868 L/s) contributes a much smaller loss than the identical real PROPORTIONAL error at a large station (e.g. real Q50~9323 L/s), simply because both get divided by the same global discharge_std -- the loss implicitly favors getting the large stations right, not a deliberate, reasoned choice. Dividing by each node's own real per_node_std puts every real station's discharge into a comparable, station-relative unit -- a real error that's proportionally severe for a given station contributes comparable loss regardless of whether that station is large or small. SCOPED TO THIS FUNCTION ONLY -- deliberately not a project-wide change, same discipline as log_space_multi_quantile_loss. confluence_mass_balance_loss and water_balance_loss need real, ADDITIVE, linear-space discharge in the ORIGINAL global scale (Q_downstream ~= sum(Q_upstream)); per-node-relative values aren't additive across nodes the way real physical discharge is, so this transform would break the physical meaning of every real mass-conservation constraint if applied there. Everything else keeps operating on real, linear, globally-standardized discharge exactly as before; only this one supervised term changes. Args: y_true_standardized: real target, standardized (z-scored) space, [n_nodes, n_horizons] or similar -- last dim broadcasts against per_node_std via its own leading node dimension. y_pred_quantiles_standardized: [..., n_quantiles], same standardized space, same leading node dimension. quantiles: quantile levels, e.g. [0.5, 0.9, 0.95, 0.99]. real_mean, real_std: this run's real, GLOBAL discharge_mean/ discharge_std (from standardize_dynamic_tensors) -- needed to recover real, physical L/s values first. per_node_std: [n_nodes] real, per-node historical std (from compute_per_node_historical_std, fit on the training period only) -- real NaN for any node with fewer than 2 real training-period observations, which correctly propagates to NaN loss for that node's terms, excluded by multi_quantile_loss's own real NaN-masking, the same way a NaN per_node_mean already gets handled for NSE. sample_weight: optional, same meaning as weighted_multi_quantile_loss's own sample_weight. None (the default) reproduces the original, unweighted behavior exactly. """ is_torch = _HAS_TORCH and isinstance(y_true_standardized, torch.Tensor) y_true_real = y_true_standardized * real_std + real_mean y_pred_real = y_pred_quantiles_standardized * real_std + real_mean # Broadcast per_node_std ([n_nodes]) against y_true_real's # [n_nodes, n_horizons, ...] shape -- reshape to add trailing # singleton dims matching however many extra dims y_true_real has # beyond the leading node dimension. # # CRITICAL, REAL, CONFIRMED FIX: replace NaN entries in per_node_std # (real, ungauged/virtual nodes with no real training-period # discharge at all -- compute_per_node_historical_std correctly # returns NaN for these) with a safe, finite placeholder BEFORE the # division, not after. A real training run confirmed this is NOT # optional: dividing by a real NaN produces a NaN y_pred_relative at # those real node positions, and even though multi_quantile_loss's # own real masking correctly excludes those NaN positions from the # FORWARD loss sum, PyTorch's backward pass still computes # d(NaN)/d(input) for every real parameter that touched that # division -- NaN * 0 is NaN, not 0, in gradient computation, a # real, well-known autograd pitfall. Confirmed directly: a real # training run showed EVERY SINGLE model parameter with a real NaN # gradient from epoch 1, batch 1 onward, tracing back to exactly # this division. The placeholder value itself doesn't matter for # correctness (any finite, positive number works, since the # corresponding real Y_true is ALSO NaN at these same node # positions and gets excluded from the forward sum regardless) -- # 1.0 is used here for simplicity and because it can never itself # be zero. if is_torch: safe_per_node_std = torch.where(torch.isnan(per_node_std), torch.ones_like(per_node_std), per_node_std) else: safe_per_node_std = np.where(np.isnan(per_node_std), 1.0, per_node_std) if is_torch: true_extra_dims = y_true_real.dim() - 1 pred_extra_dims = y_pred_real.dim() - 1 node_std_true = safe_per_node_std.reshape(-1, *([1] * true_extra_dims)) node_std_pred = safe_per_node_std.reshape(-1, *([1] * pred_extra_dims)) else: true_extra_dims = y_true_real.ndim - 1 pred_extra_dims = y_pred_real.ndim - 1 node_std_true = safe_per_node_std.reshape(-1, *([1] * true_extra_dims)) node_std_pred = safe_per_node_std.reshape(-1, *([1] * pred_extra_dims)) y_true_relative = y_true_real / node_std_true y_pred_relative = y_pred_real / node_std_pred if sample_weight is None: return multi_quantile_loss(y_true_relative, y_pred_relative, quantiles) return weighted_multi_quantile_loss(y_true_relative, y_pred_relative, quantiles, sample_weight) def weighted_multi_quantile_loss(y_true, y_pred_quantiles, quantiles: List[float], sample_weight, alpha: float = 1.0): """ Same pinball-loss sum as multi_quantile_loss, but each element is additionally weighted before averaging -- a generic weighted pinball average, not tied to any one specific weighting scheme. Two real, distinct uses in this project so far: (1) Recency weighting (the original use this was built for): sample_weight = 1 + alpha*log(1+dt) per node, forcing more loss-weight onto observations following a longer gap in that node's own history. Overlap warning: a days-since-last-real-observation channel is already wired in as a direct input feature (dynamic_features.py's compute_days_since_last_real / compress_days_since_last_real), which lets the model learn how to use staleness itself. This function instead forces more loss-weight onto post-gap examples regardless of whether the model finds that useful -- two different mechanisms carrying the same underlying signal, not guaranteed to be aligned. (2) Event-proximity weighting (train_spatiotemporal_gnn.py's --event-reweight): sample_weight = build_event_weight_tensor's output, boosting examples near a real labeled flood event. Built to directly counteract real, measured sample-count dilution -- real flood events are rare (125 across the whole 2013-2025 record), so the sheer VOLUME of quiet-day gradient updates dominates training even though pinball loss's own per-sample asymmetry already weights a flood-day underprediction more heavily than a same-size quiet-day overprediction; that per-sample weighting doesn't fix a per-example-COUNT imbalance the way this does. Treat sample_weight of all-1.0 (a no-op, exactly reproducing multi_quantile_loss) as the baseline to compare any real weighting scheme against, rather than assuming either one helps by default. Args: y_true, y_pred_quantiles, quantiles: same as multi_quantile_loss. sample_weight: either [n_nodes] (broadcast against y_true's leading node dimension) or full-shaped (same shape as y_true, e.g. [n_nodes, n_horizons]) -- already-computed per-element weights. alpha: scaling on the log(1+dt) term for the recency-weighting use case, applied by the caller when building that sample_weight -- kept here only for direct-call convenience; has no effect if sample_weight is pre-built (as it always is for event-proximity weighting). """ is_torch = _HAS_TORCH and isinstance(y_true, torch.Tensor) weight_bcast = sample_weight if (hasattr(sample_weight, "ndim") and sample_weight.ndim == y_true.ndim) else \ (sample_weight.view(-1, *([1] * (y_true.ndim - 1))) if is_torch else sample_weight.reshape(-1, *([1] * (y_true.ndim - 1)))) total = torch.zeros((), dtype=y_true.dtype, device=y_true.device) if is_torch \ else np.zeros((), dtype=y_true.dtype) total_weight = torch.zeros((), dtype=y_true.dtype, device=y_true.device) if is_torch \ else np.zeros((), dtype=y_true.dtype) for i, q in enumerate(quantiles): y_pred_i = y_pred_quantiles[..., i] abs_diff = y_true - y_pred_i raw = torch.maximum(q * abs_diff, (q - 1) * abs_diff) if is_torch else \ np.maximum(q * abs_diff, (q - 1) * abs_diff) weighted = raw * (weight_bcast if is_torch else np.broadcast_to(weight_bcast, raw.shape)) is_finite = torch.isfinite(weighted) if is_torch else np.isfinite(weighted) has_any = is_finite.any() if hasattr(is_finite, "any") else any(is_finite) if not has_any: continue w_for_sum = (weight_bcast if is_torch else np.broadcast_to(weight_bcast, raw.shape))[is_finite] total = total + weighted[is_finite].sum() total_weight = total_weight + w_for_sum.sum() if (total_weight == 0) if not is_torch else (total_weight.item() == 0): return total * float("nan") return total / total_weight