diff --git a/.gitattributes b/.gitattributes index 4aa7b789e11839467e7f9434cb824b04fb564f76..833788583edb84d91b020995e84578a6cbb4e15c 100644 --- a/.gitattributes +++ b/.gitattributes @@ -134,3 +134,56 @@ plots/safran/precipitation_cumulative.png filter=lfs diff=lfs merge=lfs -text plots/safran/temperature.png filter=lfs diff=lfs merge=lfs -text plots/station_elevations/elevation_bar.png filter=lfs diff=lfs merge=lfs -text plots/station_elevations/station_map.png filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1960.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1961.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1962.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1963.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1964.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1965.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1966.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1967.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1968.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1969.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1970.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1971.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1972.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1973.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1974.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1975.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1976.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1977.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1978.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1979.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1980.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1981.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1982.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1983.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1984.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1985.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1986.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1987.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1988.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1989.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1990.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1991.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1992.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1993.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1994.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1995.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1996.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1997.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1998.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_1999.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_2000.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_2001.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_2002.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_2003.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_2004.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_2005.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_2006.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_2007.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_2008.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_2009.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_2010.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_2011.nc filter=lfs diff=lfs merge=lfs -text +datasets/draft/era5_2012.nc filter=lfs diff=lfs merge=lfs -text diff --git a/datasets/draft/download_log.txt b/datasets/draft/download_log.txt new file mode 100644 index 0000000000000000000000000000000000000000..a14655dd2837760f312e6cc01cbbdc973e5825f3 --- /dev/null +++ b/datasets/draft/download_log.txt @@ -0,0 +1,1233 @@ +2026-08-13 23:26:10,579 INFO Request ID is e4b1d81d-4e1e-4572-981c-4039c50ceb09 +2026-08-13 23:26:10,941 INFO status has been updated to accepted +2026-08-13 23:26:45,758 INFO status has been updated to running +2026-08-13 23:29:07,088 INFO status has been updated to successful +================================================================================ +ERA5 FULL DOWNLOAD (1960-2026) +================================================================================ +Study area: [48.29, 0.33] to [49.44, 1.64] +Years: 1960-2026 (67 years) +Output: datasets/safran/ +================================================================================ + +Downloading 1960... + → Instantaneous variables... + be3cb342730c8c70ed770d341589dce3.nc: 0%| | 0.00/117k [00:00 pd.Series: + """ + Spatial join: for each edge, which incremental catchment polygon + contains its midpoint. Returns a Series indexed like edges_df, with + the polygon's unique identifier (or NaN if no polygon contains it -- + can happen near basin edges or bbox boundaries). + + NEEDS geopandas -- not verified in the environment this was written + in, see module docstring. + """ + import geopandas as gpd + from shapely.geometry import Point + + polygons = gpd.read_file(polygons_path) + if "cleabs" in polygons.columns: + poly_id_col = "cleabs" + else: + polygons = polygons.reset_index().rename(columns={"index": "_poly_idx"}) + poly_id_col = "_poly_idx" + + coord = nodes_df.set_index("station_code")[["latitude", "longitude"]] + midpoints = [] + for _, e in edges_df.iterrows(): + try: + src, tgt = coord.loc[e["source"]], coord.loc[e["target"]] + except KeyError: + midpoints.append(None) + continue + midpoints.append(Point((src["longitude"] + tgt["longitude"]) / 2, + (src["latitude"] + tgt["latitude"]) / 2)) + + edge_points = gpd.GeoDataFrame({"edge_idx": range(len(edges_df))}, geometry=midpoints, crs=polygons.crs) + joined = gpd.sjoin(edge_points.dropna(subset=["geometry"]), polygons[[poly_id_col, "geometry"]], + how="left", predicate="within") + joined = joined.drop_duplicates("edge_idx") # a midpoint exactly on a shared boundary could match >1 polygon + + result = pd.Series(index=range(len(edges_df)), dtype=object) + result.loc[joined["edge_idx"]] = joined[poly_id_col].values + return result + + +def compute_cumulative_catchment( + nodes_df: pd.DataFrame, + edges_df: pd.DataFrame, + edge_polygon_id: pd.Series, + polygon_areas: dict, +) -> pd.Series: + """ + For every node, sum the areas of distinct polygons among all edges + upstream of it (inclusive of the edge(s) arriving at the node + itself). Pure graph-traversal + arithmetic -- no geospatial + dependency, fully testable without geopandas (see the test in this + script's accompanying verification, or run this module's functions + directly against synthetic data). + + Args: + nodes_df: reach graph nodes [station_code, ...] + edges_df: reach graph edges [source, target, ...], same row + order/index as edge_polygon_id. + edge_polygon_id: from assign_edges_to_polygons (or an equivalent + manually-built mapping for testing) -- polygon ID per edge, + aligned to edges_df's index. + polygon_areas: {polygon_id: area_km2} + + Returns: + Series indexed like nodes_df, cumulative catchment area in km2 + (NaN for a node with no upstream edges assigned to any polygon + at all -- a true headwater with no BD TOPO catchment coverage, + distinct from zero). + """ + G = nx.DiGraph() + G.add_nodes_from(nodes_df["station_code"]) + edge_to_poly = {} + for i, e in edges_df.iterrows(): + G.add_edge(e["source"], e["target"]) + poly = edge_polygon_id.get(i) + if pd.notna(poly): + edge_to_poly[(e["source"], e["target"])] = poly + + results = {} + for node in nodes_df["station_code"]: + ancestors = nx.ancestors(G, node) | {node} + # every edge whose TARGET is in this node's ancestor-or-self set + # is an edge that feeds into this node's upstream drainage + polys_here = set() + for (u, v), poly in edge_to_poly.items(): + if v in ancestors: + polys_here.add(poly) + if not polys_here: + results[node] = float("nan") + else: + results[node] = sum(polygon_areas.get(p, 0.0) for p in polys_here) + + return pd.Series(results).reindex(nodes_df["station_code"]) + + +def run_for_basin(data_root: Path, basin_id: int, file_key: str) -> None: + graph_dir = data_root / "reach_graph" + enriched_path = graph_dir / f"{file_key}_nodes_enriched.csv" + edges_path = graph_dir / f"{file_key}_edges.csv" + polygons_path = data_root / "bdtopo_hydro" / "bassin_versant_topographique.geojson" + + if not enriched_path.exists() or not edges_path.exists(): + print(f"{file_key}: missing enriched nodes/edges, run build_reach_graphs.py + " + f"enrich_reach_graph.py first. Skipping.") + return + if not polygons_path.exists(): + print(f"{file_key}: no {polygons_path} found -- run download_bdtopo_hydro.py first. Skipping.") + return + + nodes_df = pd.read_csv(enriched_path) + edges_df = pd.read_csv(edges_path) + print(f"{file_key}: {len(nodes_df)} nodes, {len(edges_df)} edges") + + try: + import geopandas as gpd + except ImportError: + print(f"{file_key}: geopandas not installed -- cannot do the spatial join. Skipping.") + return + + polygons = gpd.read_file(polygons_path) + poly_id_col = "cleabs" if "cleabs" in polygons.columns else "_poly_idx" + if poly_id_col == "_poly_idx": + polygons = polygons.reset_index().rename(columns={"index": poly_id_col}) + polygons_metric = polygons.to_crs(epsg=2154) + polygon_areas = dict(zip(polygons[poly_id_col], polygons_metric.geometry.area / 1e6)) + + edge_polygon_id = assign_edges_to_polygons(edges_df, nodes_df, polygons_path) + n_assigned = edge_polygon_id.notna().sum() + print(f" {n_assigned}/{len(edges_df)} edges assigned to a polygon " + f"({len(edges_df) - n_assigned} unassigned -- near bbox/basin edges, expected)") + + cumulative = compute_cumulative_catchment(nodes_df, edges_df, edge_polygon_id, polygon_areas) + nodes_df["cumulative_catchment_area_km2"] = cumulative.values + + n_covered = nodes_df["cumulative_catchment_area_km2"].notna().sum() + print(f" cumulative_catchment_area_km2 coverage: {n_covered}/{len(nodes_df)} nodes " + f"(vs. {nodes_df['catchment_area_km2'].notna().sum()} from Hub'Eau's own reported values)") + + # Cross-check against Hub'Eau's real values where both exist, as a + # sanity check on the whole pipeline -- they measure different things + # (Hub'Eau: official gauge catchment; this: BD TOPO polygon sum) so + # exact agreement isn't expected, but they should be in the same + # ballpark for a real gauge, not off by an order of magnitude. + both = nodes_df.dropna(subset=["catchment_area_km2", "cumulative_catchment_area_km2"]) + if not both.empty: + both = both.copy() + both["ratio"] = both["cumulative_catchment_area_km2"] / both["catchment_area_km2"] + print(f" Cross-check vs Hub'Eau (real gauges only): ratio mean={both['ratio'].mean():.2f}, " + f"range=[{both['ratio'].min():.2f}, {both['ratio'].max():.2f}]") + print(both[["station_code", "catchment_area_km2", "cumulative_catchment_area_km2", "ratio"]] + .to_string(index=False)) + + nodes_df.to_csv(enriched_path, index=False) + print(f" updated {enriched_path}") + print() + + +def main() -> None: + parser = argparse.ArgumentParser(description="Compute cumulative catchment area from BD TOPO polygons") + parser.add_argument("--data-root", type=Path, default=Path("datasets")) + args = parser.parse_args() + + for basin_id, file_key in [(0, "eure"), (1, "risle")]: + run_for_basin(args.data_root, basin_id, file_key) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/scripts/enrich_reach_graph.py b/scripts/enrich_reach_graph.py new file mode 100644 index 0000000000000000000000000000000000000000..bcde032510d739aac8ec7c75447592b877c4a719 --- /dev/null +++ b/scripts/enrich_reach_graph.py @@ -0,0 +1,109 @@ +""" +Enriches the reach graph's node tables (from scripts/build_reach_graphs.py) +with every available feature loader, via node_features.py's base_nodes_df +path -- closing the gap where the reach graph and the feature pipeline +existed separately but nothing actually connected them. + +Usage: + python -m scripts.enrich_reach_graph --data-root datasets +""" +import argparse +import sys +from pathlib import Path + +import pandas as pd + +try: + from src.graph.node_features import build_node_features +except ImportError: + sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + from src.graph.node_features import build_node_features + +BASIN_FILE_NAMES = {0: "eure", 1: "risle"} + + +def enrich_basin( + data_root: Path, basin_id: int, skip_climate: bool = False, + date_range: "tuple[str, str] | None" = None, +) -> None: + file_key = BASIN_FILE_NAMES[basin_id] + graph_dir = data_root / "reach_graph" + nodes_path = graph_dir / f"{file_key}_nodes.csv" + + if not nodes_path.exists(): + print(f"{file_key}: no {nodes_path} found -- run scripts/build_reach_graphs.py first. Skipping.") + return + + base_nodes = pd.read_csv(nodes_path) + print(f"{file_key}: {len(base_nodes)} nodes loaded from {nodes_path}") + + enriched, report = build_node_features( + base_nodes_df=base_nodes, + idpr_path=data_root / "idpr.csv", + ades_path=data_root / "ades", + safran_path=None if skip_climate else data_root / "safran", + # --skip-climate exists specifically because add_safran_features has + # not been verified at reach-graph scale (~2,900 nodes) and is the + # prime suspect for a real hang there, unlike everything else in + # this pipeline which has been checked at this scale already (see + # add_groundwater_features's rewrite, which hit exactly this class + # of bug once before). Use this flag to get unblocked on everything + # else while that gets diagnosed with the real source. + catchment_path=data_root / "catchment_area.csv", + hydrometric_path=data_root / "hydrometric", + date_range=date_range, + # Safe to include here, unlike an earlier version of this script + # assumed: target_* attaches via a left-merge on station_code, + # and real gauge codes ("H...") can never collide with BD TOPO + # hydrographic node IDs ("NOEUDHYD...") or virtual infill node + # IDs ("VIRTUAL::..."), confirmed against the real station code + # namespace -- so only real gauge rows ever get a real target + # value; every confluence/virtual node correctly gets NaN, not + # a mislabeled one. + ) + print(report) + + n_climate = enriched["climate_precip_mm"].notna().sum() if "climate_precip_mm" in enriched.columns else 0 + n_catchment = enriched["catchment_area_km2"].notna().sum() if "catchment_area_km2" in enriched.columns else 0 + n_targets = enriched["target_discharge_m3s_mean"].notna().sum() if "target_discharge_m3s_mean" in enriched.columns else 0 + print(f" climate coverage: {n_climate}/{len(enriched)} nodes") + print(f" target (discharge) coverage: {n_targets}/{len(enriched)} nodes " + f"(expected: only real gauges -- {int(base_nodes['is_gauged'].sum())} in this basin)") + print(f" catchment_area coverage: {n_catchment}/{len(enriched)} nodes " + f"(expected: only real gauges have one -- see node_features.py's " + f"add_catchment_features docstring; this is what water_balance_loss " + f"can actually use right now, everything else needs a cumulative " + f"catchment estimate this script doesn't build)") + + out_path = graph_dir / f"{file_key}_nodes_enriched.csv" + enriched.to_csv(out_path, index=False) + print(f" saved to {out_path}") + print() + + +def main() -> None: + parser = argparse.ArgumentParser(description="Enrich reach graph nodes with feature loaders") + parser.add_argument("--data-root", type=Path, default=Path("datasets")) + parser.add_argument("--skip-climate", action="store_true", + help="Skip SAFRAN/ERA5 climate features -- use if that step hangs at reach-graph scale") + parser.add_argument("--start-date", type=str, default="2013-01-01", + help="Start of the training period (inclusive). Default 2013-01-01: the window " + "maximizing discharge-gauge coverage (6/8 stations, 13,084 observations) -- " + "see the date-range analysis this default came from before changing it.") + parser.add_argument("--end-date", type=str, default="2026-12-31", + help="End of the training period (inclusive).") + parser.add_argument("--no-date-filter", action="store_true", + help="Disable date filtering entirely -- each source uses its own full history") + args = parser.parse_args() + + date_range = None if args.no_date_filter else (args.start_date, args.end_date) + if date_range: + print(f"Using date range: {date_range[0]} to {date_range[1]}") + print() + + for basin_id in BASIN_FILE_NAMES: + enrich_basin(args.data_root, basin_id, skip_climate=args.skip_climate, date_range=date_range) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/src/__pycache__/test_build_graph.cpython-311.pyc b/src/__pycache__/test_build_graph.cpython-311.pyc index e5819796010e3783328ecaf359e75b5bd769d1f2..1599c77b2461d565e26056a81dbd1256635b14e5 100644 Binary files a/src/__pycache__/test_build_graph.cpython-311.pyc and b/src/__pycache__/test_build_graph.cpython-311.pyc differ diff --git a/src/data/loaders/__pycache__/safran.cpython-311.pyc b/src/data/loaders/__pycache__/safran.cpython-311.pyc index 74fc265ddd47d5260b86241c788d862051554148..1437a89b59bd7361a6d59c0a71a78bbb4a25ca11 100644 Binary files a/src/data/loaders/__pycache__/safran.cpython-311.pyc and b/src/data/loaders/__pycache__/safran.cpython-311.pyc differ diff --git a/src/data/loaders/safran.py b/src/data/loaders/safran.py index 4b690f26101d1fe5c04ded794cd616480d574316..7aba30d98f9afa57c8434ddaa2c9d94598e1c145 100644 --- a/src/data/loaders/safran.py +++ b/src/data/loaders/safran.py @@ -74,34 +74,55 @@ class SAFRANLoader(BaseDataLoader): return combined.sort_values(['date', 'station_code']).reset_index(drop=True) def _load_and_interpolate(self, nc_file: Path) -> pd.DataFrame: - """Load single NetCDF and interpolate to stations.""" + """ + Load single NetCDF and interpolate to ALL stations in one + vectorized xarray selection, not one .sel() + .to_dataframe() + call per station. + + The original per-station loop does real work with meaningful + per-call overhead (an index lookup, then a full DataFrame + conversion) once per station PER FILE. At 27 stations across + ~67 year-files (1960-2026) that's ~1,800 calls -- slow but + tolerable. At the reach graph's ~2,900 nodes it's ~193,000 + individual calls, which is what was actually hanging, not a + proportionally-worse runtime. + + Fix: xarray's vectorized ("pointwise") indexing -- when the + latitude/longitude indexers are DataArrays sharing a common + dimension name ("station" here), a single .sel() call looks up + every point at once instead of one point per call, producing a + result with a "station" dimension instead of separate + latitude/longitude dimensions. This turns ~193,000 calls into + ~67 (one per file). + + CAVEAT -- I could not actually run this against xarray/netCDF4 + in this environment (no network access to install them here), + so this is reasoned from documented xarray API behavior, not + verified execution like the rest of this project's fixes have + been. Please test carefully, ideally against a small subset of + the era5_*.nc files first, and confirm the output matches the + old per-station-loop version's shape/values before trusting it + for real training data. + """ ds = xr.open_dataset(nc_file) - - # Time coordinate name time_coord = 'valid_time' if 'valid_time' in ds.coords else 'time' - interpolated = [] - - # For each station - for _, station in self.station_coords.iterrows(): - station_code = station['station_code'] - lat, lon = station['lat'], station['lon'] - - # Interpolate - point_ds = ds.sel(latitude=lat, longitude=lon, method=self.interp_method) - point_df = point_ds.to_dataframe().reset_index() - - # Rename - point_df = point_df.rename(columns={time_coord: 'date', **self.VAR_MAP}) + station_lats = xr.DataArray(self.station_coords['lat'].values, dims='station') + station_lons = xr.DataArray(self.station_coords['lon'].values, dims='station') + station_codes = self.station_coords['station_code'].values - # Keep relevant columns - keep_cols = ['date'] + [v for v in self.VAR_MAP.values() if v in point_df.columns] - point_df = point_df[keep_cols].copy() - point_df['station_code'] = station_code + point_ds = ds.sel(latitude=station_lats, longitude=station_lons, method=self.interp_method) + # point_ds now has a 'station' dimension (length n_stations) instead + # of separate latitude/longitude dimensions -- attach real station + # codes as a coordinate so they survive the to_dataframe() below, + # rather than staying as anonymous integer positions. + point_ds = point_ds.assign_coords(station=station_codes) - interpolated.append(point_df) + df = point_ds.to_dataframe().reset_index() + df = df.rename(columns={time_coord: 'date', 'station': 'station_code', **self.VAR_MAP}) - df = pd.concat(interpolated, ignore_index=True) + keep_cols = ['date', 'station_code'] + [v for v in self.VAR_MAP.values() if v in df.columns] + df = df[keep_cols].copy() df['date'] = pd.to_datetime(df['date']).dt.tz_localize(None) ds.close() diff --git a/src/graph/__pycache__/build_graph.cpython-311.pyc b/src/graph/__pycache__/build_graph.cpython-311.pyc index b9e003115db30c679c9313c836cf91c56e430040..b28b4e580a602c0caa2a9e5fc5fd1ecd627ef022 100644 Binary files a/src/graph/__pycache__/build_graph.cpython-311.pyc and b/src/graph/__pycache__/build_graph.cpython-311.pyc differ diff --git a/src/graph/__pycache__/build_reach_graph.cpython-311.pyc b/src/graph/__pycache__/build_reach_graph.cpython-311.pyc index 8abcc6ebe2b28e76c2a97992ca279d7005fb7acb..8acd5d252d92e60f5eaff4e66f7e58e1d550dbfb 100644 Binary files a/src/graph/__pycache__/build_reach_graph.cpython-311.pyc and b/src/graph/__pycache__/build_reach_graph.cpython-311.pyc differ diff --git a/src/graph/__pycache__/node_features.cpython-311.pyc b/src/graph/__pycache__/node_features.cpython-311.pyc index 164b0711d85a53045284547e2accb756019992e6..f07e795a20e7d15e9fd3f55a31f66a0fe2f8efd8 100644 Binary files a/src/graph/__pycache__/node_features.cpython-311.pyc and b/src/graph/__pycache__/node_features.cpython-311.pyc differ diff --git a/src/graph/build_graph.py b/src/graph/build_graph.py index 91f6048a2717dcb149a383689471288ba297579a..c4ca94206ab0fbb0338311c654f9906869e26fa0 100644 --- a/src/graph/build_graph.py +++ b/src/graph/build_graph.py @@ -249,6 +249,19 @@ def build_pyg_graph( (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: @@ -257,7 +270,8 @@ def build_pyg_graph( 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 columns. + 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 @@ -274,16 +288,26 @@ def build_pyg_graph( 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`. + 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", + } 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]) ] @@ -344,6 +368,21 @@ def build_pyg_graph( data.feature_names = all_feature_cols data.basin_id = torch.tensor(nodes_df["basin_id"].values, dtype=torch.long) + 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 diff --git a/src/graph/build_reach_graph.py b/src/graph/build_reach_graph.py index 7249d97781241d615d213dd430a11485d4a97fea..48bfd01d6cb6ce52e7167b94e2991edd120b0526 100644 --- a/src/graph/build_reach_graph.py +++ b/src/graph/build_reach_graph.py @@ -214,16 +214,25 @@ def normalize_toponym(name: Optional[str]) -> Optional[str]: return n or None -def _ancestors_within_km(G: nx.DiGraph, start, max_km: float) -> Dict: +def _ancestors_within_km(G: nx.MultiDiGraph, start, max_km: float) -> Dict: """Bounded backward BFS (via predecessors) from `start`, tracking - cumulative upstream distance. Returns {node: distance_km}.""" + cumulative upstream distance. Returns {node: distance_km}. + + G is a MultiDiGraph (parallel edges between the same two nodes are + real -- e.g. two distinct tronçons that both happen to run directly + between the same pair of hydrographic nodes, exactly the shape a + short braid takes). Where multiple parallel edges exist between a + predecessor and n, the shortest is used -- the conservative choice + for "is there any path within range," not an arbitrary one. + """ visited = {start: 0.0} frontier = [start] while frontier: new_frontier = [] for n in frontier: for pred in G.predecessors(n): - edge_dist = G.edges[pred, n].get("distance_km", 0.0) or 0.0 + parallel = G.get_edge_data(pred, n) # {key: data} for every parallel edge + edge_dist = min((d.get("distance_km", 0.0) or 0.0) for d in parallel.values()) cum = visited[n] + edge_dist if cum <= max_km and (pred not in visited or cum < visited[pred]): visited[pred] = cum @@ -256,11 +265,103 @@ def is_split_rejoin(G: nx.DiGraph, node, max_upstream_km: float = 15.0) -> bool: return len(common) > 0 +def find_split_points(G: nx.DiGraph) -> Set: + """ + Nodes where a single upstream channel divides into 2+ downstream + paths -- unambiguous by construction, unlike confluences/rejoins: + a split has exactly one thing flowing in and multiple things flowing + out, so it's always the same water partitioning, regardless of what + the branches get named downstream. No name or topology check needed + here the way `find_real_confluences` needs one -- out-degree >= 2 is + a sufficient physical definition on its own. + + Useful as the other half of a mass-conservation pair with + `find_rejoin_points`: for a genuine split-then-rejoin, flow at the + split should equal flow at the corresponding rejoin, since no new + catchment area is added anywhere in between. + """ + return {n for n in G.nodes if G.out_degree(n) >= 2} + + +def find_rejoin_points(G: nx.DiGraph, max_upstream_km: float = 15.0) -> Set: + """ + Nodes where two branches recombine after a common upstream split -- + the companion classification to `find_real_confluences`, made + explicit and persistable rather than silently discarded. A rejoin + and a real confluence look identical by raw in-degree; the + distinguishing test is the same one `find_real_confluences` already + uses (a common upstream ancestor within `max_upstream_km`), just + keeping the nodes that test *positively* identifies instead of the + ones it excludes. + + These matter for a different physics constraint than a real + confluence: no new mass enters at a rejoin (or its paired split) -- + what should hold is Q_upstream_of_split ~= Q_downstream_of_rejoin, + not "new tributary inflow contributes here." + """ + rejoins = set() + for n in G.nodes: + if G.in_degree(n) >= 2 and is_split_rejoin(G, n, max_upstream_km): + rejoins.add(n) + return rejoins + + +def pair_splits_and_rejoins(G: nx.DiGraph, max_upstream_km: float = 15.0) -> pd.DataFrame: + """ + For each rejoin point, identify which upstream split it pairs with + -- the actual link a mass-conservation constraint needs + (Q_at_this_split ~= Q_at_this_rejoin), not just "splits exist + somewhere and rejoins exist somewhere" as two independent facts. + + Returns: + DataFrame [rejoin_node, split_node, upstream_distance_km, + branch_nodes] -- one row per rejoin, `branch_nodes` is the list + of nodes along the shortest path of each branch between the + split and the rejoin (so edge-level attributes for the braid are + easy to pull later). A rejoin whose branches share more than one + common ancestor within range keeps the CLOSEST one (shortest + upstream distance), since that's the most specific matching split. + """ + rows = [] + for rejoin in find_rejoin_points(G, max_upstream_km): + in_edges = list(G.in_edges(rejoin)) + branch_sources = [u for u, _ in in_edges] + ancestor_sets = [_ancestors_within_km(G, src, max_upstream_km) for src in branch_sources] + common = set(ancestor_sets[0].keys()) + for s in ancestor_sets[1:]: + common &= set(s.keys()) + common.discard(rejoin) + if not common: + continue + # closest common ancestor = the actual split point for this rejoin + split = min(common, key=lambda n: max(a.get(n, float("inf")) for a in ancestor_sets)) + dist = max(a.get(split, 0.0) for a in ancestor_sets) + + branch_nodes = [] + for src in branch_sources: + try: + path = nx.shortest_path(G, split, src) + branch_nodes.append(path) + except nx.NetworkXNoPath: + branch_nodes.append([]) + + rows.append({ + "rejoin_node": rejoin, "split_node": split, + "upstream_distance_km": dist, "branch_nodes": branch_nodes, + }) + return pd.DataFrame(rows) + + def find_real_confluences(G: nx.DiGraph, max_upstream_km: float = 15.0) -> Set: """ Nodes where a genuinely different, independently-sourced river joins -- not just where in-degree happens to be >= 2, and not a channel - split that rejoins downstream (see `is_split_rejoin`). + split that rejoins downstream (see `find_rejoin_points`, which keeps + that classification rather than discarding it -- splits and rejoins + carry real mass-conservation information of their own, just a + different kind than a tributary confluence: no new catchment area + is added at either point in a braid, whereas a real confluence + genuinely does add one). Two filters, in order: 1. In-degree >= 2 with more than one distinct (normalized) river @@ -269,10 +370,11 @@ def find_real_confluences(G: nx.DiGraph, max_upstream_km: float = 15.0) -> Set: many same-river multi-inflow nodes with no real branching involved (confirmed against real data -- distances as short as 4.6m at some flagged "confluences"). - 2. NOT a split-rejoin: the branches must not trace back to a common + 2. NOT a rejoin: the branches must not trace back to a common upstream node within `max_upstream_km`. A braided or anabranched - stretch can satisfy filter 1 (different channel names) while - still being the same water recombining, not a real tributary. + stretch can satisfy filter 1 (different channel names -- e.g. a + named "bras"/arm) while still being the same water recombining, + not a real tributary. """ confluences = set() for n in G.nodes: @@ -362,7 +464,13 @@ def build_node_link_digraph( IDs; each edge carries the tronçon's geometry, distance_km, elevation_drop_m, toponym, basin_id, and cleabs as edge data. """ - G = nx.DiGraph() + G = nx.MultiDiGraph() # parallel edges between the same node pair are real + # data, not a collision -- two distinct tronçons directly connecting the + # same two hydrographic nodes is exactly what a short braid looks like. + # A plain DiGraph silently OVERWRITES the second such edge's data on + # add_edge rather than keeping both (confirmed as a real, silent data + # loss bug, not a hypothetical -- caught while testing a synthetic + # adjacent-split-rejoin case). n_sens, n_elev, n_ambiguous = 0, 0, 0 for f in features: @@ -568,6 +676,16 @@ def build_reach_graph_tables( # rather than each re-deriving it from the flat CSV, which is exactly how the # split-rejoin exclusion and the earlier normalization fix drifted out of sync # across files before. + split_points = find_split_points(G) + rejoin_points = find_rejoin_points(G) + # Splits and rejoins are NOT discarded as noise -- they carry a real, + # different mass-conservation constraint than a tributary confluence + # (Q_at_split ~= Q_at_rejoin, no new catchment area involved, vs. a + # confluence genuinely adding one). braid_id links a rejoin to its + # paired split so that constraint is actually constructible later, + # not just "some splits and rejoins exist somewhere, unpaired." + braid_pairs = pair_splits_and_rejoins(G) + rejoin_to_split = dict(zip(braid_pairs["rejoin_node"], braid_pairs["split_node"])) if not braid_pairs.empty else {} node_rows = {} for n, data in G.nodes(data=True): @@ -591,6 +709,8 @@ def build_reach_graph_tables( "station_code": str(n), "basin_id": basin_id, "latitude": lat, "longitude": lon, "elevation_m": elev, "is_gauged": False, "is_confluence": n in real_confluences, + "is_split_point": n in split_points, "is_rejoin_point": n in rejoin_points, + "braid_id": str(rejoin_to_split[n]) if n in rejoin_to_split else None, } edge_rows = [] @@ -612,6 +732,7 @@ def build_reach_graph_tables( "station_code": code, "basin_id": basin_id, "latitude": g["latitude"], "longitude": g["longitude"], "elevation_m": g.get("elevation_m"), "is_gauged": True, "is_confluence": False, + "is_split_point": False, "is_rejoin_point": False, "braid_id": None, "snap_distance_km": g.get("snap_distance_km"), } @@ -636,14 +757,20 @@ def build_reach_graph_tables( return nodes_df, edges_df -def insert_virtual_nodes(G: nx.DiGraph, spacing_km: float = 5.0) -> Tuple[nx.DiGraph, int]: +def insert_virtual_nodes(G: nx.MultiDiGraph, spacing_km: float = 5.0) -> Tuple[nx.MultiDiGraph, int]: """ Split any edge longer than `spacing_km` into evenly-spaced segments by inserting virtual node IDs along it, so long confluence-free stretches get intermediate nodes instead of one edge covering many kilometers. - Virtual node IDs are strings like "VIRTUAL::{u}::{v}::{i}" -- distinct - from real BD TOPO node IDs by construction, so they can never collide. + Virtual node IDs are strings like "VIRTUAL::{u}::{v}::{i}::{key}" -- + distinct from real BD TOPO node IDs by construction, so they can + never collide. The edge `key` is included specifically because two + parallel tronçons between the same (u, v) pair are real, not a + collision (see build_node_link_digraph) -- each needs its own + distinctly-IDed virtual chain, and removing the correct specific + parallel edge (not an arbitrary one) requires operating on its key, + not just its endpoints. Returns: (new_digraph, n_virtual_nodes_inserted) @@ -651,7 +778,7 @@ def insert_virtual_nodes(G: nx.DiGraph, spacing_km: float = 5.0) -> Tuple[nx.DiG G2 = G.copy() n_inserted = 0 - for u, v, data in list(G.edges(data=True)): + for u, v, key, data in list(G.edges(keys=True, data=True)): dist = data["distance_km"] if dist <= spacing_km: continue @@ -659,7 +786,7 @@ def insert_virtual_nodes(G: nx.DiGraph, spacing_km: float = 5.0) -> Tuple[nx.DiG coords = data["coords"] n_points = len(coords) - G2.remove_edge(u, v) + G2.remove_edge(u, v, key) prev_node = u for i in range(1, n_segments): frac = i / n_segments @@ -670,7 +797,7 @@ def insert_virtual_nodes(G: nx.DiGraph, spacing_km: float = 5.0) -> Tuple[nx.DiG lat = p0[1] + local_t * (p1[1] - p0[1]) z = (p0[2] + local_t * (p1[2] - p0[2])) if len(p0) > 2 else None - vnode = f"VIRTUAL::{u}::{v}::{i}" + vnode = f"VIRTUAL::{u}::{v}::{key}::{i}" G2.add_node(vnode, latitude=lat, longitude=lon, elevation_m=z) seg_dist = dist / n_segments seg_drop = (data["elevation_drop_m"] / n_segments) if data["elevation_drop_m"] is not None else None diff --git a/src/graph/node_features.py b/src/graph/node_features.py index f7ac47e18c54dc071d2237098c42b8fb622218ff..41dd6ed957a794eb463ed84c39d582280d258141 100644 --- a/src/graph/node_features.py +++ b/src/graph/node_features.py @@ -31,7 +31,7 @@ don't feed the `target_*` columns back in as model inputs. """ from dataclasses import dataclass, field from pathlib import Path -from typing import List, Optional +from typing import List, Optional, Tuple import numpy as np import pandas as pd @@ -73,6 +73,18 @@ def build_base_station_table(elevations_df: pd.DataFrame, prefix_map=None) -> pd return df[["station_code", "basin_id", "latitude", "longitude", "elevation_m"]].reset_index(drop=True) +def _haversine_km_vec(lat1: float, lon1: float, lat2, lon2): + """Vectorized haversine: one (lat1, lon1) point against arrays + lat2/lon2. Used for the exact narrow-phase distance check after a + KD-tree coarse prefilter (see add_groundwater_features).""" + R = 6371.0 + lat1r, lon1r = np.radians(lat1), np.radians(lon1) + lat2r, lon2r = np.radians(lat2), np.radians(lon2) + dlat, dlon = lat2r - lat1r, lon2r - lon1r + a = np.sin(dlat / 2) ** 2 + np.cos(lat1r) * np.cos(lat2r) * np.sin(dlon / 2) ** 2 + return R * 2 * np.arcsin(np.sqrt(a)) + + def add_idpr_features(nodes_df: pd.DataFrame, idpr_path: Path) -> pd.DataFrame: """ IDPR value per station. Prefers an exact join when the IDPR data @@ -140,44 +152,108 @@ def add_catchment_features(nodes_df: pd.DataFrame, catchment_path: Path) -> pd.D def add_groundwater_features( - nodes_df: pd.DataFrame, ades_path: Path, max_distance_km: float = 20.0 + nodes_df: pd.DataFrame, ades_path: Path, max_distance_km: float = 20.0, + date_range: Optional[Tuple[str, str]] = None, ) -> pd.DataFrame: """ - Radius-averaged groundwater level/depth per station, via - ADESLoader.aggregate_to_stations (existing method — reused as-is, - not reimplemented). That method returns a per-date time series per - station; here we take each station's most recent aggregate as its - static feature value. + Radius-averaged groundwater level/depth per station. + + Does NOT use ADESLoader.aggregate_to_stations -- that method loops + per station and does a full haversine .apply() over the ENTIRE + groundwater dataframe for each one (O(n_stations * n_readings)). + At 27 stations against ~272k readings that's slow but tolerable; at + the reach graph's ~4,500 nodes it's over a billion row-wise Python + calls, not a proportionally-worse runtime but an unusable one. + + Instead: reduce to each well's most recent reading first (collapses + ~272k readings to ~100 wells -- we only need each station's latest + aggregate anyway, same as the old code kept via + `.sort_values("date").drop_duplicates(...)` after the fact), then a + KD-tree coarse prefilter narrows each station to a handful of nearby + candidate wells, with exact haversine distance computed only on + that small candidate set to get correct radius membership and + correct averages -- fast without trading away correctness for it. + + date_range: if given, (start, end) date strings (inclusive) -- a + well's readings are filtered to this window BEFORE picking its + "most recent" one, so a well with no reading in the window is + correctly excluded from that call rather than falling back to some + older reading from outside the period the rest of the graph's + features are being built for (real wells in this project report on + wildly different schedules -- see the ground-truth investigation + that originally found this "most recent" logic needed care). """ + from scipy.spatial import cKDTree + loader = ADESLoader(data_path=ades_path) gw_df = loader.load() + if date_range is not None: + start, end = date_range + gw_df = gw_df[(gw_df["date"] >= start) & (gw_df["date"] <= end)] + latest_per_well = gw_df.sort_values("date").drop_duplicates("code_bss", keep="last").reset_index(drop=True) - station_coords = nodes_df.rename(columns={"latitude": "lat", "longitude": "lon"})[ - ["station_code", "lat", "lon"] - ] - agg = loader.aggregate_to_stations(gw_df, station_coords, max_distance_km=max_distance_km) - if agg.empty: + if latest_per_well.empty: out = nodes_df.copy() out["avg_groundwater_level_m"] = np.nan out["avg_groundwater_depth_m"] = np.nan out["n_nearby_wells"] = 0 return out - latest = agg.sort_values("date").drop_duplicates("station_code", keep="last") - latest = latest[["station_code", "avg_groundwater_level_m", "avg_groundwater_depth_m", "n_wells"]] - latest = latest.rename(columns={"n_wells": "n_nearby_wells"}) + well_lat = latest_per_well["lat"].values + well_lon = latest_per_well["lon"].values + tree = cKDTree(np.column_stack([well_lon, well_lat])) + + # Coarse degree-radius, deliberately oversized (1.5x safety factor): + # degrees-to-km varies with latitude and direction (lon degrees are + # worth fewer km than lat degrees away from the equator), so this + # only needs to be a safe upper bound, not an accurate radius -- + # the exact haversine pass afterward is what determines real + # membership and distances. + coarse_radius_deg = (max_distance_km / 111.0) * 1.5 + + station_lon = nodes_df["longitude"].values + station_lat = nodes_df["latitude"].values + candidate_lists = tree.query_ball_point(np.column_stack([station_lon, station_lat]), r=coarse_radius_deg) + + avg_levels, avg_depths, n_wells_list = [], [], [] + for i, candidates in enumerate(candidate_lists): + if not candidates: + avg_levels.append(np.nan); avg_depths.append(np.nan); n_wells_list.append(0) + continue + cand_idx = np.array(candidates) + dist = _haversine_km_vec(station_lat[i], station_lon[i], well_lat[cand_idx], well_lon[cand_idx]) + within = dist <= max_distance_km + if not within.any(): + avg_levels.append(np.nan); avg_depths.append(np.nan); n_wells_list.append(0) + continue + sel = latest_per_well.iloc[cand_idx[within]] + avg_levels.append(sel["groundwater_level_m"].mean()) + avg_depths.append(sel["groundwater_depth_m"].mean() if "groundwater_depth_m" in sel.columns else np.nan) + n_wells_list.append(int(within.sum())) - out = nodes_df.merge(latest, on="station_code", how="left") - out["n_nearby_wells"] = out["n_nearby_wells"].fillna(0).astype(int) + out = nodes_df.copy() + out["avg_groundwater_level_m"] = avg_levels + out["avg_groundwater_depth_m"] = avg_depths + out["n_nearby_wells"] = n_wells_list return out -def add_safran_features(nodes_df: pd.DataFrame, safran_path: Path) -> pd.DataFrame: +def add_safran_features( + nodes_df: pd.DataFrame, safran_path: Path, date_range: Optional[Tuple[str, str]] = None, +) -> pd.DataFrame: """ Climate summary stats per station, interpolated from ERA5 via SAFRANLoader's own station-point interpolation. Requires `xarray` and actual era5_*.nc files under safran_path; returns nodes_df unchanged (with a clear message) if either is missing. + + date_range: if given, (start, end) date strings (inclusive) -- + filters ERA5 timesteps to this window BEFORE aggregating (mean for + temp/wind/solar, sum for precip/evap/snow/runoff), so e.g. + climate_precip_mm becomes "total precip over the window" rather + than "total precip over the full 1960-2026 record", which wouldn't + be comparable to a discharge target computed over a specific + training period. """ try: from ..data.loaders.safran import SAFRANLoader @@ -201,6 +277,13 @@ def add_safran_features(nodes_df: pd.DataFrame, safran_path: Path) -> pd.DataFra return nodes_df df = loader.convert_units(df) + if date_range is not None: + start, end = date_range + df = df[(df["date"] >= start) & (df["date"] <= end)] + if df.empty: + print(f"add_safran_features: skipped (no ERA5 timesteps fall within {date_range})") + return nodes_df + agg_spec = {c: "mean" for c in ["temp_C", "wind_speed_ms", "solar_Wm2"] if c in df.columns} for c in ["precip_mm", "evap_mm", "snow_mm", "runoff_mm"]: if c in df.columns: @@ -214,11 +297,19 @@ def add_safran_features(nodes_df: pd.DataFrame, safran_path: Path) -> pd.DataFra return nodes_df.merge(summary, on="station_code", how="left") -def add_hydrometric_target_stats(nodes_df: pd.DataFrame, hydrometric_path: Path) -> pd.DataFrame: +def add_hydrometric_target_stats( + nodes_df: pd.DataFrame, hydrometric_path: Path, date_range: Optional[Tuple[str, str]] = None, +) -> pd.DataFrame: """ Discharge/water-level summary stats per gauge, prefixed `target_`. Left-joined so ungauged stations get NaN rather than being dropped — that NaN is meaningful (no observations), not a gap to fill in. + + date_range: if given, (start, end) date strings (inclusive) -- + filters observations to this window before computing mean/std/count, + so the target reflects the same training period the input features + (climate, groundwater) were built for, not each source's own full + and differently-shaped historical record. """ from ..data.loaders.hydrometric import HydrometricLoader @@ -229,6 +320,10 @@ def add_hydrometric_target_stats(nodes_df: pd.DataFrame, hydrometric_path: Path) print(f"add_hydrometric_target_stats: skipped ({e})") return nodes_df + if date_range is not None: + start, end = date_range + df = df[(df["date"] >= start) & (df["date"] <= end)] + agg_spec = {} if "discharge_m3s" in df.columns: agg_spec["discharge_m3s"] = ["mean", "std", "count"] @@ -243,27 +338,62 @@ def add_hydrometric_target_stats(nodes_df: pd.DataFrame, hydrometric_path: Path) def build_node_features( - station_elevations_path: Path, + station_elevations_path: Optional[Path] = None, idpr_path: Optional[Path] = None, ades_path: Optional[Path] = None, safran_path: Optional[Path] = None, hydrometric_path: Optional[Path] = None, catchment_path: Optional[Path] = None, prefix_map=None, + base_nodes_df: Optional[pd.DataFrame] = None, + date_range: Optional[Tuple[str, str]] = None, ) -> "tuple[pd.DataFrame, NodeFeatureReport]": """ Orchestrator: build the base station table, then layer on every available loader's features. Any path left as None (or pointing to missing files) is skipped with a message rather than raising. + Two ways to get the base table: + - `station_elevations_path` (original path): builds the 27-gauge + table via StationElevationsLoader, as before. + - `base_nodes_df` (for the reach graph): pass an already-built + node table directly -- e.g. the ~4,500-node reach graph from + scripts/build_reach_graphs.py (real gauges + real confluences + + splits/rejoins + virtual infill nodes), which every add_*_features + function already works against unchanged, since they only ever + assumed [station_code, latitude, longitude, ...] columns, not a + specific row count. Extra columns already on base_nodes_df + (is_gauged, is_confluence, is_split_point, is_rejoin_point, + braid_id, snap_distance_km) pass through every merge untouched. + Exactly one of these two must be given. + + date_range: optional (start, end) date strings (inclusive), applied + to every time-varying source (groundwater, climate, hydrometric + targets) so all three describe the same period -- otherwise each + source silently aggregates over its own full, differently-shaped + historical record (ADES wells reporting from 1972-2026 on wildly + different schedules, ERA5 spanning 1960-2026, hydrometric records + with their own per-station date ranges), which makes "climate over + the period" and "discharge over the period" describe different + periods without anyone intending that. Static sources (IDPR, + elevation, catchment area) are unaffected -- they don't vary in time. + Returns: (nodes_df, report) """ - nodes_df = build_base_station_table( - StationElevationsLoader(data_path=station_elevations_path).load(), prefix_map - ) - report = NodeFeatureReport(n_stations=len(nodes_df), included=["station_elevations (base table)"]) + if (station_elevations_path is None) == (base_nodes_df is None): + raise ValueError("Provide exactly one of station_elevations_path or base_nodes_df.") + + if base_nodes_df is not None: + nodes_df = base_nodes_df.copy() + report = NodeFeatureReport(n_stations=len(nodes_df), included=["reach graph node table (base)"]) + else: + nodes_df = build_base_station_table( + StationElevationsLoader(data_path=station_elevations_path).load(), prefix_map + ) + report = NodeFeatureReport(n_stations=len(nodes_df), included=["station_elevations (base table)"]) + date_aware = {"ades_groundwater", "safran_climate", "hydrometric_targets"} for label, path, fn in [ ("idpr", idpr_path, add_idpr_features), ("catchment_area", catchment_path, add_catchment_features), @@ -275,8 +405,12 @@ def build_node_features( report.skipped.append(f"{label} (no path given or not found)") continue try: - nodes_df = fn(nodes_df, path) - report.included.append(label) + if date_range is not None and label in date_aware: + nodes_df = fn(nodes_df, path, date_range=date_range) + else: + nodes_df = fn(nodes_df, path) + report.included.append(label if date_range is None or label not in date_aware + else f"{label} (filtered to {date_range[0]}..{date_range[1]})") except Exception as e: report.skipped.append(f"{label} (error: {e})") diff --git a/src/graph/physics_losses.py b/src/graph/physics_losses.py new file mode 100644 index 0000000000000000000000000000000000000000..05632c4a7807dc0b9b03baef5b92562583217328 --- /dev/null +++ b/src/graph/physics_losses.py @@ -0,0 +1,348 @@ +""" +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) -> float: + """Mean squared residual -- works identically on numpy or torch.""" + return (residual ** 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 \ No newline at end of file