""" Tensor builder for spatio-temporal GNN training. Single Responsibility: Convert time series data to tensor format with sliding windows. """ import pandas as pd import numpy as np import torch from typing import Tuple, List, Optional, Dict from ..config.settings import TARGET_VARIABLES, SAFRAN_VARIABLES, TEMPORAL_WINDOW class SpatioTemporalTensorBuilder: """ Builds tensors for spatio-temporal GNN training using sliding windows. Output format: - X: [samples, input_window, stations, features] - Input features - Y: [samples, stations, targets, horizons] - Target variables Uses sliding window approach: - Input window: Previous T days (e.g., 30 days) - Forecasting horizons: Next 1, 3, 7, 14 days """ def __init__( self, input_window: int = TEMPORAL_WINDOW, forecast_horizons: List[int] = [1, 3, 7, 14], target_vars: List[str] = TARGET_VARIABLES, feature_vars: List[str] = SAFRAN_VARIABLES ): """ Initialize tensor builder. Args: input_window: Number of past days for input (default: 30) forecast_horizons: Days ahead to forecast (default: [1, 3, 7, 14]) target_vars: Target variables to predict (discharge, water level) feature_vars: Input feature variables (meteorological) """ self.input_window = input_window self.forecast_horizons = forecast_horizons self.target_vars = target_vars self.feature_vars = feature_vars def build_tensors( self, fused_df: pd.DataFrame, station_ids: Optional[List[str]] = None, date_col: str = "date", station_col: str = "station_id" ) -> Tuple[torch.Tensor, torch.Tensor]: """ Build input and target tensors from fused data. Args: fused_df: Fused dataframe with time series for all stations station_ids: List of station IDs (if None, use all) date_col: Name of date column station_col: Name of station column Returns: Tuple of (X, Y) tensors: - X: [samples, input_window, stations, features] - Y: [samples, stations, targets, horizons] """ # Get unique stations if station_ids is None: station_ids = sorted(fused_df[station_col].unique()) num_stations = len(station_ids) # Sort by date fused_df = fused_df.sort_values([station_col, date_col]).reset_index(drop=True) # Create station-indexed data station_data = {} for station in station_ids: station_df = fused_df[fused_df[station_col] == station].copy() station_df = station_df.sort_values(date_col).reset_index(drop=True) station_data[station] = station_df # Build samples using sliding window X_samples = [] Y_samples = [] # Get date range for sliding window all_dates = sorted(fused_df[date_col].unique()) for i in range(len(all_dates) - self.input_window - max(self.forecast_horizons)): # Get input window dates input_start_idx = i input_end_idx = i + self.input_window # Build input tensor for this window X_window = self._build_input_window( station_data, station_ids, all_dates[input_start_idx:input_end_idx], date_col ) # Build target tensor for forecast horizons Y_targets = self._build_targets( station_data, station_ids, all_dates, input_end_idx, date_col ) if X_window is not None and Y_targets is not None: X_samples.append(X_window) Y_samples.append(Y_targets) # Convert to tensors if len(X_samples) == 0: raise ValueError("No valid samples could be created. Check data quality and temporal coverage.") X = torch.tensor(np.array(X_samples), dtype=torch.float32) Y = torch.tensor(np.array(Y_samples), dtype=torch.float32) print(f"Built tensors: X shape={X.shape}, Y shape={Y.shape}") return X, Y def _build_input_window( self, station_data: Dict[str, pd.DataFrame], station_ids: List[str], window_dates: List, date_col: str ) -> Optional[np.ndarray]: """ Build input tensor for one sliding window. Args: station_data: Dictionary of station DataFrames station_ids: List of station IDs window_dates: Dates in this window date_col: Date column name Returns: Array of shape [input_window, stations, features] """ window_data = [] for date in window_dates: station_features = [] for station_id in station_ids: df = station_data[station_id] row = df[df[date_col] == date] if len(row) == 0: # Missing data - use zeros or skip features = [0.0] * len(self.feature_vars) else: features = [] for var in self.feature_vars: if var in row.columns: features.append(float(row[var].iloc[0])) else: features.append(0.0) station_features.append(features) window_data.append(station_features) return np.array(window_data) # [time_steps, stations, features] def _build_targets( self, station_data: Dict[str, pd.DataFrame], station_ids: List[str], all_dates: List, current_idx: int, date_col: str ) -> Optional[np.ndarray]: """ Build target tensor for forecast horizons. Args: station_data: Dictionary of station DataFrames station_ids: List of station IDs all_dates: All dates current_idx: Current position in date sequence date_col: Date column name Returns: Array of shape [stations, targets, horizons] """ targets = [] for station_id in station_ids: df = station_data[station_id] station_targets = [] for target_var in self.target_vars: horizon_values = [] for horizon in self.forecast_horizons: target_date_idx = current_idx + horizon if target_date_idx >= len(all_dates): horizon_values.append(0.0) continue target_date = all_dates[target_date_idx] row = df[df[date_col] == target_date] if len(row) == 0 or target_var not in row.columns: horizon_values.append(0.0) else: horizon_values.append(float(row[target_var].iloc[0])) station_targets.append(horizon_values) targets.append(station_targets) return np.array(targets) # [stations, targets, horizons] def get_tensor_shapes(self, X: torch.Tensor, Y: torch.Tensor) -> Dict: """ Get tensor shape information. Args: X: Input tensor Y: Target tensor Returns: Dictionary with shape information """ return { "X_shape": list(X.shape), "Y_shape": list(Y.shape), "num_samples": X.shape[0], "input_window": X.shape[1], "num_stations": X.shape[2], "num_features": X.shape[3], "num_targets": Y.shape[2], "num_horizons": Y.shape[3] } def save_tensors(self, X: torch.Tensor, Y: torch.Tensor, filepath: str) -> None: """ Save tensors to file. Args: X: Input tensor Y: Target tensor filepath: Path to save file """ torch.save({'X': X, 'Y': Y}, filepath) print(f"Tensors saved to {filepath}") @staticmethod def load_tensors(filepath: str) -> Tuple[torch.Tensor, torch.Tensor]: """ Load tensors from file. Args: filepath: Path to tensor file Returns: Tuple of (X, Y) tensors """ data = torch.load(filepath) return data['X'], data['Y']