""" Builds genuine [n_nodes, T] time series for the features that node_features.py's add_safran_features/add_groundwater_features/ add_hydrometric_target_stats only ever hand over as a single period- aggregated number. That aggregation was always a real, named limitation (see build_graph.py's DYNAMIC_PREFIXES docstring and the README's ยง2.3) -- this module is what actually closes it. physics_losses.py's routing_consistency_loss specifically needs a real time dimension and had nothing to consume before this existed. Three sources, three different real challenges: - discharge: already daily, already per-station -- just needs pivoting to wide form, restricted to real gauges (no fabrication for non-gauge nodes). - groundwater: wells report on wildly irregular schedules (confirmed against real data: 13 different "latest dates" among 18 wells within 20km of one station). Needs resampling to a common daily grid via forward-fill (a well's level changes slowly -- carrying the last known reading forward is standard practice for sparse level data, not an invented shortcut) BEFORE the spatial average, not after -- averaging raw irregular readings per exact calendar date is exactly the bug that made the old "aggregate_to_stations" path undercount real well coverage by 5-10x (see node_features.py's add_groundwater_features docstring). - climate: SAFRANLoader.load() already returns per-date rows before node_features.py's add_safran_features aggregates them -- this just skips that aggregation step and pivots instead. UNTESTED in this environment: no real ERA5/safran files were available here to validate against (only real ADES and hydrometric data were). """ from pathlib import Path from typing import Dict, List, Optional, Tuple import numpy as np import pandas as pd def build_discharge_timeseries( nodes_df: pd.DataFrame, hydrometric_path: Path, date_range: Tuple[str, str], freq: str = "D", ) -> pd.DataFrame: """ Real daily QmnJ discharge, pivoted to [date x station_code] wide form, restricted to the date range and reindexed onto a regular daily grid (missing days stay NaN -- no fabrication). Deliberately does NOT forward-fill discharge the way groundwater gets resampled below: discharge genuinely changes day to day and a missing daily reading should stay missing, not be papered over with the previous day's value the way a slowly-changing water table can reasonably be. Only real gauge station_codes appear as columns -- confluences and virtual nodes never had discharge observations to begin with, same principle as the period-aggregate target_* columns. """ from ..data.loaders.hydrometric import HydrometricLoader loader = HydrometricLoader(data_path=hydrometric_path) df = loader.load() if "discharge_m3s" not in df.columns: raise ValueError("Loaded hydrometric data has no discharge_m3s column") start, end = date_range df = df[(df["date"] >= start) & (df["date"] <= end)] wide = df.pivot_table(index="date", columns="station_code", values="discharge_m3s", aggfunc="mean") full_index = pd.date_range(start, end, freq=freq) wide = wide.reindex(full_index) wide.index.name = "date" return wide def build_groundwater_timeseries( nodes_df: pd.DataFrame, ades_path: Path, date_range: Tuple[str, str], max_distance_km: float = 20.0, freq: str = "D", ) -> Tuple[pd.DataFrame, pd.DataFrame]: """ Real groundwater level, resampled to a daily grid per well (forward-fill) BEFORE spatial averaging, then averaged per node using the SAME nearby-well set every day -- the spatial join (which wells are near which node) is time-invariant, so it's computed once, not repeated per date. That's what keeps this tractable at real reach-graph scale: O(n_nodes) spatial lookups total, not O(n_nodes x n_dates). Returns: (level_wide, depth_wide) -- both [date x station_code], node ordering matching nodes_df. A node with zero wells within max_distance_km gets NaN for every date, not zero -- "no nearby monitoring" is meaningfully different from "measured zero," same principle as node_features.py's static version. """ from ..data.loaders.ades import ADESLoader from scipy.spatial import cKDTree loader = ADESLoader(data_path=ades_path) gw_df = loader.load() start, end = date_range full_index = pd.date_range(start, end, freq=freq) # Resample each well independently to the common daily grid, forward-fill. wells_wide = gw_df.pivot_table(index="date", columns="code_bss", values="groundwater_level_m", aggfunc="mean") wells_wide = wells_wide.reindex(wells_wide.index.union(full_index)).sort_index().ffill() wells_wide = wells_wide.reindex(full_index) depth_wide_raw = gw_df.pivot_table(index="date", columns="code_bss", values="groundwater_depth_m", aggfunc="mean") \ if "groundwater_depth_m" in gw_df.columns else None if depth_wide_raw is not None: depth_wide_raw = depth_wide_raw.reindex(depth_wide_raw.index.union(full_index)).sort_index().ffill() depth_wide_raw = depth_wide_raw.reindex(full_index) well_coords = gw_df.drop_duplicates("code_bss").set_index("code_bss")[["lat", "lon"]] well_coords = well_coords.reindex(wells_wide.columns) # align to the pivoted columns' order tree = cKDTree(well_coords[["lon", "lat"]].values) coarse_radius_deg = (max_distance_km / 111.0) * 1.5 def _haversine_km(lat1, lon1, lat2, lon2): R = 6371.0 lat1r, lon1r, lat2r, lon2r = map(np.radians, [lat1, lon1, lat2, 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)) 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) level_cols, depth_cols = {}, {} for i, station_code in enumerate(nodes_df["station_code"]): candidates = candidate_lists[i] if not candidates: level_cols[station_code] = pd.Series(np.nan, index=full_index) depth_cols[station_code] = pd.Series(np.nan, index=full_index) continue cand_idx = np.array(candidates) cand_lat = well_coords["lat"].values[cand_idx] cand_lon = well_coords["lon"].values[cand_idx] dist = _haversine_km(station_lat[i], station_lon[i], cand_lat, cand_lon) within = cand_idx[dist <= max_distance_km] if len(within) == 0: level_cols[station_code] = pd.Series(np.nan, index=full_index) depth_cols[station_code] = pd.Series(np.nan, index=full_index) continue well_names = wells_wide.columns[within] level_cols[station_code] = wells_wide[well_names].mean(axis=1) if depth_wide_raw is not None: depth_cols[station_code] = depth_wide_raw[well_names].mean(axis=1) else: depth_cols[station_code] = pd.Series(np.nan, index=full_index) level_wide = pd.DataFrame(level_cols) depth_wide = pd.DataFrame(depth_cols) return level_wide, depth_wide def build_climate_timeseries( nodes_df: pd.DataFrame, safran_path: Path, date_range: Tuple[str, str], ) -> Dict[str, pd.DataFrame]: """ Real per-date ERA5 variables, pivoted per variable to [date x station_code] wide form -- skips node_features.py's add_safran_ features aggregation step entirely rather than reversing it. UNTESTED in this environment: no real era5_*.nc files were available here (only ADES and hydrometric real data were). The logic mirrors build_discharge_timeseries's pivot pattern directly, but please verify the real output shape/values before trusting it for training. """ from ..data.loaders.safran import SAFRANLoader station_coords = nodes_df.rename(columns={"latitude": "lat", "longitude": "lon"})[ ["station_code", "lat", "lon"] ] loader = SAFRANLoader(data_path=safran_path, station_coords=station_coords) df = loader.load() df = loader.convert_units(df) start, end = date_range df = df[(df["date"] >= start) & (df["date"] <= end)] var_cols = [c for c in df.columns if c not in ("date", "station_code")] result = {} for var in var_cols: wide = df.pivot_table(index="date", columns="station_code", values=var, aggfunc="mean") result[var] = wide return result def assemble_dynamic_tensor(nodes_df: pd.DataFrame, wide_df: pd.DataFrame) -> Tuple[np.ndarray, List[pd.Timestamp]]: """ Aligns a [date x station_code] wide DataFrame onto nodes_df's own row order, so the result matches data.x's node indexing exactly -- this is the piece that makes a build_*_timeseries output directly usable as physics_losses.py's routing_consistency_loss's `Q` argument (shape [n_nodes, T]). A node whose station_code never appears as a column (every confluence/virtual node, for discharge) gets an all-NaN row, not a dropped row -- shape stays [n_nodes, T] regardless of coverage. """ aligned = wide_df.reindex(columns=nodes_df["station_code"]) return aligned.values.T, list(wide_df.index)