""" PyTorch Geometric graph construction for the hydrometric station network. SCOPE — surface connectivity only. Groundwater (ADES well) data is deliberately NOT used here. We don't yet have strong enough evidence (correlated well hydrographs, shared BDLISA aquifer units) to justify drawing subsurface/karst edges between stations or basins, and a wrong edge would quietly corrupt a physics-informed model rather than help it. Revisit if that evidence gets built later (see the wells discussion around river_line.py / river_centerline.py). KARST CAVEAT — disappearing rivers ("pertes"): La Eure and La Risle run over Normandy chalk plateau terrain, a geological setting where a river can lose flow underground for a stretch and resurface downstream. An edge in this graph means "these two gauges are sequential along the river's real course" — the course itself (from the digitized centerline) is reliable, but that does NOT by itself verify surface flow is continuous end to end: a correctly-mapped river can still have a losing reach along it. If/when specific losing reaches are known, pass them via `known_losing_reaches` to flag (default) or exclude the affected edges, rather than silently assuming continuity everywhere. USAGE — combining with node_features.py's enriched table: `build_surface_edges` rebuilds its own minimal nodes_df (station_code, basin_id, latitude, longitude, elevation_m) from scratch, since edge construction only needs those columns. To get the enriched feature set (IDPR, groundwater, climate, ...) into the graph, merge it back on top, keyed by station_code, so node ordering stays consistent with edge_index. Note La Eure and La Risle are built as TWO SEPARATE Data objects, not one combined graph — see `build_pyg_graphs_per_basin`: from src.graph.node_features import build_node_features from src.graph.build_graph import build_surface_edges, build_pyg_graphs_per_basin enriched, feat_report = build_node_features( station_elevations_path=..., idpr_path=..., ades_path=..., hydrometric_path=..., ) base_nodes, edges_df, graph_report = build_surface_edges( enriched, centerline_dir=..., basin_file_names={0: "eure", 1: "risle"}, ) full_nodes = base_nodes.merge( enriched.drop(columns=["basin_id", "latitude", "longitude", "elevation_m"]), on="station_code", how="left", ) graphs = build_pyg_graphs_per_basin(full_nodes, edges_df) eure_graph, risle_graph = graphs[0], graphs[1] """ from dataclasses import dataclass from pathlib import Path from typing import Dict, List, Optional, Tuple import numpy as np import pandas as pd try: import torch from torch_geometric.data import Data except ImportError as e: raise ImportError( "This module requires torch and torch_geometric: " "pip install torch torch_geometric" ) from e from ..data.river_graph import ( NodeBuilder, StationNode, order_stations_by_elevation, EdgeBuilder, assign_basin_id, ) from ..data.river_centerline import ( load_centerline, snap_gauges_to_centerline, cumulative_distance_km, ) @dataclass class GraphBuildReport: """Summary of what happened during construction, for a sanity check before trusting the graph — print/log this, don't just discard it.""" basin_ids: List[int] n_nodes: int n_edges: int edges_from_centerline: int # built from real digitized geometry edges_from_fallback: int # built from elevation/latitude ordering only (weaker evidence) flagged_losing_reaches: int # edges marked verified_continuous=False excluded_losing_reaches: int # edges dropped entirely due to known_losing_reaches def __str__(self) -> str: return ( f"Graph: {self.n_nodes} nodes, {self.n_edges} edges across basins {self.basin_ids}\n" f" from real centerline geometry: {self.edges_from_centerline}\n" f" from elevation/latitude fallback (no centerline available): {self.edges_from_fallback}\n" f" flagged as possibly-discontinuous (known losing reach): {self.flagged_losing_reaches}\n" f" excluded entirely (known losing reach): {self.excluded_losing_reaches}" ) def _edges_from_real_centerline( basin_id: int, b_nodes_df: pd.DataFrame, centerline: pd.DataFrame, ) -> pd.DataFrame: """ Build sequential surface edges using real digitized centerline order (mouth -> source measured as centerline_km, so upstream = larger km). This is preferred over elevation/latitude ordering when available, since it follows the river's actual traced course rather than assuming station elevation alone determines sequence. """ gauges = snap_gauges_to_centerline(centerline, b_nodes_df) gauges = gauges.sort_values("centerline_km", ascending=False).reset_index(drop=True) edges = [] for i in range(len(gauges) - 1): a, b = gauges.iloc[i], gauges.iloc[i + 1] edges.append({ "source": a["station_code"], "target": b["station_code"], "basin_id": basin_id, "distance_km": round(float(a["centerline_km"] - b["centerline_km"]), 3), "elevation_drop_m": ( round(float(a["elevation_m"] - b["elevation_m"]), 2) if pd.notna(a["elevation_m"]) and pd.notna(b["elevation_m"]) else None ), "source_of_evidence": "real_centerline", }) return pd.DataFrame(edges) def _edges_from_fallback_ordering(basin_id: int, b_nodes: List[StationNode]) -> pd.DataFrame: """Fallback when no digitized centerline exists for this basin: order by elevation (or latitude, per-basin, if elevation is incomplete — see order_stations_by_elevation), then chain sequentially. Weaker evidence than a real traced course.""" order = order_stations_by_elevation(b_nodes, descending=True) edges_df = EdgeBuilder.build_sequential(b_nodes, order) if not edges_df.empty: edges_df["source_of_evidence"] = "elevation_or_latitude_fallback" return edges_df def build_surface_edges( elevations_df: pd.DataFrame, centerline_dir: Optional[Path] = None, basin_file_names: Optional[Dict[int, str]] = None, known_losing_reaches: Optional[List[Tuple[str, str]]] = None, exclude_losing_reaches: bool = False, prefix_map: Optional[Dict[str, int]] = None, ) -> Tuple[pd.DataFrame, pd.DataFrame, GraphBuildReport]: """ Build station nodes and surface-only edges, preferring real digitized centerline order per basin when available, falling back to elevation/ latitude ordering otherwise. Args: elevations_df: [station_code, latitude, longitude, elevation_m]. centerline_dir: directory containing {name}_centerline.csv files (see river_centerline.py). If None, every basin uses the fallback ordering. basin_file_names: {basin_id: file_key} mapping to find each basin's centerline CSV, e.g. {0: "eure", 1: "risle"}. known_losing_reaches: list of (source_station_code, target_station_code) pairs known or suspected to have discontinuous surface flow (karst losing reaches). These edges are flagged (verified_continuous=False) by default, or dropped entirely if exclude_losing_reaches=True. exclude_losing_reaches: if True, remove known losing-reach edges from the graph instead of just flagging them. prefix_map: optional override for basin assignment from station code prefix, see river_graph.assign_basin_id. Returns: (nodes_df, edges_df, report) """ basin_file_names = basin_file_names or {0: "eure", 1: "risle"} known_losing_reaches = set(known_losing_reaches or []) nodes = NodeBuilder.build_from_dataframe(elevations_df, prefix_map) if not nodes: raise ValueError("No stations matched a basin prefix — check the data or prefix_map.") basins: Dict[int, List[StationNode]] = {} for n in nodes: basins.setdefault(n.basin_id, []).append(n) all_edges = [] n_from_centerline, n_from_fallback = 0, 0 for basin_id, basin_nodes in basins.items(): b_nodes_df = pd.DataFrame([{ "station_code": n.station_id, "latitude": n.latitude, "longitude": n.longitude, "elevation_m": n.elevation, } for n in basin_nodes]) csv_path = (centerline_dir / f"{basin_file_names.get(basin_id, basin_id)}_centerline.csv" if centerline_dir else None) if csv_path and csv_path.exists(): centerline = load_centerline(csv_path) edges_df = _edges_from_real_centerline(basin_id, b_nodes_df, centerline) n_from_centerline += len(edges_df) else: edges_df = _edges_from_fallback_ordering(basin_id, basin_nodes) n_from_fallback += len(edges_df) all_edges.append(edges_df) edges_df = pd.concat(all_edges, ignore_index=True) if all_edges else pd.DataFrame() n_flagged, n_excluded = 0, 0 if not edges_df.empty: edges_df["verified_continuous"] = ~edges_df.apply( lambda r: (r["source"], r["target"]) in known_losing_reaches, axis=1 ) n_flagged = int((~edges_df["verified_continuous"]).sum()) if exclude_losing_reaches and n_flagged: n_excluded = n_flagged n_flagged = 0 edges_df = edges_df[edges_df["verified_continuous"]].reset_index(drop=True) nodes_df = pd.DataFrame([{ "station_code": n.station_id, "basin_id": n.basin_id, "latitude": n.latitude, "longitude": n.longitude, "elevation_m": n.elevation, } for n in nodes]) report = GraphBuildReport( basin_ids=sorted(basins.keys()), n_nodes=len(nodes_df), n_edges=len(edges_df), edges_from_centerline=n_from_centerline, edges_from_fallback=n_from_fallback, flagged_losing_reaches=n_flagged, excluded_losing_reaches=n_excluded, ) return nodes_df, edges_df, report def build_pyg_graph( nodes_df: pd.DataFrame, edges_df: pd.DataFrame, feature_columns: Optional[List[str]] = None, add_missingness_flags: bool = True, bidirectional: bool = False, standardize_features: bool = True, ) -> Data: """ Convert the surface-connectivity tables into a torch_geometric Data object. Basins with no cross-basin edges naturally form disconnected components in the graph (as intended — surface connectivity only). Node features (x): auto-detected numeric columns from `nodes_df` if `feature_columns` is None — this means it works directly with the plain 4-column table from `build_surface_edges` OR the enriched table from `node_features.build_node_features` (elevation, IDPR, groundwater, climate, ...). Columns named `target_*` are always excluded from `x` regardless of `feature_columns`, since those are the model's targets, not inputs — see node_features.py's docstring on why interpolating/using them as features would be label leakage. If present, target columns are attached separately as `data.y` (raw, NaN preserved — the training loop should mask NaN targets, not have them silently imputed). STRUCTURAL_COLUMNS (is_gauged, is_confluence, is_split_point, is_rejoin_point, snap_distance_km, braid_id) are ALSO always excluded from `x`, even though several are boolean — and pandas treats bool as a numeric dtype, so without this explicit exclusion they'd silently get z-scored and fed to the model as if they were physical covariates like elevation, which they are not (confirmed: this actually happened before this exclusion list existed). They're still attached to the returned Data object as their own named attributes (data.is_gauged, etc.) rather than being fully discarded — needed downstream for masking supervised loss to gauged nodes and for physics_losses.py's confluence/braid index builders, just not as model input. Edge features (edge_attr): [distance_km, elevation_drop_m, verified_continuous] Args: nodes_df, edges_df: from `build_surface_edges` (optionally with nodes_df enriched via node_features.build_node_features, re-merged onto the same station ordering — see module usage example). feature_columns: explicit list of columns to use as `x`. If None, auto-detects all numeric, non-target, non-structural columns. add_missingness_flags: if True, adds a `{col}__was_missing` binary column for any feature column that had NaNs, before mean-filling those NaNs — so the model can distinguish "no data nearby" from "value happens to be near the mean". bidirectional: if True, add a reverse edge for every upstream-> downstream edge (common in GNN practice to let information flow both ways even though physical flow is directional). The reverse edges get elevation_drop_m negated. standardize_features: if True, z-score each feature column (after mean-filling). Set False to keep raw units. Returns: A torch_geometric.data.Data with `.station_codes` (index -> station_code), `.node_id_map` (station_code -> index), `.feature_names` (x column order), `.basin_id`, and — if any `target_*` columns were present — `.y` (raw values, NaN preserved) and `.target_names`. Any STRUCTURAL_COLUMNS present on nodes_df are attached as their own same-named attributes (booleans as a bool tensor/array; snap_distance_km as float, NaN for non-gauged nodes; braid_id kept as a plain Python list, since it's station-code strings or None, not something to tensor-ify). """ station_codes = nodes_df["station_code"].tolist() node_id_map = {code: i for i, code in enumerate(station_codes)} STRUCTURAL_COLUMNS = { "is_gauged", "is_confluence", "is_split_point", "is_rejoin_point", "snap_distance_km", "braid_id", } # Feature columns that are genuinely time-varying in reality, even # though today's pipeline only ever produces a period-aggregate # value for them (see node_features.py's add_safran_features/ # add_groundwater_features/add_ndvi_features) -- distinct from # STATIC_COLUMNS below, which are physically time-invariant (a # location's elevation or IDPR index doesn't change on any # timescale relevant here). n_nearby_wells is classified static # despite living under the groundwater loader: it describes monitor # coverage (which wells exist nearby), not a water-table quantity # that itself varies day to day. DYNAMIC_PREFIXES = ("climate_", "avg_groundwater_level", "avg_groundwater_depth", "ndvi_") def _is_dynamic(col: str) -> bool: base = col[:-len("__was_missing")] if col.endswith("__was_missing") else col return any(base.startswith(p) for p in DYNAMIC_PREFIXES) target_cols = [c for c in nodes_df.columns if c.startswith("target_")] if feature_columns is None: feature_columns = [ c for c in nodes_df.columns if c not in ("station_code",) and c not in target_cols and c not in STRUCTURAL_COLUMNS and pd.api.types.is_numeric_dtype(nodes_df[c]) ] feat = nodes_df[feature_columns].copy() missingness_cols = [] if add_missingness_flags: for col in feature_columns: if feat[col].isna().any(): flag_col = f"{col}__was_missing" feat[flag_col] = feat[col].isna().astype(float) missingness_cols.append(flag_col) for col in feature_columns: if feat[col].isna().any(): fill_value = feat[col].mean() feat[col] = feat[col].fillna(0.0 if pd.isna(fill_value) else fill_value) if standardize_features: for col in feature_columns: std = feat[col].std() feat[col] = (feat[col] - feat[col].mean()) / std if std and std > 0 else 0.0 all_feature_cols = feature_columns + missingness_cols x = torch.tensor(feat[all_feature_cols].values, dtype=torch.float) if edges_df.empty: edge_index = torch.zeros((2, 0), dtype=torch.long) edge_attr = torch.zeros((0, 3), dtype=torch.float) else: src_idx = edges_df["source"].map(node_id_map).values tgt_idx = edges_df["target"].map(node_id_map).values if pd.isna(src_idx).any() or pd.isna(tgt_idx).any(): raise ValueError("Some edge endpoints are not present in nodes_df — check station codes.") elev_drop = edges_df["elevation_drop_m"].fillna(0.0).values dist_km = edges_df["distance_km"].values verified = edges_df.get("verified_continuous", pd.Series([True] * len(edges_df))).astype(float).values if bidirectional: index_pairs = np.concatenate([ np.stack([src_idx, tgt_idx]), np.stack([tgt_idx, src_idx]), ], axis=1) attr = np.concatenate([ np.stack([dist_km, elev_drop, verified], axis=1), np.stack([dist_km, -elev_drop, verified], axis=1), ], axis=0) else: index_pairs = np.stack([src_idx, tgt_idx]) attr = np.stack([dist_km, elev_drop, verified], axis=1) edge_index = torch.tensor(index_pairs, dtype=torch.long) edge_attr = torch.tensor(attr, dtype=torch.float) data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr) data.station_codes = station_codes data.node_id_map = node_id_map data.feature_names = all_feature_cols data.basin_id = torch.tensor(nodes_df["basin_id"].values, dtype=torch.long) # Static/dynamic split: `x` stays the full combined tensor (nothing # existing that reads data.x breaks), but every feature is also # tagged and split out separately -- x_static for physically # time-invariant quantities (elevation, IDPR, catchment area, # landcover), x_dynamic for quantities that vary in reality even # though the current pipeline only ever hands them over as a single # period-aggregate (climate, groundwater level/depth, NDVI). See # DYNAMIC_PREFIXES above for exactly which columns land where, and # node_features.py's add_safran_features/add_groundwater_features/ # add_ndvi_features docstrings for why "dynamic" here still means # "one aggregated number," not a real time series yet -- a genuine # multi-timestep pipeline is a separate, larger piece of future work # this split does not attempt to solve on its own. dynamic_mask = [_is_dynamic(c) for c in all_feature_cols] static_idx = [i for i, d in enumerate(dynamic_mask) if not d] dynamic_idx = [i for i, d in enumerate(dynamic_mask) if d] data.static_feature_names = [all_feature_cols[i] for i in static_idx] data.dynamic_feature_names = [all_feature_cols[i] for i in dynamic_idx] x_np_for_split = feat[all_feature_cols].values data.x_static = torch.tensor(x_np_for_split[:, static_idx], dtype=torch.float) if static_idx else None data.x_dynamic = torch.tensor(x_np_for_split[:, dynamic_idx], dtype=torch.float) if dynamic_idx else None for bool_col in ("is_gauged", "is_confluence", "is_split_point", "is_rejoin_point"): if bool_col in nodes_df.columns: setattr(data, bool_col, torch.tensor(nodes_df[bool_col].values, dtype=torch.bool)) if "snap_distance_km" in nodes_df.columns: data.snap_distance_km = torch.tensor( nodes_df["snap_distance_km"].astype(float).values, dtype=torch.float ) if "braid_id" in nodes_df.columns: # station-code strings or None -- not tensor-able, kept as a plain # list so physics_losses.py's build_braid_index can still use it # (that function already works on nodes_df directly, so this is # for convenience when only the Data object is at hand, not a # hard requirement). data.braid_id = nodes_df["braid_id"].tolist() if target_cols: data.y = torch.tensor(nodes_df[target_cols].values, dtype=torch.float) data.target_names = target_cols return data def build_pyg_graphs_per_basin( nodes_df: pd.DataFrame, edges_df: pd.DataFrame, feature_columns: Optional[List[str]] = None, add_missingness_flags: bool = True, bidirectional: bool = False, standardize_features: bool = True, ) -> Dict[int, Data]: """ Build a SEPARATE torch_geometric Data object per basin, rather than one combined graph with disconnected components. La Eure and La Risle are distinct hydrographic systems with no surface connection between them (see the KARST CAVEAT above and river_graph.py's basin assignment) — one merged Data object would blur that distinction, and PyG's own batching (`torch_geometric.data.Batch.from_data_list`) already expects a list of separate small graphs, not one pre-merged blob, so this is also the more natural shape for training. `basin_id` is dropped from the per-graph feature set automatically (it's constant within a single basin's graph and would carry zero information there — standardizing a zero-variance column is meaningless). It's still available as `data.basin_id` metadata on each graph if you need it for bookkeeping. Args: nodes_df, edges_df: from `build_surface_edges` (optionally enriched, see the module docstring's usage example). feature_columns: explicit column list PER BASIN. If None, auto-detected the same way as `build_pyg_graph`, minus `basin_id`. add_missingness_flags, bidirectional, standardize_features: passed straight through to `build_pyg_graph` for each basin. Returns: {basin_id: Data}, each with its own LOCAL node indexing (0..n-1 within that basin) and its own `.station_codes` / `.node_id_map` / `.feature_names`. """ graphs = {} for basin_id in sorted(nodes_df["basin_id"].unique()): b_nodes = nodes_df[nodes_df["basin_id"] == basin_id].reset_index(drop=True) b_edges = edges_df[edges_df["basin_id"] == basin_id].reset_index(drop=True) cols = feature_columns if cols is None: target_cols = [c for c in b_nodes.columns if c.startswith("target_")] cols = [ c for c in b_nodes.columns if c not in ("station_code", "basin_id") and c not in target_cols and pd.api.types.is_numeric_dtype(b_nodes[c]) ] graphs[basin_id] = build_pyg_graph( b_nodes, b_edges, feature_columns=cols, add_missingness_flags=add_missingness_flags, bidirectional=bidirectional, standardize_features=standardize_features, ) return graphs