from __future__ import annotations import datetime as dt import random from pathlib import Path import geopandas as gpd import numpy as np import pandas as pd import torch import torch.nn as nn import torch.nn.functional as F from tqdm import tqdm from shapely.geometry import Polygon from torch_geometric_temporal.nn.recurrent import GConvGRU import logging # Simple in-process cache for ERA5 bulk results to avoid repeated GEE requests _GEE_CACHE: dict = {} logger = logging.getLogger(__name__) _EE = None SCRIPT_DIR = Path(__file__).resolve().parent PROJECT_DIR = SCRIPT_DIR.parent WEIGHT_PATH = SCRIPT_DIR / "gnn" / "glof_gnn_weights_final.pth" EDGE_INDEX_PATH = SCRIPT_DIR / "edge_index3.pt" NODE_MAPPING_PATH = SCRIPT_DIR / "node_mapping.csv" NODE_GEOMETRY_PATH = PROJECT_DIR.parent / "gis_data" / "stage2_dataset_15.gpkg" EARTH_ENGINE_PROJECT = "mtp-phase-2" FEATURE_COLS = [ "temperature_2m", "dewpoint_temperature_2m", "surface_pressure", "total_precipitation_sum", "snowfall_sum", "snow_depth_water_equivalent", "snowmelt_sum", "surface_runoff_sum", "sub_surface_runoff_sum", ] TARGET_COLS = [ "snowmelt_sum", "surface_runoff_sum", "sub_surface_runoff_sum", ] AOI_COORDS = [ (70.203815, 38.355372), (99.306191, 34.771871), (99.161404, 24.781503), (80.628547, 26.37417), (71.072543, 34.192719), (70.203815, 38.355372), ] AOI_POLYGON = Polygon(AOI_COORDS) # Dates are interpreted as DD/MM/YYYY. TEST_DATES = ["02/02/2026", "08/04/2026", "05/05/2026", "03/02/2026"] HISTORY_DAYS = 60 FORECAST_DAYS = 3 SEQ_LEN = HISTORY_DAYS + FORECAST_DAYS NODE_HOPS = 2 RANDOM_SEED = 42 class GLOFPredictor(nn.Module): def __init__(self, node_features=9, hidden_channels=64, out_steps=3, out_vars=3, K=2): super().__init__() self.out_steps = out_steps self.out_vars = out_vars self.recurrent_1 = GConvGRU(in_channels=node_features, out_channels=hidden_channels, K=K) self.recurrent_2 = GConvGRU(in_channels=hidden_channels, out_channels=hidden_channels, K=K) self.mlp_head = nn.Sequential( nn.Linear(hidden_channels, hidden_channels), nn.ReLU(), nn.Dropout(0.2), nn.Linear(hidden_channels, out_steps * out_vars) ) self.register_buffer("edge_index", torch.empty((2, 25094), dtype=torch.long)) def forward(self, x, edge_index, edge_weight=None): if x.dim() == 3: x = x.unsqueeze(0) B, N, T, F_in = x.shape if B > 1: E = edge_index.size(1) offsets = (torch.arange(B, device=edge_index.device) * N).view(B, 1, 1) edge_index = (edge_index.unsqueeze(0) + offsets).permute(1, 0, 2).reshape(2, B * E) if edge_weight is not None: edge_weight = edge_weight.repeat(B) x = x.view(B * N, T, F_in) h1, h2 = None, None for t in range(T): h1 = self.recurrent_1(x[:, t, :], edge_index, edge_weight, H=h1) h1 = F.relu(h1) h2 = self.recurrent_2(h1, edge_index, edge_weight, H=h2) h2 = F.relu(h2) out = self.mlp_head(h2) return out.view(B, N, self.out_steps, self.out_vars) def initialize_earth_engine() -> None: ee = get_earth_engine() try: # Check if we have a service account JSON string in the environment import os import json ee_creds_json = os.environ.get("EE_SERVICE_ACCOUNT_JSON") if ee_creds_json: creds_dict = json.loads(ee_creds_json) credentials = ee.ServiceAccountCredentials(creds_dict['client_email'], key_data=ee_creds_json) ee.Initialize(credentials, project=EARTH_ENGINE_PROJECT) else: ee.Initialize(project=EARTH_ENGINE_PROJECT) except Exception as e: print(f"Failed to initialize EE: {e}. Attempting interactive auth...") ee.Authenticate() ee.Initialize(project=EARTH_ENGINE_PROJECT) def get_earth_engine(): global _EE if _EE is not None: return _EE try: import ee as earth_engine except ModuleNotFoundError as exc: if exc.name == "fcntl": raise ImportError( "The wrong PyPI package named 'ee' is installed and is shadowing " "Google Earth Engine. Run this in the geo_dl environment: " "pip uninstall ee && pip install earthengine-api" ) from exc raise if not hasattr(earth_engine, "ImageCollection"): raise ImportError( "Imported package 'ee' is not Google Earth Engine. Run: " "pip uninstall ee && pip install earthengine-api" ) _EE = earth_engine return _EE def load_graph_nodes() -> gpd.GeoDataFrame: mapping = pd.read_csv(NODE_MAPPING_PATH) lakes = gpd.read_file(NODE_GEOMETRY_PATH) merge_cols = ["GLAKE_ID", "era5_fid"] node_cols = merge_cols + ["geometry"] nodes = mapping.merge(lakes[node_cols], on=merge_cols, how="left") missing_geom = nodes["geometry"].isna().sum() if missing_geom: raise ValueError(f"{missing_geom} graph nodes could not be matched to lake geometries.") nodes_gdf = gpd.GeoDataFrame(nodes, geometry="geometry", crs=lakes.crs) nodes_gdf = nodes_gdf.sort_values("node_index").reset_index(drop=True) if not nodes_gdf.geometry.geom_type.eq("Point").all(): nodes_gdf["geometry"] = nodes_gdf.geometry.centroid return nodes_gdf def choose_random_glake_nodes( nodes_gdf: gpd.GeoDataFrame, polygon: Polygon, count: int, edge_index: torch.Tensor | None = None, ) -> gpd.GeoDataFrame: nodes_wgs84 = nodes_gdf.to_crs("EPSG:4326").copy() candidates = nodes_wgs84[nodes_wgs84.geometry.within(polygon)] if edge_index is not None: connected_nodes = set(torch.unique(edge_index.cpu()).numpy().astype(int).tolist()) candidates = candidates[candidates["node_index"].isin(connected_nodes)] if candidates.empty: raise ValueError("No connected trained graph nodes / GLAKE_IDs were found inside the AOI polygon.") candidate_indices = candidates.index.tolist() if len(candidate_indices) >= count: chosen_indices = random.sample(candidate_indices, count) else: print( f"Only {len(candidate_indices)} GLAKE_IDs found inside AOI; " "sampling with replacement." ) chosen_indices = random.choices(candidate_indices, k=count) return nodes_wgs84.loc[chosen_indices].reset_index(drop=True) def build_k_hop_subgraph(edge_index: torch.Tensor, center_node: int, hops: int = 2): src = edge_index[0].cpu().numpy() dst = edge_index[1].cpu().numpy() selected = {int(center_node)} frontier = {int(center_node)} for _ in range(hops): if not frontier: break frontier_arr = np.fromiter(frontier, dtype=np.int64) mask = np.isin(src, frontier_arr) | np.isin(dst, frontier_arr) neighbors = set(src[mask].astype(int)).union(set(dst[mask].astype(int))) frontier = neighbors - selected selected.update(neighbors) node_ids = sorted(selected) node_to_local = {node_id: i for i, node_id in enumerate(node_ids)} edge_mask = np.isin(src, node_ids) & np.isin(dst, node_ids) sub_src = src[edge_mask] sub_dst = dst[edge_mask] if len(sub_src) == 0: sub_edge_index = torch.empty((2, 0), dtype=torch.long) else: sub_edges = [[node_to_local[int(s)], node_to_local[int(d)]] for s, d in zip(sub_src, sub_dst)] sub_edge_index = torch.tensor(sub_edges, dtype=torch.long).t().contiguous() center_local = node_to_local[int(center_node)] return node_ids, sub_edge_index, center_local def parse_date_ddmmyyyy(date_text: str) -> dt.date: return dt.datetime.strptime(date_text, "%d/%m/%Y").date() def fetch_era5_land_daily_for_nodes( nodes_gdf: gpd.GeoDataFrame, target_date: dt.date, history_days: int = HISTORY_DAYS, forecast_days: int = FORECAST_DAYS, ) -> tuple[np.ndarray, list[dt.date]]: """Fetch 90 history days. The next 3 forecast days are naturally NaN for history array filling. Optimized logic: 1. Rounds coordinates to 1 decimal place (~11km, ERA5 resolution) to prevent redundant queries. 2. Uses Earth Engine .map() to process all 90 days in a single backend operation, bringing data back in one getInfo(). """ ee = get_earth_engine() # We fetch up to today/target_date. We need strictly history_days worth of past data. # We add 1 day so that endDate is exclusive and covers target_date. coll = ee.ImageCollection("ECMWF/ERA5_LAND/DAILY_AGGR").select(FEATURE_COLS) latest_img = coll.sort('system:time_start', False).first() if latest_img is None: raise RuntimeError("No ERA5-Land images found in collection") latest_date_str = latest_img.date().format('YYYY-MM-dd').getInfo() latest_date = dt.date.fromisoformat(latest_date_str) # We fetch history_days up to the latest available date. Days after latest_date # (including target_date if it is newer) will remain as NaN and be imputed later. fetch_end_date = ee.Date(latest_date.isoformat()).advance(1, 'day') fetch_start_date = ee.Date(latest_date.isoformat()).advance(-history_days + 1, 'day') # all_dates still cover the historical window plus forecast days so downstream # code keeps the 90 + 3 shape and we silently impute the missing forward days. start_date = latest_date - dt.timedelta(days=history_days - 1) all_dates = [start_date + dt.timedelta(days=i) for i in range(history_days + forecast_days)] nodes_wgs84 = nodes_gdf.to_crs("EPSG:4326") # Optimization: Deduplicate coordinates for ERA5 resolution (approx 0.1 degree) unique_coords = {} node_to_coord = {} for local_idx, row in enumerate(nodes_wgs84.itertuples()): lon, lat = float(row.geometry.x), float(row.geometry.y) # Snap to 0.1 degree intervals (~11km spacing of ERA5-Land) snapped_lon, snapped_lat = round(lon, 1), round(lat, 1) coord_key = f"{snapped_lon}_{snapped_lat}" node_to_coord[local_idx] = coord_key if coord_key not in unique_coords: unique_coords[coord_key] = {"lon": snapped_lon, "lat": snapped_lat} ee_features = [] for coord_key, coords in unique_coords.items(): ee_features.append( ee.Feature( ee.Geometry.Point([coords["lon"], coords["lat"]]), {"coord_key": coord_key} ) ) point_fc = ee.FeatureCollection(ee_features) values = np.full((len(all_dates), len(nodes_wgs84), len(FEATURE_COLS)), np.nan, dtype=np.float32) base_collection = coll.filterDate(fetch_start_date, fetch_end_date) def extract_stats(image): date_str = image.date().format("YYYY-MM-dd") stats = image.reduceRegions(collection=point_fc, reducer=ee.Reducer.mean(), scale=11132) return stats.map(lambda f: f.set("date", date_str)) # Chunk the fetch so we can show a progress bar while still reducing the # number of getInfo() calls. We fetch in chunks of days (default 10). chunk_days = 10 cache_key = (tuple(sorted(unique_coords.keys())), start_date.isoformat(), latest_date.isoformat(), chunk_days) if cache_key in _GEE_CACHE: print(f"[ERA5] Reusing cached live fetch for {len(unique_coords)} unique pixels over {history_days} days...") logger.debug("Reusing cached ERA5 bulk results for key=%s", cache_key) results = _GEE_CACHE[cache_key] else: print(f"[ERA5] Fetching live ERA5 for {len(unique_coords)} unique pixels over {history_days} days...") logger.info("Fetching ERA5 for %d unique pixels over %d days (chunks=%d)...", len(unique_coords), history_days, chunk_days) results_features = [] # build list of python dates for chunking python_start = start_date total_steps = history_days * len(unique_coords) with tqdm(total=history_days, desc="Fetching ERA5 (days)", leave=True) as day_pbar: for i in range(0, history_days, chunk_days): chunk_start = python_start + dt.timedelta(days=i) chunk_end = min(python_start + dt.timedelta(days=i + chunk_days), latest_date + dt.timedelta(days=1)) # filter the collection for this chunk chunk_coll = coll.filterDate(chunk_start.isoformat(), chunk_end.isoformat()) chunk_results = chunk_coll.map(extract_stats).flatten().getInfo() # extend features features = chunk_results.get("features", []) results_features.extend(features) # advance progress by the number of days fetched in the chunk day_pbar.update((chunk_end - chunk_start).days) results = {"features": results_features} _GEE_CACHE[cache_key] = results print(f"[ERA5] Completed live fetch for {len(unique_coords)} unique pixels over {history_days} days.") # Map the unified results back to the individual lake nodes # First, create a dictionary of day -> coord_key -> features stats_dict = {} for feature in results.get("features", []): props = feature.get("properties", {}) date_str = props.get("date") key = props.get("coord_key") if date_str and key: if date_str not in stats_dict: stats_dict[date_str] = {} stats_dict[date_str][key] = props # Parse and fill array for t_idx, day in enumerate(all_dates): day_str = day.isoformat() if day_str not in stats_dict: continue day_data = stats_dict[day_str] for local_idx in range(len(nodes_wgs84)): coord_key = node_to_coord[local_idx] props = day_data.get(coord_key, {}) for f_idx, col in enumerate(FEATURE_COLS): val = props.get(col) if val is not None: values[t_idx, local_idx, f_idx] = float(val) return values, all_dates def build_model_input(raw_window: np.ndarray) -> tuple[torch.Tensor, np.ndarray]: """Apply the same normalization logic used in training. Training used node-wise mean/std from the 90-day history window, not one global scaler. The model target stayed unnormalized, so predictions are compared directly with raw ERA5 hydrology values. """ if raw_window.shape[0] != SEQ_LEN: raise ValueError(f"Expected {SEQ_LEN} timesteps, got {raw_window.shape[0]}.") target = raw_window[HISTORY_DAYS:, :, 6:9].copy() mean = np.nanmean(raw_window[:HISTORY_DAYS], axis=0, keepdims=True) std = np.nanstd(raw_window[:HISTORY_DAYS], axis=0, keepdims=True) + 1e-8 normalized_window = (raw_window - mean) / std history = normalized_window[:HISTORY_DAYS] forecast_input = normalized_window[HISTORY_DAYS:].copy() forecast_input[:, :, 6:9] = 0.0 x_seq = np.concatenate([history, forecast_input], axis=0) x_seq = np.nan_to_num(x_seq, nan=0.0, posinf=0.0, neginf=0.0) x_seq = x_seq.transpose(1, 0, 2) x_tensor = torch.from_numpy(np.ascontiguousarray(x_seq, dtype=np.float32)).unsqueeze(0) return x_tensor, target def masked_metrics(pred: np.ndarray, target: np.ndarray) -> dict[str, float]: mask = ~np.isnan(target) if not mask.any(): return {"mae": float("nan"), "rmse": float("nan")} diff = pred[mask] - target[mask] return { "mae": float(np.mean(np.abs(diff))), "rmse": float(np.sqrt(np.mean(diff**2))), } def load_model_and_graph(device: torch.device): edge_index = torch.load(EDGE_INDEX_PATH, map_location="cpu").long() model = GLOFPredictor(node_features=9, hidden_channels=64, out_steps=3, out_vars=3, K=2).to(device) model.load_state_dict(torch.load(WEIGHT_PATH, map_location=device)) model.eval() return model, edge_index _CACHED_MODEL_AND_GRAPH = None def predict_for_lake(glake_id: str, target_date: dt.date = None) -> str: global _CACHED_MODEL_AND_GRAPH if target_date is None: target_date = dt.date.today() initialize_earth_engine() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") if _CACHED_MODEL_AND_GRAPH is None: _CACHED_MODEL_AND_GRAPH = load_model_and_graph(device) model, edge_index = _CACHED_MODEL_AND_GRAPH nodes_gdf = load_graph_nodes() # Find the node for the given GLAKE_ID lake_row = nodes_gdf[nodes_gdf["GLAKE_ID"] == glake_id] if lake_row.empty: return f"Error: No trained graph node found for GLAKE_ID '{glake_id}'." center_node = int(lake_row["node_index"].iloc[0]) node_ids, sub_edge_index, center_local = build_k_hop_subgraph(edge_index, center_node, hops=NODE_HOPS) sub_nodes = nodes_gdf.set_index("node_index").loc[node_ids].reset_index() try: raw_window, all_dates = fetch_era5_land_daily_for_nodes(sub_nodes, target_date) x_tensor, target = build_model_input(raw_window) except Exception as e: return f"Error fetching ERA5 data or building input: {e}" x_tensor = x_tensor.to(device) sub_edge_index = sub_edge_index.to(device) with torch.no_grad(): pred = model(x_tensor, sub_edge_index).squeeze(0).cpu().numpy() # FIX: Clip negative predictions to 0.0 for physical consistency pred = np.maximum(pred, 0.0) center_pred = pred[center_local] forecast_dates = all_dates[HISTORY_DAYS:] lines = [f"Forecast for {glake_id} starting {target_date.isoformat()}:"] for step_idx, forecast_day in enumerate(forecast_dates): lines.append(f" Date: {forecast_day.isoformat()}") for var_idx, var_name in enumerate(TARGET_COLS): val = float(center_pred[step_idx, var_idx]) val_str = f"{val:.4f}" lines.append(f" - {var_name}: {val_str}") return "\n".join(lines) def run_inference_check() -> pd.DataFrame: random.seed(RANDOM_SEED) initialize_earth_engine() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model, edge_index = load_model_and_graph(device) nodes_gdf = load_graph_nodes() print(f"Model loaded from {WEIGHT_PATH}") print(f"Graph loaded from {EDGE_INDEX_PATH}: {edge_index.size(1):,} edges") print(f"Running on {device}") selected_nodes = choose_random_glake_nodes( nodes_gdf, AOI_POLYGON, count=len(TEST_DATES), edge_index=edge_index, ) rows = [] for date_text, selected_node in zip(TEST_DATES, selected_nodes.itertuples()): target_date = parse_date_ddmmyyyy(date_text) center_node = int(selected_node.node_index) glake_id = selected_node.GLAKE_ID lon = float(selected_node.geometry.x) lat = float(selected_node.geometry.y) node_ids, sub_edge_index, center_local = build_k_hop_subgraph(edge_index, center_node, hops=NODE_HOPS) sub_nodes = nodes_gdf.set_index("node_index").loc[node_ids].reset_index() print("\n" + "-" * 72) print(f"Requested date: {date_text} -> forecast start {target_date.isoformat()}") print(f"Random GLAKE_ID: {glake_id} | node_index={center_node}") print(f"Lake centroid: lon={lon:.5f}, lat={lat:.5f}") print(f"Subgraph: {len(node_ids)} nodes, {sub_edge_index.size(1)} edges, hops={NODE_HOPS}") raw_window, all_dates = fetch_era5_land_daily_for_nodes(sub_nodes, target_date) x_tensor, target = build_model_input(raw_window) x_tensor = x_tensor.to(device) sub_edge_index = sub_edge_index.to(device) with torch.no_grad(): pred = model(x_tensor, sub_edge_index).squeeze(0).cpu().numpy() # FIX: Clip negative predictions to 0.0 for physical consistency pred = np.maximum(pred, 0.0) center_pred = pred[center_local] center_target = target[:, center_local, :] metrics = masked_metrics(center_pred, center_target) forecast_dates = all_dates[HISTORY_DAYS:] print(f"Central-node MAE: {metrics['mae']:.6f}, RMSE: {metrics['rmse']:.6f}") for step_idx, forecast_day in enumerate(forecast_dates): for var_idx, var_name in enumerate(TARGET_COLS): rows.append( { "requested_date": date_text, "forecast_date": forecast_day.isoformat(), "glake_id": glake_id, "node_index": center_node, "node_lon": lon, "node_lat": lat, "subgraph_nodes": len(node_ids), "subgraph_edges": int(sub_edge_index.size(1)), "target_variable": var_name, "prediction": float(center_pred[step_idx, var_idx]), "actual": float(center_target[step_idx, var_idx]) if not np.isnan(center_target[step_idx, var_idx]) else np.nan, "absolute_error": float(abs(center_pred[step_idx, var_idx] - center_target[step_idx, var_idx])) if not np.isnan(center_target[step_idx, var_idx]) else np.nan, } ) case_df = pd.DataFrame(rows).tail(FORECAST_DAYS * len(TARGET_COLS)) print(case_df[["forecast_date", "target_variable", "prediction", "actual", "absolute_error"]]) results = pd.DataFrame(rows) output_path = SCRIPT_DIR / "inference_accuracy_check.csv" results.to_csv(output_path, index=False) print(f"\nSaved inference accuracy results to {output_path}") return results if __name__ == "__main__": run_inference_check()