Spaces:
Running on Zero
Running on Zero
| """ | |
| 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}") | |
| 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'] | |