Files
entrix_case_challange/forecasting_model/utils/config_model.py
2025-05-02 14:36:19 +02:00

151 lines
7.5 KiB
Python

from pydantic import BaseModel, Field, field_validator, model_validator
from typing import Optional, List, Union, Literal
from enum import Enum
# --- Nested Configs ---
class WaveletTransformConfig(BaseModel):
"""Configuration for optional wavelet transform features."""
apply: bool = False
target_or_feature: Literal['target', 'feature'] = "target"
wavelet_type: str = "db4"
level: int = Field(3, gt=0) # Level must be positive
use_coeffs: List[str] = ["approx", "detail_1"]
class ClippingConfig(BaseModel):
"""Configuration for optional feature clipping."""
apply: bool = False
clip_min: float = -5.0
clip_max: float = 5.0
@model_validator(mode='after')
def check_clip_range(self) -> 'ClippingConfig':
if self.apply and self.clip_max <= self.clip_min:
raise ValueError(f'clip_max ({self.clip_max}) must be greater than clip_min ({self.clip_min}) when clipping is applied.')
return self
# --- Main Config Sections ---
class DataConfig(BaseModel):
"""Configuration related to data loading and initial preparation."""
data_path: str = Field(..., description="Path to the input CSV data file.")
# --- Raw Data Specifics ---
raw_datetime_col: str = Field(..., description="Name of the raw datetime column in the CSV (e.g., 'MTU (CET/CEST)')")
raw_target_col: str = Field(..., description="Name of the raw target/price column in the CSV (e.g., 'Day-ahead Price [EUR/MWh]')")
raw_datetime_format: str = '%d.%m.%Y %H:%M' # Example, make it configurable if needed
# --- Standardized Names & Processing ---
datetime_col: str = Field(..., description="Standardized name for the datetime index after processing (e.g., 'Timestamp')")
target_col: str = Field(..., description="Standardized name for the target column after processing (e.g., 'Price')")
expected_frequency: Optional[str] = Field('h', description="Expected pandas frequency string (e.g., 'h', 'D', '15min'). If null, no frequency check/setting is performed.")
fill_initial_target_nans: bool = Field(True, description="Forward/backward fill NaNs in the target column immediately after loading?")
class FeatureConfig(BaseModel):
"""Configuration for feature engineering and preprocessing."""
sequence_length: int = Field(..., gt=0)
forecast_horizon: int = Field(..., gt=0)
lags: List[int] = []
rolling_window_sizes: List[int] = []
use_time_features: bool = True
sinus_curve: bool = False # Added
cosin_curve: bool = False # Added
wavelet_transform: Optional[WaveletTransformConfig] = None
fill_nan: Optional[Union[str, float, int]] = 'ffill' # Added (e.g., 'ffill', 0)
clipping: ClippingConfig = ClippingConfig() # Default instance
scaling_method: Optional[Literal['standard', 'minmax']] = 'standard' # Added literal validation
@field_validator('lags', 'rolling_window_sizes')
@classmethod
def check_positive_list_values(cls, v: List[int]) -> List[int]:
if any(val <= 0 for val in v):
raise ValueError('Lists lags/rolling_window_sizes must contain only positive values')
return v
class ModelConfig(BaseModel):
"""Configuration for the forecasting model architecture."""
# input_size: Optional[int] = Field(None, gt=0) # Removed: Determined dynamically
hidden_size: int = Field(..., gt=0)
num_layers: int = Field(..., gt=0)
dropout: float = Field(..., ge=0.0, le=1.0)
use_residual_skips: bool = False
# Add forecast_horizon here to ensure LightningModule gets it directly
forecast_horizon: Optional[int] = Field(None, gt=0) # Will be set from FeatureConfig
class TrainingConfig(BaseModel):
"""Configuration for the training process (PyTorch Lightning)."""
batch_size: int = Field(..., gt=0)
epochs: int = Field(..., gt=0) # Max epochs
learning_rate: float = Field(..., gt=0.0)
loss_function: Literal['MSE', 'MAE'] = 'MSE'
# device: str = 'auto' # Handled by PL Trainer accelerator/devices args
early_stopping_patience: Optional[int] = Field(None, ge=1) # Patience must be >= 1 if set
scheduler_step_size: Optional[int] = Field(None, gt=0)
scheduler_gamma: Optional[float] = Field(None, gt=0.0, lt=1.0)
gradient_clip_val: Optional[float] = Field(None, ge=0.0) # Added
num_workers: int = Field(0, ge=0) # Added
precision: Literal[16, 32, 64, 'bf16'] = 32 # Added
class CrossValidationConfig(BaseModel):
"""Configuration for time series cross-validation."""
n_splits: int = Field(..., gt=0)
test_size_fraction: float = Field(..., gt=0.0, lt=1.0, description="Fraction of the fixed training window size for the test set.")
val_size_fraction: float = Field(..., gt=0.0, lt=1.0, description="Fraction of the fixed training window size for the validation set.")
initial_train_size: Optional[Union[int, float]] = Field(None, gt=0.0, description="Size of the fixed training window (absolute number or fraction of total data > 0). If null, estimated automatically.")
class EvaluationConfig(BaseModel):
"""Configuration for the final evaluation process."""
# metrics: List[str] = ['MAE', 'RMSE'] # Defined internally now
eval_batch_size: int = Field(..., gt=0)
save_plots: bool = True
plot_sample_size: Optional[int] = Field(1000, gt=0) # Max points for plots
class OptunaConfig(BaseModel):
"""Optional configuration for Optuna hyperparameter optimization."""
enabled: bool = False
n_trials: int = Field(20, gt=0)
storage: Optional[str] = None # e.g., "sqlite:///output/hpo_results/study.db"
direction: Literal['minimize', 'maximize'] = 'minimize'
metric_to_optimize: str = 'val_mae_orig_scale'
pruning: bool = True
# --- Top-Level Configuration Model ---
class MainConfig(BaseModel):
"""Main configuration model nesting all sections."""
project_name: str = "TimeSeriesForecasting"
random_seed: Optional[int] = 42 # Added top-level seed
data: DataConfig
features: FeatureConfig
model: ModelConfig # ModelConfig no longer contains input_size
training: TrainingConfig
cross_validation: CrossValidationConfig
evaluation: EvaluationConfig
optuna: Optional[OptunaConfig] = OptunaConfig() # Added optional Optuna config
@model_validator(mode='after')
def check_forecast_horizon_consistency(self) -> 'MainConfig':
# Ensure model config gets forecast horizon from features config if not set
if self.features and self.model:
if self.model.forecast_horizon is None:
# If model config doesn't have it, set it from features config
self.model.forecast_horizon = self.features.forecast_horizon
elif self.model.forecast_horizon != self.features.forecast_horizon:
# If both are set but differ, raise error
raise ValueError(
f"ModelConfig forecast_horizon ({self.model.forecast_horizon}) must match "
f"FeatureConfig forecast_horizon ({self.features.forecast_horizon})."
)
# After potential setting, ensure model.forecast_horizon is actually set
if self.model and (self.model.forecast_horizon is None or self.model.forecast_horizon <= 0):
raise ValueError("ModelConfig requires a positive forecast_horizon (must be set in features config if not set explicitly in model config).")
# Input size check is removed as it's not part of static config anymore
return self
class Config:
# Example configuration for Pydantic itself
validate_assignment = True # Re-validate on assignment
# extra = 'forbid' # Forbid extra fields not defined in schema