""" Physics-informed loss terms for streamflow prediction on the reach graph. Four constraints, each tied to real structure the graph now actually has (see src/graph/build_reach_graph.py): 1. confluence_mass_balance_loss -- new mass genuinely enters (is_confluence) 2. split_rejoin_conservation_loss -- no new mass, paired via braid_id 3. routing_consistency_loss -- travel-time lag from distance_km/elevation_drop_m 4. water_balance_loss -- P - ET - Q - deltaS ~= 0 per node All four apply graph-wide, not just at the 27 gauged nodes -- that's the actual mechanism by which sparse labels generalize to ~4,500 ungauged nodes, not an incidental detail. No model exists yet, so these are standalone functions operating on whatever Q tensor a model eventually produces. They use only operations that behave identically on a torch.Tensor or a plain numpy array (indexing, elementwise arithmetic, sum, mean) so the same code path is testable now with numpy and will work unchanged with real torch tensors and real gradients once a model exists -- verified by testing this module against numpy inputs directly. """ from typing import Dict, List, Optional, Tuple import numpy as np import pandas as pd try: import torch _HAS_TORCH = True except ImportError: _HAS_TORCH = False def _mse(residual): """ Mean squared residual, NaN-masked -- works identically on numpy or torch. Real ground-truth Q is naturally, heavily NaN (only real gauges with real observations ever have a value; confirmed against real data: ~93.5% NaN for the reach graph's discharge tensor) -- without masking, ANY single NaN anywhere in the residual poisons the entire mean to NaN, which isn't a rare edge case for this data, it's the normal shape of it. Returns NaN only if truly nothing usable exists (every entry NaN), which is a real "no data" signal worth surfacing, not silently averaging to 0 and implying perfect physics satisfaction when there was actually no evidence either way. """ if _HAS_TORCH and isinstance(residual, torch.Tensor): mask = ~torch.isnan(residual) if not mask.any(): return residual.sum() * float("nan") return (residual[mask] ** 2).mean() mask = ~np.isnan(residual) if not mask.any(): return np.nan return (residual[mask] ** 2).mean() # --------------------------------------------------------------------------- # 1. Confluence mass balance: Q_confluence ~= sum(Q_upstream_branches) # --------------------------------------------------------------------------- def build_confluence_index(nodes_df: pd.DataFrame, edges_df: pd.DataFrame) -> List[Tuple[int, List[int]]]: """ Precompute, once per graph (not per training step), which node indices feed into each real confluence. Returns [(confluence_idx, [upstream_idx, ...]), ...] using positional indices into nodes_df (0..n-1), matching how a model's output tensor would be laid out. Splitting this out from the loss function itself means the loss can just do array indexing every step -- the graph structure doesn't change between training steps, so there's no reason to recompute which nodes are involved in each confluence on every call. """ code_to_idx = {code: i for i, code in enumerate(nodes_df["station_code"])} confluence_codes = set(nodes_df[nodes_df["is_confluence"]]["station_code"]) pairs = [] for conf_code in confluence_codes: upstream = edges_df[edges_df["target"] == conf_code]["source"].tolist() upstream_idx = [code_to_idx[u] for u in upstream if u in code_to_idx] if len(upstream_idx) >= 2: pairs.append((code_to_idx[conf_code], upstream_idx)) return pairs def confluence_mass_balance_loss(Q, confluence_index: List[Tuple[int, List[int]]]): """ For each real confluence, predicted discharge there should approximately equal the sum of its upstream branches' predicted discharge -- new mass genuinely enters at a confluence (an independent tributary catchment), so this is a straightforward sum, unlike the split/rejoin case below. Ignores travel time between the branches and the confluence (an instantaneous-mass approximation) -- see routing_consistency_loss for the piece that accounts for lag separately. Args: Q: predicted discharge, shape [n_nodes] (one timestep) or [n_nodes, T] (multiple timesteps, this loss applies per timestep the same way). confluence_index: from build_confluence_index. Returns: Scalar loss (0.0, on the same array type as Q, if no confluences). """ if not confluence_index: return Q.sum() * 0.0 # zero, but keeps dtype/type consistent (torch-safe) residuals = [] for conf_idx, upstream_idx in confluence_index: upstream_sum = Q[upstream_idx[0]] for idx in upstream_idx[1:]: upstream_sum = upstream_sum + Q[idx] residuals.append(Q[conf_idx] - upstream_sum) if _HAS_TORCH and isinstance(Q, torch.Tensor): residual_stack = torch.stack(residuals) else: residual_stack = np.stack(residuals) return _mse(residual_stack) # --------------------------------------------------------------------------- # 2. Split/rejoin conservation: Q_split ~= Q_rejoin (no new mass between them) # --------------------------------------------------------------------------- def build_braid_index(nodes_df: pd.DataFrame) -> List[Tuple[int, int]]: """ Precompute (split_idx, rejoin_idx) pairs from the saved braid_id column (see build_reach_graph.py's pair_splits_and_rejoins). Same precompute-once rationale as build_confluence_index. """ code_to_idx = {code: i for i, code in enumerate(nodes_df["station_code"])} pairs = [] for _, row in nodes_df[nodes_df["braid_id"].notna()].iterrows(): rejoin_code, split_code = row["station_code"], row["braid_id"] if rejoin_code in code_to_idx and split_code in code_to_idx: pairs.append((code_to_idx[split_code], code_to_idx[rejoin_code])) return pairs def split_rejoin_conservation_loss(Q, braid_index: List[Tuple[int, int]]): """ For each matched split/rejoin pair, predicted discharge should be approximately equal at both ends -- the same water dividing into parallel channels and recombining adds no new mass, unlike a real confluence (see confluence_mass_balance_loss). This is a genuinely different physical constraint, not a weaker version of the same one: a model that learned "sum inflows" generically would get this wrong, since a rejoin's two branches together should equal the SPLIT's single value, not add something new on top of it. Args: Q: predicted discharge, shape [n_nodes] or [n_nodes, T]. braid_index: from build_braid_index. Returns: Scalar loss (0.0 if no braids in this graph). """ if not braid_index: return Q.sum() * 0.0 split_idx = [s for s, _ in braid_index] rejoin_idx = [r for _, r in braid_index] residual = Q[split_idx] - Q[rejoin_idx] return _mse(residual) # --------------------------------------------------------------------------- # 3. Routing: travel-time lag from real channel distance and slope # --------------------------------------------------------------------------- def estimate_travel_time_hours( distance_km, elevation_drop_m, min_velocity_ms: float = 0.1, max_velocity_ms: float = 3.0, velocity_coefficient: float = 1.0, ) -> float: """ Rough channel-flow velocity from slope, in the spirit of Manning's equation's slope dependence (v ~ sqrt(slope)) without the channel geometry/roughness terms Manning's actually needs, which we don't have real data for -- explicitly an approximation, not a full hydraulic solve. Slope = elevation_drop_m / (distance_km * 1000). Clamped to [min_velocity_ms, max_velocity_ms] since a near-zero or negative slope (a virtually flat reach, or a data artifact) would otherwise give a nonsensical near-infinite or negative travel time. Returns: Travel time in hours for water to traverse this edge. """ distance_m = distance_km * 1000.0 slope = np.clip(elevation_drop_m / np.maximum(distance_m, 1.0), 1e-6, None) velocity = np.clip(velocity_coefficient * np.sqrt(slope) * 10.0, min_velocity_ms, max_velocity_ms) return distance_m / velocity / 3600.0 def build_routing_index( nodes_df: pd.DataFrame, edges_df: pd.DataFrame, timestep_hours: float = 24.0, ) -> List[Tuple[int, int, int]]: """ Precompute (upstream_idx, downstream_idx, lag_timesteps) for every edge, rounding each edge's estimated travel time to the nearest whole timestep -- e.g. a 30-hour travel time at a 24-hour (daily) timestep rounds to a 1-step lag. An edge whose travel time rounds to 0 is still included (same-timestep routing, lag=0). """ code_to_idx = {code: i for i, code in enumerate(nodes_df["station_code"])} pairs = [] for _, e in edges_df.iterrows(): if e["source"] not in code_to_idx or e["target"] not in code_to_idx: continue drop = e["elevation_drop_m"] if pd.notna(e["elevation_drop_m"]) else 0.1 hours = estimate_travel_time_hours(e["distance_km"], max(drop, 0.1)) lag = int(round(hours / timestep_hours)) pairs.append((code_to_idx[e["source"]], code_to_idx[e["target"]], lag)) return pairs def routing_consistency_loss(Q, routing_index: List[Tuple[int, int, int]]): """ Q at a downstream node at time t should approximately equal Q at its upstream node at time (t - lag), lag coming from real distance and slope (build_routing_index) -- not just "conserve mass at the same instant," which routing_consistency_loss's siblings above assume as a simplification. This is the piece that makes that simplification less necessary over time: a well-trained model satisfying this loss is learning the actual travel-time behavior of each reach. Args: Q: predicted discharge, shape [n_nodes, T] -- REQUIRES a time dimension, unlike the other three losses, since travel-time lag is meaningless for a single instant. routing_index: from build_routing_index. Returns: Scalar loss (0.0 if no edges have a usable lag within Q's time range). """ T = Q.shape[1] residuals = [] for up_idx, down_idx, lag in routing_index: if lag >= T: continue # this edge's travel time exceeds the whole prediction window if lag == 0: residuals.append(Q[down_idx, :] - Q[up_idx, :]) else: residuals.append(Q[down_idx, lag:] - Q[up_idx, :-lag]) if not residuals: return Q.sum() * 0.0 if _HAS_TORCH and isinstance(Q, torch.Tensor): residual_cat = torch.cat(residuals) else: residual_cat = np.concatenate(residuals) return _mse(residual_cat) # --------------------------------------------------------------------------- # 4. Water balance: P - ET - Q - deltaS ~= 0, per node # --------------------------------------------------------------------------- def water_balance_loss( Q_m3s, precip_mm, evap_mm, catchment_area_km2, period_days: float = 365.0, delta_storage_m3: Optional[object] = None, ): """ Precipitation minus evapotranspiration minus discharge minus storage change should balance to ~0, in volume terms, over the given period. UNIT CONVERSION (the easy part to get subtly wrong): 1 mm of depth over 1 km^2 is 1000 m^3 (1 km^2 = 1e6 m^2, 1 mm = 1e-3 m, 1e6 * 1e-3 = 1e3). P and ET (mm, over the period) get converted to m^3 via catchment_area_km2 before comparing against Q, which is converted from a rate (m^3/s) to a volume by multiplying by the period length in seconds. delta_storage_m3 defaults to zero (a steady-state approximation) -- we have no direct storage measurement (soil moisture, groundwater volume change) in this project's data, only groundwater LEVEL at sparse wells, which isn't the same thing as a basin-wide storage volume. Treating deltaS as strictly zero is a real, named approximation, not a hidden one -- pass a nonzero delta_storage_m3 if a proxy for it becomes available later (e.g. derived from groundwater level trend where well coverage allows it). Args: Q_m3s: predicted discharge, shape [n_nodes] (period-average rate). precip_mm, evap_mm: node features, already available. catchment_area_km2: from catchment.py -- NaN for ungauged/unknown catchments, in which case that node is excluded from this loss entirely (silently including it with a wrong/zero area would corrupt the term, not just add noise). period_days: length of the period P/ET/Q are aggregated over. delta_storage_m3: optional storage change; zero-array default. Returns: Scalar loss, computed only over nodes with a real catchment area. """ valid = ~np.isnan(catchment_area_km2) if not _HAS_TORCH or not isinstance(catchment_area_km2, torch.Tensor) \ else ~torch.isnan(catchment_area_km2) period_seconds = period_days * 86400.0 Q_volume_m3 = Q_m3s * period_seconds P_volume_m3 = precip_mm * catchment_area_km2 * 1000.0 ET_volume_m3 = evap_mm * catchment_area_km2 * 1000.0 dS = delta_storage_m3 if delta_storage_m3 is not None else (Q_m3s * 0.0) residual = P_volume_m3 - ET_volume_m3 - Q_volume_m3 - dS residual_valid = residual[valid] if (residual_valid.shape[0] if hasattr(residual_valid, "shape") else len(residual_valid)) == 0: return Q_m3s.sum() * 0.0 return _mse(residual_valid) # --------------------------------------------------------------------------- # Combined loss # --------------------------------------------------------------------------- def physics_informed_loss( Q_supervised_pred, Q_supervised_true, gauged_mask, Q_full, confluence_index, braid_index, weights: Optional[Dict[str, float]] = None, Q_timeseries=None, routing_index=None, precip_mm=None, evap_mm=None, catchment_area_km2=None, ) -> Dict[str, float]: """ Combines the supervised loss (masked to gauged nodes) with all physics terms that have the inputs to compute (routing and water balance are optional -- they need a time dimension / climate data respectively, which not every training step may have on hand). Returns a dict of every individual term plus 'total', rather than just the summed scalar -- so it's possible to see which physics term is actually driving the loss during training, not just that "the loss" went up or down. """ weights = weights or {"confluence": 1.0, "split_rejoin": 1.0, "routing": 1.0, "water_balance": 1.0} supervised_residual = (Q_supervised_pred - Q_supervised_true)[gauged_mask] losses = {"supervised": _mse(supervised_residual)} losses["confluence"] = confluence_mass_balance_loss(Q_full, confluence_index) losses["split_rejoin"] = split_rejoin_conservation_loss(Q_full, braid_index) if Q_timeseries is not None and routing_index is not None: losses["routing"] = routing_consistency_loss(Q_timeseries, routing_index) if precip_mm is not None and evap_mm is not None and catchment_area_km2 is not None: losses["water_balance"] = water_balance_loss(Q_full, precip_mm, evap_mm, catchment_area_km2) total = losses["supervised"] for name, w_key in [("confluence", "confluence"), ("split_rejoin", "split_rejoin"), ("routing", "routing"), ("water_balance", "water_balance")]: if name in losses: total = total + weights.get(w_key, 1.0) * losses[name] losses["total"] = total return losses