""" 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