intermediate backup

optimizer/forecasting/ensemble.py

@@ -0,0 +1,188 @@
+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
+
+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 EnsembleProvider(ForecastProvider):
+    """Provides forecasts using an ensemble of trained LSTM models."""
+
+    def __init__(
+        self,
+        fold_artifacts: List[Dict[str, Any]],
+        ensemble_method: str,
+        ensemble_feature_config: FeatureConfig, # Assumed consistent across folds by loading logic
+        ensemble_target_col: str, # Assumed consistent
+    ):
+        if not fold_artifacts:
+            raise ValueError("EnsembleProvider requires at least one fold artifact.")
+
+        self.fold_artifacts = fold_artifacts
+        self.ensemble_method = ensemble_method
+        # Store common config for reference, but use fold-specific details in get_forecast
+        self.ensemble_feature_config = ensemble_feature_config
+        self.ensemble_target_col = ensemble_target_col
+        self.common_forecast_horizons = sorted(ensemble_feature_config.forecast_horizon) # Assumed consistent
+
+        # Calculate max lookback needed across all folds
+        max_lookback = 0
+        for i, fold in enumerate(fold_artifacts):
+            try:
+                fold_feature_config = fold['feature_config']
+                fold_seq_len = fold_feature_config.sequence_length
+
+                feature_lookback = 0
+                if fold_feature_config.lags:
+                    feature_lookback = max(feature_lookback, max(fold_feature_config.lags))
+                if fold_feature_config.rolling_window_sizes:
+                     feature_lookback = max(feature_lookback, max(w - 1 for w in fold_feature_config.rolling_window_sizes))
+
+                fold_total_lookback = fold_seq_len + feature_lookback
+                max_lookback = max(max_lookback, fold_total_lookback)
+            except KeyError as e:
+                 raise ValueError(f"Fold artifact {i} is missing expected key: {e}") from e
+            except Exception as e:
+                 raise ValueError(f"Error processing fold artifact {i} for lookback calculation: {e}") from e
+
+        self._required_lookback = max_lookback
+        logger.debug(f"EnsembleProvider initialized with {len(fold_artifacts)} folds. Method: '{ensemble_method}'. Required lookback: {self._required_lookback}")
+
+        if ensemble_method not in ['mean', 'median']:
+             raise ValueError(f"Unsupported ensemble method: {ensemble_method}. Use 'mean' or 'median'.")
+
+    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 forecasts from each fold model, interpolates, and aggregates.
+        """
+        logger.debug(f"EnsembleProvider: Generating forecast for {optimization_horizon_hours} hours using {self.ensemble_method}.")
+        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
+
+        fold_forecasts_interpolated = []
+        last_actual_price = historical_data_slice[self.ensemble_target_col].iloc[-1] # Common anchor for all folds
+
+        for i, fold_artifact in enumerate(self.fold_artifacts):
+            fold_id = fold_artifact.get("fold_id", i + 1)
+            try:
+                fold_model: LSTMForecastLightningModule = fold_artifact['model_instance']
+                fold_feature_config: FeatureConfig = fold_artifact['feature_config']
+                fold_target_scaler: Optional[Any] = fold_artifact['target_scaler']
+                fold_target_col: str = fold_artifact['main_forecasting_config'].data.target_col # Use fold specific target
+                fold_seq_len = fold_feature_config.sequence_length
+                fold_horizons = sorted(fold_feature_config.forecast_horizon)
+
+                # Calculate lookback needed *for this specific fold* to check slice length
+                fold_feature_lookback = 0
+                if fold_feature_config.lags: fold_feature_lookback = max(fold_feature_lookback, max(fold_feature_config.lags))
+                if fold_feature_config.rolling_window_sizes: fold_feature_lookback = max(fold_feature_lookback, max(w - 1 for w in fold_feature_config.rolling_window_sizes))
+                fold_total_lookback = fold_seq_len + fold_feature_lookback
+
+                if len(historical_data_slice) < fold_total_lookback:
+                    logger.warning(f"Fold {fold_id}: Skipping fold. Insufficient historical data in slice for this fold's lookback ({fold_total_lookback} needed).")
+                    continue
+
+                # 1. Feature Engineering (using fold's config)
+                # Slice needs to be long enough for this fold's total lookback.
+                # The input slice `historical_data_slice` should already be long enough based on max_lookback.
+                engineered_df_fold = engineer_features(historical_data_slice.copy(), fold_target_col, fold_feature_config)
+
+                if engineered_df_fold.isnull().any().any():
+                    logger.warning(f"Fold {fold_id}: NaNs found after feature engineering. Attempting fill.")
+                    engineered_df_fold = engineered_df_fold.ffill().bfill()
+                    if engineered_df_fold.isnull().any().any():
+                        logger.error(f"Fold {fold_id}: NaNs persist after fill. Skipping fold.")
+                        continue
+
+                # 2. Create *one* input sequence (using fold's sequence length)
+                if len(engineered_df_fold) < fold_seq_len:
+                    logger.error(f"Fold {fold_id}: Engineered data ({len(engineered_df_fold)}) is shorter than fold sequence length ({fold_seq_len}). Skipping fold.")
+                    continue
+
+                input_sequence_data_fold = engineered_df_fold.iloc[-fold_seq_len:].copy()
+                feature_columns_fold = [col for col in engineered_df_fold.columns if col != fold_target_col] # Example
+                if not feature_columns_fold: feature_columns_fold = engineered_df_fold.columns.tolist()
+                input_sequence_np_fold = input_sequence_data_fold[feature_columns_fold].values
+
+                if input_sequence_np_fold.shape != (fold_seq_len, len(feature_columns_fold)):
+                     logger.error(f"Fold {fold_id}: Input sequence has wrong shape. Expected ({fold_seq_len}, {len(feature_columns_fold)}), got {input_sequence_np_fold.shape}. Skipping fold.")
+                     continue
+
+                input_tensor_fold = torch.FloatTensor(input_sequence_np_fold).unsqueeze(0)
+
+                # 3. Run Inference (using fold's model)
+                fold_model.eval()
+                with torch.no_grad():
+                    predictions_scaled_fold = fold_model(input_tensor_fold) # Shape (1, num_fold_horizons)
+
+                if predictions_scaled_fold.ndim != 2 or predictions_scaled_fold.shape[0] != 1 or predictions_scaled_fold.shape[1] != len(fold_horizons):
+                     logger.error(f"Fold {fold_id}: Prediction output shape mismatch. Expected (1, {len(fold_horizons)}), got {predictions_scaled_fold.shape}. Skipping fold.")
+                     continue
+
+                predictions_scaled_np_fold = predictions_scaled_fold.squeeze(0).cpu().numpy()
+
+                # 4. Inverse Transform (using fold's scaler)
+                predictions_original_scale_fold = predictions_scaled_np_fold
+                if fold_target_scaler:
+                    try:
+                        predictions_original_scale_fold = fold_target_scaler.inverse_transform(predictions_scaled_np_fold.reshape(-1, 1)).flatten()
+                    except Exception as e:
+                        logger.error(f"Fold {fold_id}: Failed to apply inverse transform: {e}. Skipping fold.", exc_info=True)
+                        continue
+
+                # 5. Interpolate (using fold's horizons)
+                interpolated_forecast_fold = interpolate_forecast(
+                    native_horizons=fold_horizons,
+                    native_predictions=predictions_original_scale_fold,
+                    target_horizon=optimization_horizon_hours,
+                    last_known_actual=last_actual_price
+                )
+
+                if interpolated_forecast_fold is not None:
+                     fold_forecasts_interpolated.append(interpolated_forecast_fold)
+                     logger.debug(f"Fold {fold_id}: Successfully generated interpolated forecast.")
+                else:
+                     logger.warning(f"Fold {fold_id}: Interpolation failed. Skipping fold.")
+
+            except Exception as e:
+                logger.error(f"Error processing ensemble fold {fold_id}: {e}", exc_info=True)
+                continue # Skip this fold on error
+
+        # --- Aggregation ---
+        if not fold_forecasts_interpolated:
+            logger.error("No successful forecasts generated from any ensemble folds.")
+            return None
+
+        logger.debug(f"Aggregating forecasts from {len(fold_forecasts_interpolated)} folds using '{self.ensemble_method}'.")
+        stacked_predictions = np.stack(fold_forecasts_interpolated, axis=0) # Shape (n_folds, target_horizon)
+
+        if self.ensemble_method == 'mean':
+            final_ensemble_forecast = np.mean(stacked_predictions, axis=0)
+        elif self.ensemble_method == 'median':
+            final_ensemble_forecast = np.median(stacked_predictions, axis=0)
+        else:
+            # Should be caught in __init__, but double-check
+            logger.error(f"Internal error: Invalid ensemble method '{self.ensemble_method}' during aggregation.")
+            return None
+
+        logger.debug(f"EnsembleProvider: Successfully generated forecast.")
+        return final_ensemble_forecast