Spaces:
Sleeping
Sleeping
| """ | |
| 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 |