Files
entrix_case_challange/optimizer/forecasting/single_model.py
2025-05-03 20:46:14 +02:00

151 lines
7.5 KiB
Python

import logging
from typing import List, Dict, Any, Optional
import numpy as np
import pandas as pd
import torch
from sklearn.preprocessing import StandardScaler, MinMaxScaler
# Imports from our project structure
from .base import ForecastProvider
from forecasting_model.utils import FeatureConfig
from forecasting_model.train.model import LSTMForecastLightningModule
from forecasting_model import engineer_features
from optimizer.forecasting.utils import interpolate_forecast
logger = logging.getLogger(__name__)
class SingleModelProvider(ForecastProvider):
"""Provides forecasts using a single trained LSTM model."""
def __init__(
self,
model_instance: LSTMForecastLightningModule,
feature_config: FeatureConfig,
target_col: str,
target_scaler: Optional[Any], # BaseEstimator, TransformerMixin -> more specific if possible
# input_size: int # Not needed directly if model instance is configured
):
self.model = model_instance
self.feature_config = feature_config
self.target_col = target_col
self.target_scaler = target_scaler
self.sequence_length = feature_config.sequence_length
self.forecast_horizons = sorted(feature_config.forecast_horizon) # Ensure sorted
# Calculate required lookback for feature engineering
feature_lookback = 0
if feature_config.lags:
feature_lookback = max(feature_lookback, max(feature_config.lags))
if feature_config.rolling_window_sizes:
# Rolling window of size W needs W-1 previous points
feature_lookback = max(feature_lookback, max(w - 1 for w in feature_config.rolling_window_sizes))
# Total lookback: sequence length for model input + feature engineering needs
# We need `sequence_length` points for the *last* input sequence.
# The first point of that sequence needs `feature_lookback` points before it.
# So, total points needed before the *end* of the input sequence is sequence_length + feature_lookback.
# Since the input sequence ends *before* the first forecast point (t=1),
# we need `sequence_length + feature_lookback` points before t=1.
self._required_lookback = self.sequence_length + feature_lookback
logger.debug(f"SingleModelProvider initialized. Required lookback: {self._required_lookback} (SeqLen: {self.sequence_length}, FeatLookback: {feature_lookback})")
def get_required_lookback(self) -> int:
return self._required_lookback
def get_forecast(
self,
historical_data_slice: pd.DataFrame,
optimization_horizon_hours: int
) -> np.ndarray | None:
"""
Generates forecast using the single model and interpolates to hourly resolution.
"""
logger.debug(f"SingleModelProvider: Generating forecast for {optimization_horizon_hours} hours.")
if len(historical_data_slice) < self._required_lookback:
logger.error(f"Insufficient historical data provided. Need {self._required_lookback}, got {len(historical_data_slice)}.")
return None
try:
# 1. Feature Engineering
# Use the provided slice which already includes the lookback.
engineered_df = engineer_features(historical_data_slice.copy(), self.target_col, self.feature_config)
# Check for NaNs after feature engineering before creating sequences
if engineered_df.isnull().any().any():
logger.warning("NaNs found after feature engineering. Attempting to fill with ffill/bfill.")
# Be careful about filling target vs features if needed
engineered_df = engineered_df.ffill().bfill()
if engineered_df.isnull().any().any():
logger.error("NaNs persist after fill. Cannot create sequences.")
return None
# 2. Create *one* input sequence ending at the last point of the historical slice
# This sequence is used to predict starting from the next hour (t=1)
if len(engineered_df) < self.sequence_length:
logger.error(f"Engineered data ({len(engineered_df)}) is shorter than sequence length ({self.sequence_length}).")
return None
input_sequence_data = engineered_df.iloc[-self.sequence_length:].copy()
# Convert sequence data to numpy array (excluding target if model expects it that way)
# Assuming model takes all engineered features as input
# TODO: Verify the exact features the model expects (target included/excluded?)
# Assuming all columns except maybe the original target are features
feature_columns = [col for col in engineered_df.columns if col != self.target_col] # Example
if not feature_columns: feature_columns = engineered_df.columns.tolist() # Use all if target wasn't dropped
input_sequence_np = input_sequence_data[feature_columns].values
if input_sequence_np.shape != (self.sequence_length, len(feature_columns)):
logger.error(f"Input sequence has wrong shape. Expected ({self.sequence_length}, {len(feature_columns)}), got {input_sequence_np.shape}")
return None
input_tensor = torch.FloatTensor(input_sequence_np).unsqueeze(0) # Add batch dim
# 3. Run Inference
self.model.eval()
with torch.no_grad():
# Model output shape: (1, num_horizons)
predictions_scaled = self.model(input_tensor)
if predictions_scaled.ndim != 2 or predictions_scaled.shape[0] != 1 or predictions_scaled.shape[1] != len(self.forecast_horizons):
logger.error(f"Model prediction output shape mismatch. Expected (1, {len(self.forecast_horizons)}), got {predictions_scaled.shape}.")
return None
predictions_scaled_np = predictions_scaled.squeeze(0).cpu().numpy() # Shape: (num_horizons,)
# 4. Inverse Transform
predictions_original_scale = predictions_scaled_np
if self.target_scaler:
try:
# Scaler expects shape (n_samples, n_features), even if n_features=1
predictions_original_scale = self.target_scaler.inverse_transform(predictions_scaled_np.reshape(-1, 1)).flatten()
logger.debug("Applied inverse transform to predictions.")
except Exception as e:
logger.error(f"Failed to apply inverse transform: {e}", exc_info=True)
# Decide whether to return scaled or None. Returning None is safer.
return None
# 5. Interpolate
# Use the last actual price from the input data as the anchor point t=0
last_actual_price = historical_data_slice[self.target_col].iloc[-1]
interpolated_forecast = interpolate_forecast(
native_horizons=self.forecast_horizons,
native_predictions=predictions_original_scale,
target_horizon=optimization_horizon_hours,
last_known_actual=last_actual_price
)
if interpolated_forecast is None:
logger.error("Interpolation step failed.")
return None
logger.debug(f"SingleModelProvider: Successfully generated forecast.")
return interpolated_forecast
except Exception as e:
logger.error(f"Error during single model forecast generation: {e}", exc_info=True)
return None