SIENTIAPDE-1222
SIENTIAPDE-1214: Enhance MLFlow and tests with datetime index handling and logging improvements - Added a new method in MLFlow to detect and parse datetime indices in DataFrames, ensuring proper format and raising errors for invalid types. - Updated prediction workflows to utilize the new datetime index handling, improving data integrity during transformations. - Enhanced logging in model_repository to include detailed data outputs for better traceability. - Adjusted timeout settings in prediction workflows for improved execution time management. - Updated tests.ipynb to include additional checks for index types and outputs for better validation of functionality.
This commit is contained in:
vitor-aignosi
Binary file not shown.
@@ -1,824 +0,0 @@
|
||||
"""
|
||||
Base model classes and interfaces.
|
||||
|
||||
This module defines base classes with consistent interfaces for all models,
|
||||
promoting modular model development. It includes essential functionality
|
||||
for model fitting, prediction, evaluation, saving, and loading, while
|
||||
abstracting common behaviors into base classes.
|
||||
|
||||
While full compatibility with scikit-learn is not guaranteed, the base
|
||||
classes provide a consistent interface for model fitting, prediction,
|
||||
evaluation, saving, and loading, which should be sufficient for most
|
||||
use cases.
|
||||
|
||||
Key components:
|
||||
- **Model**: Abstract base class for all models, providing core utilities and
|
||||
interfaces.
|
||||
- **TimeSeriesModel**: Abstract base class for time series models, adding
|
||||
time-based functionality.
|
||||
- **UnivariateTimeSeriesModel**: Base class for univariate time series models.
|
||||
- **MultivariateTimeSeriesModel**: Base class for multivariate time series
|
||||
models that use exogenous features.
|
||||
|
||||
The module also includes utility functions such as `ensure_fitted`, which
|
||||
ensures models are fitted before calling certain methods.
|
||||
|
||||
Modules in this package should inherit from these base classes and implement
|
||||
the required methods.
|
||||
|
||||
TODO:
|
||||
- Add methods to create lagged features for time series models.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
import joblib
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Optional, List, Sequence, Tuple, cast, Any, Protocol
|
||||
from contextlib import contextmanager
|
||||
|
||||
import pandas as pd
|
||||
import numpy as np
|
||||
import shap
|
||||
import matplotlib.pyplot as plt
|
||||
from rich.console import Console
|
||||
|
||||
from sklearn.base import BaseEstimator, RegressorMixin
|
||||
from sklearn.utils.validation import check_array, check_X_y
|
||||
from sklearn.exceptions import NotFittedError
|
||||
|
||||
console = Console()
|
||||
|
||||
|
||||
class PredictorProtocol(Protocol):
|
||||
"""Protocol for models with predict method and optional imputation."""
|
||||
|
||||
def predict(self, X: Any) -> Any: ...
|
||||
def _impute_missing_values(self, X: Any) -> Any: ...
|
||||
|
||||
|
||||
def ensure_fitted(method):
|
||||
"""
|
||||
Decorator to ensure the model is fitted before calling the method.
|
||||
|
||||
Raises:
|
||||
sklearn.exceptions.NotFittedError: If the model is not fitted.
|
||||
Usage:
|
||||
@ensure_fitted
|
||||
def predict(self, X): # Or other methods requiring fit
|
||||
pass
|
||||
"""
|
||||
|
||||
def wrapper(self, *args, **kwargs):
|
||||
is_fitted = self.__sklearn_is_fitted__()
|
||||
if not is_fitted:
|
||||
raise NotFittedError(
|
||||
f"This {self.__class__.__name__} instance is not fitted yet. "
|
||||
"Call 'fit' with appropriate arguments before using this "
|
||||
"method."
|
||||
)
|
||||
return method(self, *args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
class Model(BaseEstimator, ABC):
|
||||
"""
|
||||
Abstract base class for all models.
|
||||
|
||||
Provides core utilities, input validation, and interface consistency
|
||||
for time series models. Compatible with scikit-learn workflows.
|
||||
"""
|
||||
|
||||
def __init__(self, name: Optional[str] = None, random_seed: int = 42):
|
||||
"""
|
||||
Initialize the model with a unique name and random seed.
|
||||
|
||||
Args:
|
||||
name: Optional identifier; auto-generated if None.
|
||||
random_seed: Seed for reproducibility.
|
||||
"""
|
||||
self.name = name or f"{self.__class__.__name__}_{uuid.uuid4().hex}"
|
||||
self.random_seed = random_seed
|
||||
self.feature_names_in_: Optional[List[str]] = None
|
||||
self.n_features_in_: Optional[int] = None
|
||||
self._is_fitted = False
|
||||
|
||||
def fit(
|
||||
self,
|
||||
y: pd.Series,
|
||||
X: Optional[pd.DataFrame] = None,
|
||||
X_val: Optional[pd.DataFrame] = None,
|
||||
y_val: Optional[pd.Series] = None,
|
||||
) -> "Model":
|
||||
"""
|
||||
Trains the model.
|
||||
|
||||
Handles basic input validation for y and sets internal fitted
|
||||
state after calling _fit_logic.
|
||||
|
||||
'X' is optional to account for univariate time series models.
|
||||
|
||||
Args:
|
||||
y: The target variable.
|
||||
X: Optional exogenous variables.
|
||||
X_val: Optional validation feature matrix.
|
||||
y_val: Optional validation target series.
|
||||
Raises:
|
||||
TypeError: If y is not a pandas Series.
|
||||
If X is provided, it must be a pandas DataFrame.
|
||||
If X_val and y_val are provided, they must be pandas DataFrames
|
||||
and Series respectively.
|
||||
|
||||
Returns:
|
||||
Self for chaining.
|
||||
"""
|
||||
if not isinstance(y, pd.Series):
|
||||
raise TypeError("Input 'y' (target) must be a pandas Series.")
|
||||
|
||||
self._fit_logic(y, X, X_val, y_val)
|
||||
self._is_fitted = True
|
||||
return self
|
||||
|
||||
@abstractmethod
|
||||
def _fit_logic(
|
||||
self,
|
||||
y: pd.Series,
|
||||
X: Optional[pd.DataFrame] = None,
|
||||
X_val: Optional[pd.DataFrame] = None,
|
||||
y_val: Optional[pd.Series] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Core fitting logic to be implemented by subclasses with
|
||||
optional validation data.
|
||||
|
||||
Args:
|
||||
y: The target variable.
|
||||
X: Optional exogenous variables.
|
||||
X_val: Validation feature matrix (optional).
|
||||
y_val: Validation target series (optional).
|
||||
"""
|
||||
raise NotImplementedError("Subclasses must implement _fit_logic().")
|
||||
|
||||
@ensure_fitted
|
||||
@abstractmethod
|
||||
def predict(self, X: Optional[pd.DataFrame] = None) -> Sequence:
|
||||
"""
|
||||
Predict values.
|
||||
|
||||
Args:
|
||||
X: Optional features for prediction. For univariate models
|
||||
not using exogenous variables, this might be None or
|
||||
contain future timestamps. Multivariate models will
|
||||
require X.
|
||||
|
||||
Returns:
|
||||
NumPy array or similar sequence of predictions.
|
||||
"""
|
||||
raise NotImplementedError("Subclasses must implement predict().")
|
||||
|
||||
def fit_predict(
|
||||
self,
|
||||
y: pd.Series,
|
||||
X: Optional[pd.DataFrame] = None,
|
||||
X_val: Optional[pd.DataFrame] = None,
|
||||
y_val: Optional[pd.Series] = None,
|
||||
) -> Sequence:
|
||||
"""
|
||||
Fits model and returns predictions on the same data.
|
||||
|
||||
Args:
|
||||
y: The target time series.
|
||||
X: Optional exogenous variables.
|
||||
|
||||
Returns:
|
||||
Predictions for the input data.
|
||||
"""
|
||||
return self.fit(y, X, X_val, y_val).predict(X)
|
||||
|
||||
def save(self, path: str) -> None:
|
||||
"""
|
||||
Saves model to disk using joblib.
|
||||
"""
|
||||
joblib.dump(self, path)
|
||||
|
||||
@classmethod
|
||||
def load(cls, path: str) -> "Model":
|
||||
"""
|
||||
Loads model from disk using joblib.
|
||||
"""
|
||||
return joblib.load(path)
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"{self.__class__.__name__}(name={self.name})"
|
||||
|
||||
def __sklearn_is_fitted__(self) -> bool:
|
||||
"""
|
||||
Check fitted status and return a Boolean value.
|
||||
"""
|
||||
return hasattr(self, "_is_fitted") and self._is_fitted
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"{self.__class__.__name__}(name={self.name})"
|
||||
|
||||
@contextmanager
|
||||
def model_state_preservation(self):
|
||||
"""Context manager to preserve model state during operations."""
|
||||
original_state = self._get_state_snapshot()
|
||||
try:
|
||||
yield
|
||||
except Exception:
|
||||
self._restore_state_snapshot(original_state)
|
||||
raise
|
||||
|
||||
def _get_state_snapshot(self) -> dict:
|
||||
"""Get snapshot of current model state."""
|
||||
return {
|
||||
"name": self.name,
|
||||
"is_fitted": getattr(self, "_is_fitted", False),
|
||||
"feature_names": self.feature_names_in_,
|
||||
"n_features": self.n_features_in_,
|
||||
}
|
||||
|
||||
def _restore_state_snapshot(self, snapshot: dict) -> None:
|
||||
"""Restore model state from snapshot."""
|
||||
self.name = snapshot["name"]
|
||||
self._is_fitted = snapshot["is_fitted"]
|
||||
self.feature_names_in_ = snapshot["feature_names"]
|
||||
self.n_features_in_ = snapshot["n_features"]
|
||||
|
||||
def get_params_dict(self) -> dict:
|
||||
"""Get model parameters as dictionary for logging/serialization."""
|
||||
return {
|
||||
"name": self.name,
|
||||
"random_seed": self.random_seed,
|
||||
"n_features_in_": self.n_features_in_,
|
||||
}
|
||||
|
||||
def summary(self) -> str:
|
||||
"""Generate a summary string of the model."""
|
||||
params = self.get_params_dict()
|
||||
fitted_status = (
|
||||
"✓ Fitted" if self.__sklearn_is_fitted__() else "✗ Not fitted"
|
||||
)
|
||||
|
||||
summary_lines = [
|
||||
f"Model: {self.__class__.__name__}",
|
||||
f"Status: {fitted_status}",
|
||||
f"Features: {params.get('n_features_in_', 'Unknown')}",
|
||||
]
|
||||
|
||||
return "\n".join(summary_lines)
|
||||
|
||||
|
||||
class TimeSeriesModel(Model):
|
||||
"""
|
||||
Abstract base class for time series forecasting models.
|
||||
|
||||
Extends the base Model class with specific methods for time series
|
||||
data handling and evaluation.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: Optional[str] = None,
|
||||
time_col: str = "ds",
|
||||
target_col: str = "y",
|
||||
random_seed: int = 42,
|
||||
n_lags: int = 0,
|
||||
sampling_freq: Optional[str] = None,
|
||||
):
|
||||
super().__init__(name=name, random_seed=random_seed)
|
||||
self.time_col = time_col
|
||||
self.target_col = target_col
|
||||
self.n_lags = n_lags
|
||||
self.sampling_freq = sampling_freq
|
||||
|
||||
self.training_series_: Optional[pd.Series] = None
|
||||
self.model_: Optional[BaseEstimator] = None
|
||||
|
||||
# Validate configuration
|
||||
self._validate_configuration()
|
||||
|
||||
def _validate_configuration(self) -> None:
|
||||
"""Validate model configuration."""
|
||||
if self.n_lags < 0:
|
||||
raise ValueError("n_lags must be non-negative")
|
||||
|
||||
def _validate_y(self, y: pd.Series) -> np.ndarray:
|
||||
"""
|
||||
Validates the target variable (y) for the model.
|
||||
|
||||
Ensures y is a pandas Series and checks its name against
|
||||
the expected target column name. Converts y to a NumPy array.
|
||||
The series name can be None, but if it is set, it should match
|
||||
the expected target column name.
|
||||
|
||||
Args:
|
||||
y: The target variable as a pandas Series.
|
||||
|
||||
Returns:
|
||||
A NumPy array of the target variable.
|
||||
|
||||
Raises:
|
||||
TypeError: If y is not a pandas Series.
|
||||
"""
|
||||
# Check if y is a pandas Series
|
||||
if not isinstance(y, pd.Series):
|
||||
raise TypeError("Input 'y' (target) must be a pandas Series.")
|
||||
if (y.name is not None) and (y.name != self.target_col):
|
||||
raise ValueError(
|
||||
f"Expected target column name '{self.target_col}', "
|
||||
f"but got '{y.name}'."
|
||||
)
|
||||
return check_array(y, ensure_2d=False)
|
||||
|
||||
@ensure_fitted
|
||||
@abstractmethod
|
||||
def backtest(
|
||||
self,
|
||||
y: pd.Series,
|
||||
X: Optional[pd.DataFrame] = None,
|
||||
retrain_every: int = 50,
|
||||
reuse_previous_execution: bool = False,
|
||||
) -> pd.Series:
|
||||
"""
|
||||
Performs backtesting on the time series data.
|
||||
|
||||
Args:
|
||||
y: The target time series data.
|
||||
X: Optional exogenous features.
|
||||
retrain_every: Number of steps after which to retrain the model.
|
||||
reuse_previous_execution: Whether to reuse the previous execution
|
||||
of a backtest. If True, any overlapping data between the
|
||||
previous execution and the current execution will be used
|
||||
without retraining the model.
|
||||
Returns:
|
||||
Series of predictions for each step in the time series.
|
||||
"""
|
||||
|
||||
raise NotImplementedError("Subclasses must implement backtest().")
|
||||
|
||||
def get_params_dict(self) -> dict:
|
||||
"""Get model parameters as dictionary for logging/serialization."""
|
||||
base_params = super().get_params_dict()
|
||||
ts_params = {
|
||||
"time_col": self.time_col,
|
||||
"target_col": self.target_col,
|
||||
"n_lags": self.n_lags,
|
||||
"sampling_freq": self.sampling_freq,
|
||||
}
|
||||
return {**base_params, **ts_params}
|
||||
|
||||
|
||||
class UnivariateTimeSeriesModel(TimeSeriesModel, RegressorMixin):
|
||||
"""
|
||||
Base class for univariate time series models.
|
||||
|
||||
Only supports regression settings. Concrete subclasses
|
||||
must implement `_fit_logic` and `predict`.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
@ensure_fitted
|
||||
def forecast(self, forecast_horizon: int) -> Sequence:
|
||||
"""
|
||||
Forecast into the future for a given number of steps.
|
||||
|
||||
Args:
|
||||
forecast_horizon: Number of future time steps to forecast.
|
||||
|
||||
Returns:
|
||||
Sequence of forecasted values.
|
||||
"""
|
||||
raise NotImplementedError("Subclasses must implement forecast().")
|
||||
|
||||
|
||||
class MultivariateTimeSeriesModel(TimeSeriesModel):
|
||||
"""
|
||||
Base class for multivariate time series models.
|
||||
|
||||
This class provides a foundation for time series models that utilize
|
||||
multiple exogenous features (X) to predict a target variable (y).
|
||||
It supports both regression and classification tasks.
|
||||
|
||||
Attributes:
|
||||
selected_features_: List of feature names selected for the model.
|
||||
learning_task: Type of learning task ('regression', 'binary',
|
||||
'multiclass').
|
||||
differentiate_target: Whether to apply differencing to make series
|
||||
stationary.
|
||||
bins: Bin edges for multiclass classification target
|
||||
transformation.
|
||||
|
||||
Example:
|
||||
>>> class MyModel(MultivariateTimeSeriesModel):
|
||||
... def _fit_logic(self, y, X=None, **kwargs):
|
||||
... # Implementation here
|
||||
... pass
|
||||
... def predict(self, X=None):
|
||||
... # Implementation here
|
||||
... return predictions
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: Optional[str] = None,
|
||||
time_col: str = "ds",
|
||||
target_col: str = "y",
|
||||
random_seed: int = 42,
|
||||
n_lags: int = 0,
|
||||
sampling_freq: Optional[str] = None,
|
||||
differentiate_target: bool = False,
|
||||
bins: Optional[List[float]] = None,
|
||||
learning_task: Optional[str] = None,
|
||||
):
|
||||
# Set attributes before calling parent constructor
|
||||
# This is needed because parent constructor calls
|
||||
# _validate_configuration
|
||||
self.selected_features_: Optional[List[str]] = None
|
||||
self.learning_task: Optional[str] = learning_task
|
||||
self.differentiate_target = differentiate_target
|
||||
self.bins = bins
|
||||
self.model_: Optional[PredictorProtocol] = None
|
||||
|
||||
super().__init__(
|
||||
name=name,
|
||||
time_col=time_col,
|
||||
target_col=target_col,
|
||||
random_seed=random_seed,
|
||||
n_lags=n_lags,
|
||||
sampling_freq=sampling_freq,
|
||||
)
|
||||
|
||||
# Additional validation for multivariate models
|
||||
self._validate_learning_task()
|
||||
|
||||
def _validate_learning_task(self) -> None:
|
||||
"""Validate learning task configuration."""
|
||||
valid_tasks = {"regression", "binary", "multiclass", None}
|
||||
if self.learning_task not in valid_tasks:
|
||||
raise ValueError(
|
||||
f"Invalid learning_task: {self.learning_task}. "
|
||||
+ f"Must be one of {valid_tasks}"
|
||||
)
|
||||
|
||||
if self.learning_task == "multiclass" and not self.bins:
|
||||
raise ValueError(
|
||||
"bins must be provided for multiclass learning_task"
|
||||
)
|
||||
|
||||
def _get_default_loss_function(
|
||||
self, provided_loss: Optional[str]
|
||||
) -> str:
|
||||
"""
|
||||
Get default loss function based on learning task.
|
||||
|
||||
Args:
|
||||
provided_loss: User-provided loss function (takes precedence)
|
||||
|
||||
Returns:
|
||||
str: Appropriate loss function for the learning task
|
||||
"""
|
||||
if provided_loss is not None:
|
||||
return provided_loss
|
||||
|
||||
if self.learning_task == "regression":
|
||||
return "RMSE"
|
||||
elif self.learning_task == "binary":
|
||||
return "Logloss"
|
||||
elif self.learning_task == "multiclass":
|
||||
return "MultiClass"
|
||||
else:
|
||||
return "RMSE"
|
||||
|
||||
def _validate_configuration(self) -> None:
|
||||
"""Validate model configuration."""
|
||||
super()._validate_configuration()
|
||||
|
||||
if self.differentiate_target and self.learning_task in [
|
||||
"binary",
|
||||
"multiclass",
|
||||
]:
|
||||
console.print(
|
||||
"[yellow]Warning: Using differentiation with classification "
|
||||
+ "tasks may not be appropriate[/yellow]"
|
||||
)
|
||||
|
||||
@ensure_fitted
|
||||
def feature_importance(self) -> Optional[pd.DataFrame]:
|
||||
"""
|
||||
Returns feature importance if implemented by subclass.
|
||||
|
||||
Returns:
|
||||
A DataFrame with feature names and their importance scores,
|
||||
or None if not applicable.
|
||||
"""
|
||||
return None
|
||||
|
||||
def _validate_X_y(
|
||||
self, X: pd.DataFrame, y: pd.Series, allow_nan: bool = True
|
||||
) -> Tuple[np.ndarray, np.ndarray]:
|
||||
"""
|
||||
Validates input features (X) and target (y).
|
||||
|
||||
Infers and sets `feature_names_in_` and `n_features_in_`.
|
||||
This method should be called within the `_fit_logic` of
|
||||
concrete subclasses that use exogenous features.
|
||||
|
||||
Args:
|
||||
X: DataFrame of input features.
|
||||
y: Series for the target variable.
|
||||
allow_nan: If True, allows NaN values in X and y.
|
||||
Raises:
|
||||
TypeError: If X is not a DataFrame or y is not a Series.
|
||||
ValueError: If the number of features in X does not match
|
||||
the expected number of features.
|
||||
|
||||
Returns:
|
||||
Tuple of validated NumPy arrays (X_array, y_array).
|
||||
"""
|
||||
if allow_nan:
|
||||
X_array, y_array = check_X_y(X, y, force_all_finite=False)
|
||||
else:
|
||||
X_array, y_array = check_X_y(X, y, force_all_finite=True)
|
||||
|
||||
if hasattr(X, "columns"):
|
||||
console.log(
|
||||
f"Validating input features with columns: {X.columns.tolist()}"
|
||||
)
|
||||
self.feature_names_in_ = list(X.columns)
|
||||
else:
|
||||
console.log(
|
||||
"Input features do not have column names, using default names."
|
||||
)
|
||||
self.feature_names_in_ = [
|
||||
f"feature_{i}" for i in range(X_array.shape[1])
|
||||
]
|
||||
|
||||
self.n_features_in_ = X_array.shape[1]
|
||||
return X_array, y_array
|
||||
|
||||
def _validate_X(
|
||||
self, X: pd.DataFrame, allow_nan: bool = True
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Validates input features (X) before prediction or scoring.
|
||||
|
||||
Ensures consistency with features seen during fit. This should
|
||||
be called by concrete subclasses in `predict`, `score`, etc.
|
||||
|
||||
Args:
|
||||
X: DataFrame of input features.
|
||||
allow_nan: If True, allows NaN values in X.
|
||||
|
||||
Returns:
|
||||
Validated NumPy array of X.
|
||||
"""
|
||||
if allow_nan:
|
||||
X_array = check_array(X, force_all_finite=False)
|
||||
else:
|
||||
X_array = check_array(X, force_all_finite=True)
|
||||
# If the model has been fitted, ensure the input features
|
||||
# match the features seen during fit.
|
||||
if self.feature_names_in_ is not None:
|
||||
if not set(self.feature_names_in_).issubset(X.columns):
|
||||
raise ValueError(
|
||||
"Input features do not match the features seen during fit."
|
||||
+ f" Expected features: {self.feature_names_in_}, "
|
||||
+ f"but got: {list(X.columns)}."
|
||||
)
|
||||
|
||||
return X_array
|
||||
|
||||
def _transform_target_to_multiclass(
|
||||
self, y: pd.Series, bins: Optional[List[float]] = None
|
||||
) -> pd.Series:
|
||||
"""
|
||||
Transforms the target variable into a multiclass classification
|
||||
target.
|
||||
|
||||
If bins are provided, uses pd.cut to categorize the target into
|
||||
discrete classes. If not, binarize the target at zero (0).
|
||||
|
||||
Args:
|
||||
y: The target variable as a pandas Series.
|
||||
bins: Optional list of bin edges for categorization.
|
||||
|
||||
Returns:
|
||||
A pandas Series with transformed classification targets.
|
||||
"""
|
||||
if bins is not None:
|
||||
# pd.cut returns a Categorical, convert to Series with integer
|
||||
# codes
|
||||
categories = pd.cut(y, bins=bins, labels=False)
|
||||
return pd.Series(categories, index=y.index)
|
||||
|
||||
return (y > 0).astype(int)
|
||||
|
||||
def _transform_target_to_binary(
|
||||
self, y: pd.Series, threshold: float = 0.0
|
||||
) -> pd.Series:
|
||||
"""
|
||||
Transforms the target variable into a binary classification target.
|
||||
|
||||
Binarizes the target at the specified threshold (default is 0.0).
|
||||
|
||||
Args:
|
||||
y: The target variable as a pandas Series.
|
||||
threshold: The threshold for binarization.
|
||||
|
||||
Returns:
|
||||
A pandas Series with binary classification targets.
|
||||
"""
|
||||
return (y > threshold).astype(int)
|
||||
|
||||
def _preprocess_data(
|
||||
self,
|
||||
y: pd.Series,
|
||||
X: Optional[pd.DataFrame] = None,
|
||||
X_val: Optional[pd.DataFrame] = None,
|
||||
y_val: Optional[pd.Series] = None,
|
||||
) -> Tuple[
|
||||
pd.Series,
|
||||
Optional[pd.DataFrame],
|
||||
Optional[pd.Series],
|
||||
Optional[pd.DataFrame],
|
||||
]:
|
||||
"""
|
||||
Internal method to handle common data preprocessing operations.
|
||||
|
||||
Args:
|
||||
y: The target time series data
|
||||
X: The feature matrix (including exogenous features)
|
||||
X_val: Validation feature matrix (optional)
|
||||
y_val: Validation target series (optional)
|
||||
|
||||
Returns:
|
||||
A tuple containing:
|
||||
- processed y series
|
||||
- processed X dataframe (optional)
|
||||
- processed y_val series (optional)
|
||||
- processed X_val dataframe (optional)
|
||||
"""
|
||||
# Apply differentiation if enabled
|
||||
if self.differentiate_target:
|
||||
y = y.diff().dropna()
|
||||
if X is not None:
|
||||
X = X.loc[y.index]
|
||||
|
||||
# Transform target for classification if needed
|
||||
if self.learning_task == "binary":
|
||||
y = self._transform_target_to_binary(y)
|
||||
elif self.learning_task == "multiclass":
|
||||
y = self._transform_target_to_multiclass(y, self.bins)
|
||||
|
||||
# Process validation data if provided
|
||||
if y_val is not None:
|
||||
if X_val is None:
|
||||
raise ValueError(
|
||||
"Validation features (X_val) must be provided if "
|
||||
+ "validation target (y_val) is given."
|
||||
)
|
||||
y_val = y_val.loc[X_val.index]
|
||||
if self.differentiate_target:
|
||||
y_val = y_val.diff().dropna()
|
||||
X_val = X_val.loc[y_val.index]
|
||||
if self.learning_task == "binary":
|
||||
y_val = self._transform_target_to_binary(y_val)
|
||||
elif self.learning_task == "multiclass":
|
||||
y_val = self._transform_target_to_multiclass(y_val, self.bins)
|
||||
|
||||
# Filter features if selected_features_ is set
|
||||
if X is not None and self.selected_features_ is not None:
|
||||
X = cast(pd.DataFrame, X[self.selected_features_].copy())
|
||||
if X_val is not None:
|
||||
X_val = cast(
|
||||
pd.DataFrame, X_val[self.selected_features_].copy()
|
||||
)
|
||||
|
||||
return y, X, y_val, X_val
|
||||
|
||||
def _prepare_shap_data(self, X: pd.DataFrame) -> pd.DataFrame:
|
||||
"""Prepare data for SHAP analysis."""
|
||||
X_processed = X.copy()
|
||||
|
||||
# Remove target column if present
|
||||
if self.target_col in X_processed.columns:
|
||||
X_processed = X_processed.drop(columns=[self.target_col])
|
||||
|
||||
# Filter selected features
|
||||
if self.selected_features_ is not None:
|
||||
X_processed = cast(
|
||||
pd.DataFrame, X_processed[self.selected_features_].copy()
|
||||
)
|
||||
|
||||
return X_processed
|
||||
|
||||
def _create_shap_explainer(self, X: pd.DataFrame) -> Any:
|
||||
"""Create appropriate SHAP explainer based on model type."""
|
||||
if self.model_ is None:
|
||||
raise ValueError("Model is not fitted yet.")
|
||||
|
||||
if hasattr(self.model_, "coef_"): # Linear models
|
||||
try:
|
||||
# Handle missing values if model supports it
|
||||
X_clean = self._handle_missing_values_for_shap(X)
|
||||
return shap.LinearExplainer(self.model_, X_clean)
|
||||
except Exception as e:
|
||||
console.print(
|
||||
f"[yellow]Warning: Linear explainer failed: {e}, "
|
||||
+ "using KernelExplainer[/yellow]"
|
||||
)
|
||||
background = shap.maskers.Independent(X, max_samples=100)
|
||||
return shap.KernelExplainer(
|
||||
self.model_.predict,
|
||||
background,
|
||||
)
|
||||
else:
|
||||
# Non-linear models
|
||||
if self.learning_task == "binary":
|
||||
return shap.TreeExplainer(
|
||||
self.model_, X, model_output="probability"
|
||||
)
|
||||
else:
|
||||
return shap.Explainer(self.model_, X)
|
||||
|
||||
def _handle_missing_values_for_shap(self, X: pd.DataFrame) -> pd.DataFrame:
|
||||
"""Handle missing values for SHAP analysis."""
|
||||
# Use type ignore for optional method
|
||||
if hasattr(self.model_, "_impute_missing_values"):
|
||||
return self.model_._impute_missing_values(X) # type: ignore
|
||||
else:
|
||||
return X.dropna()
|
||||
|
||||
def _generate_and_save_plot(
|
||||
self, explainer: Any, X: pd.DataFrame, path: str
|
||||
) -> None:
|
||||
"""Generate and save SHAP plot."""
|
||||
shap_values = explainer(X)
|
||||
|
||||
shap.plots.beeswarm(shap_values, show=False)
|
||||
shap_fig = plt.gcf()
|
||||
shap_fig.set_size_inches(10, 6)
|
||||
shap_fig.suptitle(f"SHAP Beeswarm Plot for {self.name}", fontsize=16)
|
||||
shap_fig.tight_layout()
|
||||
shap_fig.savefig(path)
|
||||
plt.clf()
|
||||
plt.close()
|
||||
|
||||
@ensure_fitted
|
||||
def shap_beeswarm_plot(self, X: pd.DataFrame, path: str) -> None:
|
||||
"""
|
||||
Generates a SHAP beeswarm plot for the model's predictions.
|
||||
|
||||
Args:
|
||||
X: DataFrame of input features.
|
||||
path: Path to save the plot file.
|
||||
|
||||
Raises:
|
||||
ValueError: If model is not fitted.
|
||||
Exception: If SHAP plot generation fails.
|
||||
"""
|
||||
if self.model_ is None:
|
||||
raise ValueError("Model is not fitted yet.")
|
||||
|
||||
try:
|
||||
# Prepare data
|
||||
X_processed = self._prepare_shap_data(X)
|
||||
|
||||
# Create explainer and generate plot
|
||||
explainer = self._create_shap_explainer(X_processed)
|
||||
self._generate_and_save_plot(explainer, X_processed, path)
|
||||
|
||||
except Exception as e:
|
||||
console.print(
|
||||
f"[red]Error: Failed to generate SHAP plot: {e}[/red]"
|
||||
)
|
||||
raise
|
||||
|
||||
def get_params_dict(self) -> dict:
|
||||
"""Get model parameters as dictionary for logging/serialization."""
|
||||
base_params = super().get_params_dict()
|
||||
mv_params = {
|
||||
"learning_task": self.learning_task,
|
||||
"differentiate_target": self.differentiate_target,
|
||||
"selected_features_": self.selected_features_,
|
||||
}
|
||||
return {**base_params, **mv_params}
|
||||
|
||||
def summary(self) -> str:
|
||||
"""Generate a summary string of the model."""
|
||||
params = self.get_params_dict()
|
||||
fitted_status = (
|
||||
"✓ Fitted" if self.__sklearn_is_fitted__() else "✗ Not fitted"
|
||||
)
|
||||
|
||||
summary_lines = [
|
||||
f"Model: {self.__class__.__name__}",
|
||||
f"Status: {fitted_status}",
|
||||
f"Features: {params.get('n_features_in_', 'Unknown')}",
|
||||
f"Task: {params.get('learning_task', 'regression')}",
|
||||
f"Selected Features: {len(self.selected_features_) if self.selected_features_ else 'All'}",
|
||||
]
|
||||
|
||||
return "\n".join(summary_lines)
|
||||
|
||||
Reference in New Issue
Block a user