Spaces:
Running on Zero
Running on Zero
File size: 8,782 Bytes
a74054f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 | """
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']
|