Spaces:
Running on Zero
Running on Zero
| """ | |
| The main spatiotemporal model: predicts discharge/water level at ANY | |
| node (gauged or ungauged) for a requested horizon, using the real reach | |
| graph topology + static/dynamic features built throughout this project. | |
| Architecture, per the strategy discussion: | |
| - Spatial: directed message passing over the real reach graph, | |
| respecting real flow direction. SUM aggregation at each node | |
| (not mean) -- deliberate, not a default left unconsidered: a | |
| confluence's downstream discharge is the SUM of its upstream | |
| branches (see physics_losses.py's confluence_mass_balance_loss), | |
| so the message-passing aggregation itself should reflect that | |
| physical reality, not average it away. | |
| - Temporal: per-node GRU over the lookback window, fed by the | |
| spatial layer's output at each timestep (not raw features | |
| directly) -- matches the "Option A" architecture decided early in | |
| this project (spatial + temporal kept separate) for sample | |
| efficiency given how sparse real labels are. | |
| - Decoder: horizon-conditioned, SHARED weights across both basins, | |
| with a learned basin embedding as an additional input -- one | |
| backbone for the general "how does water move through a river | |
| network" relationships, basin-specific behavior expressed only at | |
| the final decoding stage. Quantile output (Q, H), not a point | |
| estimate -- a plain-MSE-trained earlier version caught 0 of 203 | |
| real flood events, since MSE optimizes toward the distribution's | |
| center, exactly wrong for flood risk. | |
| """ | |
| from typing import Dict, List, Optional, Tuple | |
| import numpy as np | |
| import pandas as pd | |
| try: | |
| import torch | |
| import torch.nn as nn | |
| from torch_geometric.nn import MessagePassing | |
| _HAS_TORCH = True | |
| except ImportError: | |
| _HAS_TORCH = False | |
| nn = None | |
| MessagePassing = object # placeholder base so class definitions below don't NameError at import time | |
| # --------------------------------------------------------------------------- | |
| # Data preparation -- pure numpy/pandas, testable without torch | |
| # --------------------------------------------------------------------------- | |
| def prepare_graph_training_windows( | |
| x_static: np.ndarray, dynamic_tensors: Dict[str, np.ndarray], target_tensor: np.ndarray, | |
| dates: List[pd.Timestamp], lookback_days: int, horizons: List[int], | |
| dt_tensors: Optional[Dict[str, np.ndarray]] = None, | |
| per_node_fill_values: Optional[Dict[str, np.ndarray]] = None, | |
| forward_fill_tensors: Optional[Dict[str, np.ndarray]] = None, | |
| ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, List[pd.Timestamp]]: | |
| """ | |
| Builds graph-wide training windows -- genuinely different from | |
| prepare_training_windows in climate_anomaly_model.py, which could | |
| require an ENTIRE lookback window to be real per (independent) cell | |
| and just skip incomplete ones. That doesn't work at graph scale: | |
| only 6-8 of ~4,500 nodes have ANY real discharge at all, so | |
| requiring every node's window to be fully real would throw away | |
| virtually the entire graph on every single anchor date. | |
| Instead: missing dynamic values are filled with 0 (after the | |
| caller's own z-scoring, so 0 represents "no information", not a | |
| physically meaningful reading) and paired with an explicit | |
| missingness channel per variable -- same __was_missing principle | |
| already used throughout node_features.py's static features, now | |
| applied per-timestep. The model sees BOTH the (filled) value and | |
| whether it was real, and message passing is what's actually | |
| supposed to let real information at gauged nodes propagate outward | |
| to fill the gap physically, rather than the raw input pretending | |
| a filled zero is a real observation. | |
| CRITICAL, CONFIRMED FIX -- per_node_fill_values: plain 0.0 is | |
| standardized using the GLOBAL mean/std across every real gauge | |
| together, so it's only genuinely "neutral" for a node whose own | |
| real scale happens to sit near that global average. Confirmed | |
| directly as a real, severe bug, not a theoretical concern: for one | |
| real, small station (real Q50~983 L/s, far below the real global | |
| discharge_mean~4168.71), a real data gap filled the model's input | |
| with a value ~60x larger than that station ever actually sees -- | |
| traced directly via a real false-spike investigation, and this | |
| exact real station showed the worst real low-flow PBIAS measured | |
| all project (-240.1%). The missingness flag alone evidently wasn't | |
| enough for the model to fully discount a fill value that large and | |
| misleading. per_node_fill_values, when provided, gives each real | |
| node its OWN real historical mean (in the same global-standardized | |
| units the rest of the tensor is already in, via | |
| compute_per_node_historical_mean applied to the ALREADY- | |
| standardized tensor -- fit on the real training period only, same | |
| leakage discipline as discharge_mean/std themselves) as its fill | |
| value instead of the single, global 0.0 -- genuinely neutral for | |
| THAT node's own real typical range, not the whole network's. | |
| A real node with no real training-period observations at all | |
| (a true virtual/confluence node, or per_node_fill_values not | |
| provided) correctly falls back to the original global 0.0 -- there | |
| is no real, per-node information to prefer for such a node. | |
| forward_fill_tensors, when provided, takes priority over | |
| per_node_fill_values as the fill source for a missing timestep. | |
| CRITICAL, CONFIRMED FIX #2 -- per_node_fill_values (a station-wide | |
| mean or median) is only a genuinely neutral fill for a node's | |
| TYPICAL day; it is the wrong reference point for a gap that falls | |
| in the middle of a real, ongoing low- (or high-) flow episode. | |
| Confirmed directly: even a real per-node MEDIAN fill (1142.0, for | |
| station H605641401) was ~17x that station's real low-flow values | |
| (57-70 L/s) during a real gap inside a real low-flow episode, and | |
| a direct causal ablation (zeroing the missingness flag and | |
| days-since channel for this exact example) confirmed the model | |
| could not fully compensate: predicted Q0.95 was still ~28x the | |
| real observed value even with no "this was filled" signal reaching | |
| it at all, because the filled VALUE itself was implausible for the | |
| real local conditions at that moment -- not a flag-weighting | |
| problem. forward_fill_tensors (see dynamic_features.py's | |
| compute_forward_filled_tensor) instead carries the real, | |
| most-recently-known value at that specific node forward, which is | |
| directly informative about real LOCAL, CURRENT conditions in a way | |
| a station-wide statistic structurally cannot be. Needs no | |
| train-only-fit discipline (unlike per_node_fill_values) since | |
| forward-fill only ever looks backward in time -- no leakage is | |
| possible by construction. | |
| Fill priority per missing timestep, per variable: | |
| 1. forward_fill_tensors[var] at that [node, date] position, if | |
| provided AND real (i.e. some real observation has occurred at | |
| that node before that date). | |
| 2. per_node_fill_values[var] for that node, if provided AND real | |
| (station-wide fallback when there's no real prior observation | |
| yet for forward-fill to carry forward -- e.g. the very start | |
| of a node's real record). | |
| 3. global 0.0 (original fallback, for a node with no real | |
| information at all from either source). | |
| dt_tensors, when provided, adds a THIRD channel per variable: real, | |
| log-compressed days-since-last-real-observation (see | |
| dynamic_features.py's compute_days_since_last_real). A binary | |
| missing flag alone can't distinguish a single-day gap from a | |
| months-long one; this is the fix. Must cover every key in | |
| dynamic_tensors if provided at all -- a partial dt_tensors dict | |
| would silently give some variables 2 channels and others 3, | |
| genuinely breaking the interleaving downstream code relies on, so | |
| this is enforced explicitly rather than left to fail confusingly | |
| later. | |
| Args: | |
| x_static: [n_nodes, n_static_features] -- constant across the | |
| whole window (graph structure/static features don't change | |
| over the lookback period). | |
| dynamic_tensors: {variable_name: [n_nodes, T]} -- e.g. from | |
| dynamic_features.py's assemble_dynamic_tensor, one entry | |
| per dynamic variable (discharge, groundwater level, etc.) | |
| target_tensor: [n_nodes, T] -- the real quantity being | |
| predicted (discharge or water level), same real/NaN | |
| structure as the dynamic tensors (only real gauges with | |
| real observations have a value). | |
| dates: the T dates target_tensor/dynamic_tensors are indexed by. | |
| lookback_days, horizons: as elsewhere in this project. | |
| dt_tensors: optional {variable_name: [n_nodes, T]}, already | |
| log-compressed, matching dynamic_tensors' own keys exactly. | |
| Returns: | |
| X_dynamic: [n_examples, lookback_days, n_nodes, C] -- interleaved | |
| per dynamic variable, (value, was_missing) if dt_tensors is | |
| None, (value, was_missing, days_since_real) if provided -- | |
| C = 2*n_dynamic_vars or 3*n_dynamic_vars accordingly. | |
| Y: [n_examples, n_nodes, n_horizons] -- target values, NaN where | |
| no real observation exists (real gauges only, at whatever | |
| horizons they actually have data for). | |
| Y_mask: [n_examples, n_nodes, n_horizons] boolean, True where Y | |
| is real -- returned explicitly rather than relying on | |
| np.isnan downstream, since the loss function needs exactly | |
| this same mask and computing it twice invites the two | |
| copies drifting apart. | |
| anchor_dates: the "today" date for each example. | |
| """ | |
| var_names = list(dynamic_tensors.keys()) | |
| if dt_tensors is not None: | |
| missing_dt_vars = set(var_names) - set(dt_tensors.keys()) | |
| if missing_dt_vars: | |
| raise ValueError(f"dt_tensors is missing entries for {missing_dt_vars} -- must cover " | |
| f"every key in dynamic_tensors, or be None entirely, not a partial set.") | |
| channels_per_var = 3 if dt_tensors is not None else 2 | |
| n_nodes = x_static.shape[0] | |
| n_dates = len(dates) | |
| max_horizon = max(horizons) | |
| X_list, Y_list, anchor_dates = [], [], [] | |
| for t in range(lookback_days, n_dates - max_horizon): | |
| window_channels = [] | |
| for var in var_names: | |
| raw_window = dynamic_tensors[var][:, t - lookback_days:t] # [n_nodes, lookback] | |
| missing = np.isnan(raw_window) | |
| # Fill priority: forward_fill_tensors (real, local, | |
| # most-recent value at THIS node) > per_node_fill_values | |
| # (station-wide mean/median, only when forward-fill has no | |
| # real prior observation yet to carry) > global 0.0 (no | |
| # real information at all). See this function's docstring | |
| # ("CRITICAL, CONFIRMED FIX #2") for why forward-fill takes | |
| # priority over the station-wide statistic. | |
| if forward_fill_tensors is not None and var in forward_fill_tensors: | |
| ff_window = forward_fill_tensors[var][:, t - lookback_days:t] # [n_nodes, lookback] | |
| if per_node_fill_values is not None and var in per_node_fill_values: | |
| node_fill = per_node_fill_values[var] | |
| node_fill_broadcast = np.where( | |
| np.isnan(node_fill), 0.0, node_fill, | |
| ).reshape(-1, 1) | |
| stat_fallback = np.broadcast_to(node_fill_broadcast, ff_window.shape) | |
| else: | |
| stat_fallback = np.zeros_like(ff_window) | |
| # ff_window is NaN only where this node has no real | |
| # observation yet anywhere before that date -- fall | |
| # back to the station-wide stat (or 0.0) only there. | |
| fill_source = np.where(np.isnan(ff_window), stat_fallback, ff_window) | |
| filled = np.where(missing, fill_source, raw_window) | |
| elif per_node_fill_values is not None and var in per_node_fill_values: | |
| # Real, per-node fill values -- [n_nodes], broadcast | |
| # across the lookback dimension. A node whose own real | |
| # fill value is itself NaN (no real training-period | |
| # observations at all for this variable at this node) | |
| # correctly falls back to the original, global 0.0 -- | |
| # np.where below prefers the per-node value only where | |
| # it's real. | |
| node_fill = per_node_fill_values[var] | |
| node_fill_broadcast = np.where( | |
| np.isnan(node_fill), 0.0, node_fill, | |
| ).reshape(-1, 1) | |
| filled = np.where(missing, node_fill_broadcast, raw_window) | |
| else: | |
| filled = np.where(missing, 0.0, raw_window) | |
| window_channels.append(filled) | |
| window_channels.append(missing.astype(np.float32)) | |
| if dt_tensors is not None: | |
| dt_window = dt_tensors[var][:, t - lookback_days:t] | |
| window_channels.append(dt_window) | |
| # stack into [lookback, n_nodes, channels_per_var*n_vars] | |
| stacked = np.stack(window_channels, axis=-1).transpose(1, 0, 2) | |
| targets = np.stack([target_tensor[:, t + h - 1] for h in horizons], axis=-1) # [n_nodes, n_horizons] | |
| if np.isnan(targets).all(): | |
| continue # no real supervision anywhere in the graph for this anchor date -- physics losses | |
| # still apply at training time even without this window, so skipping here only affects | |
| # the supervised-loss examples, not physics-loss coverage (which the training loop applies | |
| # to every anchor date's model output, not just ones with real targets) | |
| X_list.append(stacked) | |
| Y_list.append(targets) | |
| anchor_dates.append(dates[t]) | |
| if not X_list: | |
| return (np.empty((0, lookback_days, n_nodes, channels_per_var * len(var_names))), | |
| np.empty((0, n_nodes, len(horizons))), | |
| np.empty((0, n_nodes, len(horizons)), dtype=bool), []) | |
| X = np.stack(X_list) | |
| Y = np.stack(Y_list) | |
| Y_mask = ~np.isnan(Y) | |
| return X, Y, Y_mask, anchor_dates | |
| # --------------------------------------------------------------------------- | |
| # Model -- requires torch + torch_geometric, NOT executed in this | |
| # environment. Syntax-checked only. | |
| # --------------------------------------------------------------------------- | |
| if _HAS_TORCH: | |
| class RiverMessagePassing(MessagePassing): | |
| """ | |
| One round of message passing over the real reach graph. | |
| aggr='add' (SUM), not 'mean' -- deliberate: a confluence's | |
| downstream value should reflect the SUM of what flows in from | |
| each upstream branch, matching physics_losses.py's | |
| confluence_mass_balance_loss directly, not an average that | |
| would understate a real confluence's combined inflow. | |
| Edge features (distance_km, elevation_drop_m, | |
| verified_continuous) condition the message, not just the | |
| source node's features -- a longer, steeper edge should | |
| propagate differently than a short, flat one, and | |
| verified_continuous=False (the bétoire stretch) is exactly the | |
| kind of signal that should suppress how strongly water is | |
| "assumed" to pass through that edge. | |
| """ | |
| def __init__(self, in_channels: int, out_channels: int, edge_dim: int): | |
| super().__init__(aggr="add") | |
| self.edge_mlp = nn.Sequential(nn.Linear(edge_dim, out_channels), nn.ReLU()) | |
| self.node_mlp = nn.Linear(in_channels, out_channels) | |
| # update() concatenates the ORIGINAL input x (width=in_channels) | |
| # with aggr_out (width=out_channels) -- this must be sized | |
| # in_channels + out_channels, not out_channels * 2. Confirmed as | |
| # a real, not hypothetical, bug: out_channels*2 happened to be | |
| # correct for spatial2 (in_channels == out_channels == hidden_dim | |
| # there, by coincidence) but broke immediately on real data for | |
| # spatial1, where in_channels (static_dim+dynamic_dim=7) != | |
| # out_channels (hidden_dim=64) -- real error: | |
| # "mat1 and mat2 shapes cannot be multiplied (4525x71 and 128x64)", | |
| # 71 = 7 + 64, the actual concatenated width, not 128 = 64*2. | |
| self.update_mlp = nn.Sequential( | |
| nn.Linear(in_channels + out_channels, out_channels), nn.ReLU(), | |
| ) | |
| def forward(self, x: "torch.Tensor", edge_index: "torch.Tensor", edge_attr: "torch.Tensor") -> "torch.Tensor": | |
| return self.propagate(edge_index, x=x, edge_attr=edge_attr) | |
| def message(self, x_j: "torch.Tensor", edge_attr: "torch.Tensor") -> "torch.Tensor": | |
| edge_feat = self.edge_mlp(edge_attr) | |
| return self.node_mlp(x_j) * edge_feat | |
| def update(self, aggr_out: "torch.Tensor", x: "torch.Tensor") -> "torch.Tensor": | |
| return self.update_mlp(torch.cat([x, aggr_out], dim=-1)) | |
| class TypeAwareInputEncoder(nn.Module): | |
| """ | |
| Gauge and non-gauge nodes get genuinely separate input | |
| processing -- not the same shared encoder applied to a 0-filled | |
| fabrication for nodes that never had the data. Two encoders: | |
| - shared_encoder: processes features every node genuinely has | |
| (elevation, IDPR, cumulative catchment area) -- computed | |
| identically for every node, no type distinction needed here | |
| because the data itself isn't type-specific. | |
| - gauge_encoder: processes features that only exist for real | |
| gauges (landcover, geology, NDVI, cavite proximity, real | |
| discharge/groundwater history). Technically evaluated on | |
| EVERY node's input tensor (including non-gauge nodes' 0-filled | |
| placeholders, for simple batched tensor shapes), but the | |
| result is then REPLACED for non-gauge nodes by a single | |
| learned placeholder embedding via torch.where -- gradients | |
| only flow to gauge_encoder's parameters through real gauge | |
| nodes (torch.where's backward routes gradient to whichever | |
| branch was actually selected), so gauge-only parameters never | |
| get spurious updates from fabricated non-gauge inputs, even | |
| though the forward pass computes (and discards) that encoding | |
| for every node. This replaces the model having to LEARN to | |
| discount a "0 + missingness flag" signal with an architectural | |
| guarantee that non-gauge nodes never see gauge-encoder output | |
| presented as if it were real. | |
| """ | |
| def __init__(self, shared_static_dim: int, gauge_input_dim: int, embed_dim: int): | |
| super().__init__() | |
| self.shared_encoder = nn.Linear(shared_static_dim, embed_dim) | |
| self.gauge_encoder = nn.Sequential(nn.Linear(gauge_input_dim, embed_dim), nn.ReLU()) | |
| self.non_gauge_placeholder = nn.Parameter(torch.zeros(embed_dim)) | |
| def forward( | |
| self, x_shared_static: "torch.Tensor", x_gauge_static: "torch.Tensor", | |
| x_gauge_dynamic_t: "torch.Tensor", is_gauged: "torch.Tensor", | |
| ) -> "torch.Tensor": | |
| """ | |
| Args: | |
| x_shared_static: [n_nodes, shared_static_dim] -- constant across the window. | |
| x_gauge_static: [n_nodes, gauge_static_dim] -- constant across the window, | |
| real only for gauge nodes (0-filled elsewhere). | |
| x_gauge_dynamic_t: [n_nodes, gauge_dynamic_dim] -- THIS timestep's | |
| gauge-only dynamic history (real only for gauge nodes). | |
| is_gauged: [n_nodes] bool tensor. | |
| Returns: [n_nodes, embed_dim] combined node representation for this timestep. | |
| """ | |
| shared_embed = self.shared_encoder(x_shared_static) | |
| gauge_input = torch.cat([x_gauge_static, x_gauge_dynamic_t], dim=-1) | |
| gauge_embed_computed = self.gauge_encoder(gauge_input) | |
| gauge_embed = torch.where( | |
| is_gauged.unsqueeze(-1), gauge_embed_computed, | |
| self.non_gauge_placeholder.unsqueeze(0).expand_as(gauge_embed_computed), | |
| ) | |
| return shared_embed + gauge_embed | |
| class _CausalTemporalCNN(nn.Module): | |
| """ | |
| Real, causal, dilated 1D convolution stack over the lookback | |
| window -- an alternative temporal encoder for the real, direct | |
| GRU-vs-LSTM-vs-CNN ablation this project is building (same | |
| graph layers, same physics losses, same everything else; | |
| temporal_type is the only thing that changes). A genuinely | |
| different mechanism from a GRU/LSTM (a fixed, dilated | |
| receptive field vs. recurrent memory), not a near-duplicate. | |
| CAUSALITY IS THE REAL, LOAD-BEARING PROPERTY HERE: a GRU/LSTM's | |
| forward-only recurrence guarantees a given timestep's hidden | |
| state never depends on future inputs. A naive Conv1d with | |
| symmetric padding does NOT have that property -- it would let | |
| future timesteps leak into the "final" summary this class | |
| produces, which is exactly the kind of real look-ahead bug this | |
| project has been careful to verify against everywhere else. | |
| Fixed here via explicit LEFT-only (causal) padding: pad by | |
| (kernel_size-1)*dilation on the left, none on the right, so | |
| the output at position t only ever convolves over real inputs | |
| at positions <= t. Directly, numerically verified (not just | |
| reasoned about) in the same session this was built. | |
| """ | |
| def __init__(self, in_dim: int, out_dim: int, kernel_size: int = 3, | |
| dilations: Tuple[int, ...] = (1, 2, 4, 8)): | |
| super().__init__() | |
| self.kernel_size = kernel_size | |
| self.dilations = list(dilations) | |
| self.pads = [(kernel_size - 1) * d for d in self.dilations] | |
| convs = [] | |
| channels = in_dim | |
| for dilation in self.dilations: | |
| convs.append(nn.Conv1d(channels, out_dim, kernel_size, dilation=dilation)) | |
| channels = out_dim | |
| self.convs = nn.ModuleList(convs) | |
| def forward(self, seq): | |
| # seq: [n_nodes, lookback, in_dim] -> Conv1d wants | |
| # [n_nodes, in_dim, lookback] (channels-first). | |
| x = seq.transpose(1, 2) | |
| for conv, pad in zip(self.convs, self.pads): | |
| # Explicit LEFT-only pad -- this is the real causality | |
| # mechanism, not Conv1d's own `padding=` argument (which | |
| # pads both sides equally and would leak future steps). | |
| x = nn.functional.pad(x, (pad, 0)) | |
| x = torch.relu(conv(x)) | |
| # [n_nodes, out_dim, lookback] -- the LAST real timestep is | |
| # this stack's causal summary of the whole real sequence up | |
| # to now, matching a GRU/LSTM's own final-hidden-state | |
| # semantics, for a fair, comparable ablation. | |
| return x[:, :, -1] | |
| class SpatiotemporalGNN(nn.Module): | |
| """ | |
| Full model: type-aware input encoding (see TypeAwareInputEncoder | |
| above), two RiverMessagePassing layers per timestep (two hops -- | |
| lets a node see two edges' worth of upstream influence per step, | |
| e.g. a confluence's own upstream confluence), a per-node temporal | |
| encoder across the lookback window (temporal_type: "gru" | |
| (default), "lstm", or "cnn" -- a real, direct ablation of just | |
| this one component, everything else identical), and a | |
| horizon+basin-conditioned decoder with shared weights across | |
| both basins. | |
| """ | |
| def __init__( | |
| self, shared_static_dim: int, gauge_static_dim: int, gauge_dynamic_dim: int, | |
| edge_dim: int = 5, embed_dim: int = 32, hidden_dim: int = 64, gru_hidden: int = 64, | |
| n_basins: int = 2, basin_embed_dim: int = 8, horizon_embed_dim: int = 8, n_targets: int = 2, | |
| quantiles: Optional[List[float]] = None, forecast_embed_dim: int = 8, | |
| temporal_type: str = "gru", | |
| ): | |
| super().__init__() | |
| self.quantiles = quantiles or [0.5, 0.9, 0.95, 0.99] | |
| n_quantiles = len(self.quantiles) | |
| self.input_encoder = TypeAwareInputEncoder( | |
| shared_static_dim, gauge_static_dim + gauge_dynamic_dim, embed_dim, | |
| ) | |
| self.spatial1 = RiverMessagePassing(embed_dim, hidden_dim, edge_dim) | |
| self.spatial2 = RiverMessagePassing(hidden_dim, hidden_dim, edge_dim) | |
| # temporal_type is a real, direct ablation switch -- "gru" | |
| # (default, unchanged behavior), "lstm", or "cnn" | |
| # (_CausalTemporalCNN, defined above). Everything else in | |
| # this model (graph layers, decoder, physics losses used | |
| # during training) stays identical regardless of this | |
| # choice, so a real comparison across the three isolates | |
| # just the temporal-encoding mechanism, not a mix of | |
| # changes. | |
| self.temporal_type = temporal_type | |
| if temporal_type == "gru": | |
| self.temporal = nn.GRU(hidden_dim, gru_hidden, batch_first=True) | |
| elif temporal_type == "lstm": | |
| self.temporal = nn.LSTM(hidden_dim, gru_hidden, batch_first=True) | |
| elif temporal_type == "cnn": | |
| self.temporal = _CausalTemporalCNN(hidden_dim, gru_hidden) | |
| else: | |
| raise ValueError(f"temporal_type must be 'gru', 'lstm', or 'cnn', got {temporal_type!r}") | |
| self.basin_embed = nn.Embedding(n_basins, basin_embed_dim) | |
| self.horizon_embed = nn.Linear(1, horizon_embed_dim) | |
| # Real forecasted precipitation, per node per horizon -- NOT | |
| # part of the temporal encoder, deliberately: a forecast is | |
| # future-looking information tied to a SPECIFIC horizon (the | |
| # forecast for "3 days out" and "30 days out" are genuinely | |
| # different numbers), whereas the encoder only ever sees one | |
| # shared lookback window of the PAST, applied identically | |
| # across every horizon in a single forward pass. Feeding a | |
| # horizon-specific future value into that shared-history | |
| # pathway would misrepresent it as something it isn't. The | |
| # decoder already conditions per-horizon (via horizon_embed), | |
| # so this is where per-horizon future information belongs. | |
| # | |
| # value+missingness pair, same pattern used throughout this | |
| # project for every other real-but-sometimes-absent input -- | |
| # forward() defaults this to "fully missing" when no real | |
| # forecast tensor is passed in, so this is backward | |
| # compatible with every existing training run and doesn't | |
| # require the real forecast dataset to exist yet. | |
| self.forecast_embed = nn.Sequential(nn.Linear(2, forecast_embed_dim), nn.ReLU()) | |
| self.decoder = nn.Sequential( | |
| nn.Linear(gru_hidden + basin_embed_dim + horizon_embed_dim + forecast_embed_dim, hidden_dim), | |
| nn.ReLU(), | |
| nn.Linear(hidden_dim, n_targets * n_quantiles), | |
| ) | |
| self.n_targets = n_targets | |
| # Manning's roughness -- no real measured value in this | |
| # project's data, learned as a single shared scalar. | |
| # Sigmoid-reparameterized (not log-parameterized) so n is | |
| # ALWAYS within [0.01, 0.3] by construction -- replaces an | |
| # earlier log-parameterization + hard clamp that had a | |
| # confirmed, real failure mode: torch.clamp's gradient is | |
| # exactly zero outside its range (verified directly via | |
| # finite differences), so once the raw parameter drifted | |
| # past the clamp boundary, nothing could ever pull it back. | |
| # Confirmed as the real cause of n landing at 0.38-1.0 | |
| # across several real training runs, never converging. | |
| # manning_n_logit=-2.60269 reproduces the same n=0.03 | |
| # starting point the old log-parameterization used. | |
| self.manning_n_logit = nn.Parameter(torch.tensor(-2.60269)) | |
| # Bank height for the Manning's relaxation gate -- no real | |
| # measurement exists, initialized at 2.5m as a physically | |
| # reasonable starting guess for these rivers (not derived | |
| # from data), same log-positivity trick manning_n_logit | |
| # used to use. Left as-is for now: real training runs show | |
| # this consistently landing at 0.10m, not drifting the way | |
| # n did -- no confirmed evidence this one needs the same fix. | |
| self.log_bank_height_m = nn.Parameter(torch.log(torch.tensor(2.5))) | |
| def forward( | |
| self, x_shared_static: "torch.Tensor", x_gauge_static: "torch.Tensor", | |
| x_gauge_dynamic_seq: "torch.Tensor", is_gauged: "torch.Tensor", | |
| edge_index: "torch.Tensor", edge_attr: "torch.Tensor", | |
| basin_id: "torch.Tensor", horizons: "torch.Tensor", | |
| forecast_precip: Optional["torch.Tensor"] = None, | |
| forecast_missing: Optional["torch.Tensor"] = None, | |
| ) -> "torch.Tensor": | |
| """ | |
| Args: | |
| x_shared_static: [n_nodes, shared_static_dim] -- constant across the window. | |
| x_gauge_static: [n_nodes, gauge_static_dim] -- constant across the window. | |
| x_gauge_dynamic_seq: [lookback, n_nodes, gauge_dynamic_dim] -- varies per | |
| timestep (lookback-first, matching prepare_graph_training_windows's | |
| real output shape -- see the note lower in this docstring about why | |
| that axis order specifically was worth calling out here before). | |
| is_gauged: [n_nodes] bool tensor. | |
| edge_index, edge_attr: real graph structure, same for every timestep. | |
| basin_id: [n_nodes] long tensor, 0 or 1. | |
| horizons: [n_horizons] -- requested lead times. | |
| forecast_precip: [n_nodes, n_horizons], optional -- real | |
| forecasted precipitation for each node at each | |
| requested horizon (e.g. from a real ECMWF forecast, | |
| mapped station-to-grid-cell the same way the climate | |
| work already does). None (the default) means no real | |
| forecast dataset has been wired in yet -- the model | |
| falls back to treating every value as missing, which | |
| is exactly today's behavior, so existing training | |
| runs are unaffected by this parameter existing. | |
| forecast_missing: [n_nodes, n_horizons], optional -- | |
| 1.0 where forecast_precip has no real value, 0.0 | |
| where it does. Required if forecast_precip is given | |
| (real missingness must be explicit, not inferred), | |
| ignored if forecast_precip is None. | |
| Returns: | |
| [n_nodes, n_horizons, n_targets, n_quantiles] forecast -- | |
| monotonically increasing along the last axis by | |
| construction (see below), not just intended to be. | |
| """ | |
| lookback = x_gauge_dynamic_seq.shape[0] | |
| embeddings = [] | |
| for t in range(lookback): | |
| x_t = self.input_encoder(x_shared_static, x_gauge_static, x_gauge_dynamic_seq[t, :, :], is_gauged) | |
| h = self.spatial1(x_t, edge_index, edge_attr) | |
| h = self.spatial2(h, edge_index, edge_attr) | |
| embeddings.append(h) | |
| seq = torch.stack(embeddings, dim=1) # [n_nodes, lookback, hidden_dim] | |
| # Real per-temporal_type handling -- GRU returns a single | |
| # final hidden state, LSTM returns a (h_n, c_n) tuple (real | |
| # cell state alongside hidden state, an LSTM-specific | |
| # difference from GRU), CNN (_CausalTemporalCNN) already | |
| # returns the right [n_nodes, gru_hidden] shape directly, no | |
| # further unpacking needed. All three converge to the same | |
| # real node_hidden shape the decoder below expects. | |
| if self.temporal_type == "gru": | |
| _, h_n = self.temporal(seq) | |
| node_hidden = h_n.squeeze(0) # [n_nodes, gru_hidden] | |
| elif self.temporal_type == "lstm": | |
| _, (h_n, c_n) = self.temporal(seq) | |
| node_hidden = h_n.squeeze(0) # [n_nodes, gru_hidden] | |
| else: # "cnn" | |
| node_hidden = self.temporal(seq) # [n_nodes, gru_hidden] | |
| basin_feat = self.basin_embed(basin_id) # [n_nodes, basin_embed_dim] | |
| n_nodes = node_hidden.shape[0] | |
| n_horizons = horizons.shape[0] | |
| node_hidden_exp = node_hidden.unsqueeze(1).expand(n_nodes, n_horizons, -1) | |
| basin_feat_exp = basin_feat.unsqueeze(1).expand(n_nodes, n_horizons, -1) | |
| horizon_feat = self.horizon_embed( | |
| horizons.float().view(1, n_horizons, 1).expand(n_nodes, n_horizons, 1) | |
| ) | |
| if forecast_precip is None: | |
| # No real forecast dataset wired in yet -- fully missing, | |
| # matching today's behavior exactly (see docstring). | |
| forecast_precip = torch.zeros(n_nodes, n_horizons, device=node_hidden.device) | |
| forecast_missing = torch.ones(n_nodes, n_horizons, device=node_hidden.device) | |
| forecast_input = torch.stack([forecast_precip, forecast_missing], dim=-1) # [n_nodes, n_horizons, 2] | |
| forecast_feat = self.forecast_embed(forecast_input) | |
| decoder_input = torch.cat([node_hidden_exp, basin_feat_exp, horizon_feat, forecast_feat], dim=-1) | |
| raw = self.decoder(decoder_input) # [n_nodes, n_horizons, n_targets * n_quantiles] | |
| n_quantiles = len(self.quantiles) | |
| raw = raw.view(n_nodes, n_horizons, self.n_targets, n_quantiles) | |
| # Raw decoder output bounded before the monotonicity | |
| # transform -- an untrained model's raw output is | |
| # essentially unconstrained noise (no output activation | |
| # bounds it), and that noise now feeds pinball loss | |
| # directly. Confirmed as a real, not hypothetical, gap: | |
| # Manning's equation got exactly this kind of clamp (on | |
| # depth, n, hydraulic_radius) after a real training run | |
| # produced non-finite validation loss, but that same | |
| # reasoning was never applied to the decoder's own output, | |
| # which is upstream of every loss in this model, not just | |
| # Manning's. +-50 in standardized units is a generous | |
| # ceiling -- real discharge only ever reached ~18x its | |
| # median even at extreme floods, which standardizes to well | |
| # under this bound -- so this loses no real information | |
| # while giving a hard numerical floor under everything | |
| # downstream. | |
| raw = torch.clamp(raw, min=-50.0, max=50.0) | |
| # Monotonicity across quantiles is architecturally guaranteed, | |
| # not hoped for -- a real, well-known failure mode of naive | |
| # independently-trained quantile regression is "quantile | |
| # crossing" (e.g. the predicted 90th percentile coming out | |
| # lower than the predicted median), which is physically | |
| # nonsensical for a real cumulative distribution. The lowest | |
| # quantile (self.quantiles[0], e.g. the median) is predicted | |
| # directly and unconstrained; every subsequent quantile is | |
| # the previous one plus a softplus-transformed offset, which | |
| # is always >= 0 by construction -- so quantiles[k] >= | |
| # quantiles[k-1] for every k, always, regardless of what | |
| # gradient descent does to the underlying raw outputs. | |
| first = raw[..., 0:1] | |
| increments = torch.nn.functional.softplus(raw[..., 1:]) | |
| quantile_preds = torch.cat([first, first + torch.cumsum(increments, dim=-1)], dim=-1) | |
| return quantile_preds | |
| def masked_mse_graph(pred: "torch.Tensor", target: "torch.Tensor", mask: "torch.Tensor") -> "torch.Tensor": | |
| """ | |
| Supervised loss, masked to real observations only -- reusing | |
| the exact NaN-masking principle already validated for | |
| physics_losses.py and the climate anomaly model's loss, applied | |
| here to the [n_nodes, n_horizons, n_targets] prediction shape. | |
| Real gauges only ever supply a tiny fraction of [n_nodes, | |
| n_horizons] cells; everywhere else this loss contributes | |
| nothing, and only the physics losses (evaluated separately, | |
| graph-wide, on the same forward pass's output) provide training | |
| signal. | |
| """ | |
| if not mask.any(): | |
| return pred.sum() * float("nan") | |
| return ((pred[mask] - target[mask]) ** 2).mean() | |
| else: | |
| RiverMessagePassing = None | |
| SpatiotemporalGNN = None | |
| masked_mse_graph = None |