235 lines
8.1 KiB
Python
235 lines
8.1 KiB
Python
"""
|
|
Standardized preprocessor for time series models used with XGBoost wrappers.
|
|
|
|
Args:
|
|
- None
|
|
|
|
Return:
|
|
None
|
|
"""
|
|
|
|
import pandas as pd
|
|
from sklearn.base import BaseEstimator, TransformerMixin
|
|
from sklearn.preprocessing import MinMaxScaler, StandardScaler
|
|
|
|
|
|
class TimeSeriesPreprocessor(BaseEstimator, TransformerMixin):
|
|
def __init__(
|
|
self,
|
|
target: str | None = None,
|
|
tags_list: list[str] | None = None,
|
|
scaler_method: str | None = None,
|
|
window_size: int = 0,
|
|
use_filtering: bool = False,
|
|
transform_mode: str = 'all',
|
|
arpr_config: dict[str, str] | None = None,
|
|
) -> None:
|
|
"""
|
|
Initialize the time series preprocessor for feature engineering and scaling.
|
|
|
|
Args:
|
|
- target (str | None): Name of the target column. If None, it must be
|
|
set before calling `fit`.
|
|
- tags_list (list[str] | None): Explicit list of feature columns. If
|
|
None, inferred during `fit` by excluding target (and CI column when
|
|
AR/PR filtering is enabled).
|
|
- scaler_method (str | None): Scaling strategy, either `'MinMax'`,
|
|
`'Standard'` or None for no scaling.
|
|
- window_size (int): Rolling window size used to generate lagged
|
|
statistics for the target. Zero disables rolling features.
|
|
- use_filtering (bool): Whether to filter rows based on a CI column
|
|
during `fit` when `arpr_config` is provided.
|
|
- transform_mode (str): Prediction mode, `'all'` to keep all rows or
|
|
`'latest'` to return only the last row.
|
|
- arpr_config (dict[str, str] | None): Optional configuration mapping
|
|
`'ci_col'`, `'ar_col'` and `'pr_col'` names used for AR/PR substitution.
|
|
|
|
Return:
|
|
None
|
|
"""
|
|
self.target = target
|
|
self.tags_list = tags_list
|
|
self.scaler_method = scaler_method
|
|
self.window_size = window_size
|
|
self.use_filtering = use_filtering
|
|
self.transform_mode = transform_mode
|
|
|
|
if arpr_config is None:
|
|
arpr_config = {}
|
|
|
|
self.ci_col = arpr_config.get('ci_col')
|
|
self.ar_col = arpr_config.get('ar_col')
|
|
self.pr_col = arpr_config.get('pr_col')
|
|
|
|
if self.ci_col is None or self.ar_col is None or self.pr_col is None:
|
|
self.ar_filter = False
|
|
else:
|
|
self.ar_filter = True
|
|
|
|
# Internal state
|
|
self.scaler = None
|
|
|
|
def _adjust_time_index(self, x: pd.DataFrame) -> pd.DataFrame:
|
|
"""
|
|
Standardize the time index so it is compatible with XGBoost expectations.
|
|
|
|
Args:
|
|
- x (pd.DataFrame): Input DataFrame that may contain a time column.
|
|
|
|
Return:
|
|
pd.DataFrame: Copy of the input with the index set to `Timestamp` or
|
|
`timestamp` when present; otherwise the original index is preserved.
|
|
"""
|
|
df = x.copy()
|
|
time_cols = ['Timestamp', 'timestamp']
|
|
for col in time_cols:
|
|
if col in df.columns:
|
|
df[col] = pd.to_datetime(df[col])
|
|
df = df.set_index(col)
|
|
break
|
|
return df
|
|
|
|
def _generate_rolling_features(self, df: pd.DataFrame) -> pd.DataFrame:
|
|
"""
|
|
Generate rolling statistics for the target column.
|
|
|
|
The method computes rolling mean, max, min and standard deviation for
|
|
the configured target when `window_size` is greater than zero.
|
|
|
|
Args:
|
|
- df (pd.DataFrame): Input DataFrame indexed by time and containing
|
|
the target column.
|
|
|
|
Return:
|
|
pd.DataFrame: DataFrame with new rolling feature columns added. When
|
|
the window is larger than the number of rows, the result may be empty
|
|
after `dropna`.
|
|
"""
|
|
|
|
if self.target is None:
|
|
raise ValueError('target is not set. Set it before calling fit.')
|
|
|
|
if self.window_size > 0:
|
|
df = df.sort_index()
|
|
if df.shape[0] < self.window_size:
|
|
# If we don't have enough data for the window, we might have issues
|
|
# but we'll try to calculate what we can.
|
|
pass
|
|
|
|
df['rolling_mean'] = df[self.target].rolling(window=self.window_size).mean()
|
|
df['rolling_max'] = df[self.target].rolling(window=self.window_size).max()
|
|
df['rolling_min'] = df[self.target].rolling(window=self.window_size).min()
|
|
df['rolling_std'] = df[self.target].rolling(window=self.window_size).std()
|
|
df = df.dropna()
|
|
return df
|
|
|
|
def _apply_ar_substitution(self, df: pd.DataFrame) -> pd.DataFrame:
|
|
"""
|
|
Apply AR/PR substitution based on the configured CI (confidence) column.
|
|
|
|
When AR/PR filtering is enabled and the CI column is present, rows
|
|
where `ci_col == 0` have the AR column replaced by the PR column.
|
|
|
|
Args:
|
|
- df (pd.DataFrame): Input DataFrame containing AR, PR and CI columns.
|
|
|
|
Return:
|
|
pd.DataFrame: DataFrame with AR values updated where CI indicates PR
|
|
substitution, or unchanged when filtering is disabled or misconfigured.
|
|
"""
|
|
if self.ar_filter:
|
|
ar_col = self.ar_col
|
|
pr_col = self.pr_col
|
|
ci_col = self.ci_col
|
|
|
|
if all([ar_col, pr_col, ci_col]) and ci_col in df.columns:
|
|
mask = df[ci_col] == 0
|
|
if mask.any():
|
|
df.loc[mask, ar_col] = df.loc[mask, pr_col]
|
|
return df
|
|
|
|
def fit(self, x: pd.DataFrame) -> None:
|
|
"""
|
|
Fit the preprocessor on the full dataset.
|
|
|
|
This step may filter rows, generate rolling features, infer `tags_list`
|
|
when it is not provided, and fit the underlying scaler if configured.
|
|
|
|
Args:
|
|
- x (pd.DataFrame): Full dataset containing the target and feature
|
|
columns, and optionally CI/AR/PR columns.
|
|
|
|
Return:
|
|
None
|
|
"""
|
|
if self.target is None:
|
|
raise ValueError('target is not set. Set it before calling fit.')
|
|
|
|
df = self._adjust_time_index(x)
|
|
|
|
# Optional filtering during fit
|
|
if self.use_filtering and self.ar_filter:
|
|
ci_col = self.ci_col
|
|
if ci_col in df.columns:
|
|
df = df[df[ci_col] == 1]
|
|
|
|
# Generate rolling features if needed
|
|
df = self._generate_rolling_features(df)
|
|
|
|
# Infer tags_list if not provided
|
|
if self.tags_list is None:
|
|
exclude_cols = [self.target]
|
|
if self.ar_filter and self.ci_col is not None:
|
|
exclude_cols.append(self.ci_col)
|
|
self.tags_list = [c for c in df.columns if c not in exclude_cols]
|
|
|
|
# Fit scaler
|
|
if self.scaler_method == 'MinMax':
|
|
self.scaler = MinMaxScaler()
|
|
elif self.scaler_method == 'Standard':
|
|
self.scaler = StandardScaler()
|
|
|
|
if self.scaler:
|
|
# Only fit on tags_list_
|
|
self.scaler.fit(df[self.tags_list])
|
|
|
|
def transform(self, x: pd.DataFrame) -> pd.DataFrame:
|
|
"""
|
|
Transform input data using the fitted preprocessor.
|
|
|
|
The transformation applies the same steps used during `fit` (time
|
|
indexing, AR/PR substitution and rolling features), then selects the
|
|
feature columns and optionally applies scaling and latest-row selection.
|
|
|
|
Args:
|
|
- x (pd.DataFrame): Input data to transform.
|
|
|
|
Return:
|
|
pd.DataFrame: Transformed feature matrix ready to be passed to the
|
|
model wrapper.
|
|
"""
|
|
if self.tags_list is None:
|
|
raise ValueError('tags_list is not set. Call fit first.')
|
|
|
|
df = self._adjust_time_index(x)
|
|
|
|
# Apply AR/PR substitution if configured
|
|
df = self._apply_ar_substitution(df)
|
|
|
|
# Generate rolling features
|
|
df = self._generate_rolling_features(df)
|
|
|
|
# Select features
|
|
df = df[self.tags_list]
|
|
|
|
# Apply scaling
|
|
if self.scaler:
|
|
data_scaled = self.scaler.transform(df)
|
|
df = pd.DataFrame(data=data_scaled, columns=self.tags_list, index=df.index)
|
|
|
|
# Transform mode: only latest sample
|
|
if self.transform_mode == 'latest':
|
|
df = df.tail(1)
|
|
|
|
return df
|