SIENTIAPDE-1214

SIENTIAPDE-1214: Refactor MLFlow and model repository methods to use model_config dictionary

- Updated MLFlow class methods to accept model_config instead of model_retention for improved flexibility.
- Modified model_repository methods to handle model_config, extracting necessary parameters for transformation and prediction.
- Adjusted predictions_batch and prediction_process workflows to utilize model_config for better configuration management.
- Commented out the previous sientia-mlops-library dependency in requirements.txt for clarity.
This commit is contained in:
vitor-aignosi
2025-09-12 14:23:35 -03:00
parent 4a4043c354
commit 6df04b72e5
35 changed files with 5333 additions and 18 deletions

View File

@@ -94,7 +94,7 @@ class MLFlow(BaseActivity):
self.info('Transforming data...', metadata)
data = DataFrame(input_data['data'])
model_name = input_data['model_name']
model_retention = input_data['model_retention']
model_config = input_data.get('model_config', {})
self.debug("Raw input data:", metadata)
self.debug(data, metadata)
@@ -117,10 +117,11 @@ class MLFlow(BaseActivity):
# Request transformation from MLFlow model
response_data = self.model_monitoring_repository.transform(
model_name, data, model_retention)
model_name, data, model_config
)
self.debug("Transform response data:", metadata)
self.debug(json.dumps(response_data, indent=4), metadata)
self.debug(response_data, metadata)
self.info("Data transformed successfully", metadata)
@@ -160,7 +161,7 @@ class MLFlow(BaseActivity):
self.info('Predicting data...', metadata)
data = DataFrame(input_data['data'])
model_name = input_data['model_name']
model_retention = input_data['model_retention']
model_config = input_data.get('model_config', {})
self.debug(data, metadata)
@@ -169,7 +170,8 @@ class MLFlow(BaseActivity):
# Request prediction from MLFlow model
response_data = self.model_monitoring_repository.predict(
model_name, data, model_retention)
model_name, data, model_config
)
self.debug("Prediction response data:", metadata)
self.debug(json.dumps(response_data, indent=4), metadata)

View File

@@ -25,7 +25,7 @@ class MLFlowRepository():
username=username, password=password,
logger=logger)
def transform(self, model_name: str, data: pd.DataFrame, model_retention: int):
def transform(self, model_name: str, data: pd.DataFrame, model_config: dict) -> dict:
"""
Transform data using a model.
@@ -39,11 +39,19 @@ class MLFlowRepository():
"""
try:
model_retention = model_config.get('model_retention', 0)
flavor = model_config.get('transform_flavor', 'sklearn')
compressed = model_config.get('is_compressed', False)
retention_target = model_config.get('retention_target', 'model')
transform_keyword = model_config.get(
'transform_function_keyword', 'predict')
return {
'success': True,
'content': self.model_serving.get_cached_transform(
model_name, data, model_retention).to_dict()
model_name, data, model_retention, flavor,
compressed, retention_target, transform_keyword
).to_dict()
}
except Exception as e:
@@ -55,7 +63,7 @@ class MLFlowRepository():
}
}
def predict(self, model_name: str, data: pd.DataFrame, model_retention: int):
def predict(self, model_name: str, data: pd.DataFrame, model_config: dict) -> dict:
"""
Predict data using a model.
@@ -68,11 +76,17 @@ class MLFlowRepository():
- dict: A dictionary containing the predicted data.
"""
try:
model_retention = model_config.get('retention_minutes', 0)
flavor = model_config.get('predict_flavor', 'pyfunc')
compressed = model_config.get('compressed', False)
retention_target = model_config.get('retention_target', 'model')
input_index = data.index
start_time = datetime.now()
data = self.model_serving.get_cached_predict(
model_name, data, model_retention)
model_name, data, model_retention, flavor,
compressed, retention_target
)
end_time = datetime.now()
data = pd.DataFrame(data, columns=['prediction'])

View File

@@ -113,7 +113,7 @@ class PredictionsBatch():
'POLICY': 'STOP'
}
}),
'model_retention': input_data.get('model_retention', 60),
'model_config': input_data.get('model_config', {}),
'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']),
'opc_output_config': input_data.get('opc_output_config', {})
}

View File

@@ -79,7 +79,7 @@ class PredictionProcess():
data = input_data['data']
model_id = input_data['model_id']
model_name = input_data['model_name']
model_retention = input_data['model_retention']
model_config = input_data.get('model_config', {})
# Get last timestamp for incremental processing
last_timestamp = await workflow.execute_local_activity_method(
@@ -120,7 +120,7 @@ class PredictionProcess():
**metadata,
'data': data,
'model_name': model_name,
'model_retention': model_retention
'model_config': model_config
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
@@ -172,7 +172,7 @@ class PredictionProcess():
**metadata,
'data': transformed_data,
'model_name': model_name,
'model_retention': model_retention
'model_config': model_config
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
@@ -209,7 +209,7 @@ class PredictionProcess():
'timestamp': last_timestamp,
'model_id': model_id,
'model_name': model_name,
'model_retention': model_retention,
'model_config': model_config,
'opc_output_config': input_data['opc_output_config'],
'schema': input_data['schema'],
'table_name': input_data['table_name'],
@@ -248,7 +248,7 @@ class PredictionProcess():
table_name = input_data['table_name']
model_id = input_data['model_id']
model_name = input_data['model_name']
model_retention = input_data['model_retention']
model_config = input_data.get('model_config', {})
path_flag = path_flag.upper() if path_flag else ''
@@ -282,7 +282,7 @@ class PredictionProcess():
'timestamp': last_timestamp,
'model_id': model_id,
'model_name': model_name,
'model_retention': model_retention,
'model_config': model_config,
'schema': schema,
'table_name': table_name,
'comment': comment,

View File

@@ -4,5 +4,6 @@ sqlalchemy
asyncua
redis
git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.4.5
git+ssh://git@github.com/Aignosi/sientia-mlops-library.git@0.38.13
# git+ssh://git@github.com/Aignosi/sientia-mlops-library.git@0.38.13
/home/grezewave/Documents/projects/sientia/sientia-mlops-library
prometheus-client

View File

@@ -165,7 +165,7 @@
},
{
"cell_type": "code",
"execution_count": 8,
"execution_count": 1,
"id": "c61be7ab",
"metadata": {},
"outputs": [

View File

@@ -0,0 +1,19 @@
artifact_path: transformer_pyfunc
flavors:
python_function:
artifacts:
transformer:
path: artifacts/training_transformer.pkl
uri: /tmp/tmpnzkqz3v0/training_transformer.pkl
cloudpickle_version: 2.2.1
code: code
env:
conda: conda.yaml
virtualenv: python_env.yaml
loader_module: mlflow.pyfunc.model
python_model: python_model.pkl
python_version: 3.10.16
mlflow_version: 2.7.1
model_uuid: 6a4a99079d234d0da2b8091532d55a34
run_id: c2edec4dfafd4ad8bba257d62d25cd43
utc_time_created: '2025-09-09 12:30:04.885248'

View File

@@ -0,0 +1,172 @@
"""
Module with utility functions for data processing.
"""
import pandas as pd
from typing import Dict, List
from rich.console import Console
console = Console()
def create_lagged_target(
data: pd.DataFrame,
target_column: str,
lags: List[int],
drop_nans: bool = True,
) -> pd.DataFrame:
"""
Creates lagged versions of a target column in a DataFrame.
For each lag value in the provided list, a new column is created with
the naming pattern: target_column + "_lag_" + lag_value.
Args:
data: Input DataFrame containing the target column.
target_column: Name of the target column to create lags for.
lags: List of lag values (integers between 1 and len(data)-1).
drop_nans: Whether rows with nulls generated by the lag creation
process should be dropped. Defaults to True.
Returns:
DataFrame with original columns plus the newly created lag columns.
"""
if data.empty:
console.log("[red]Warning: Input data is empty.")
return data
result = data.copy()
if target_column not in result.columns:
raise ValueError(f"Target column '{target_column}' not found in data.")
# Validate lag values
max_lag = len(data) - 1
valid_lags = [lag for lag in lags if 1 <= lag <= max_lag]
if len(valid_lags) < len(lags):
invalid_lags = set(lags) - set(valid_lags)
console.log(
f"[yellow]Warning: Ignoring invalid lag values: {invalid_lags}. "
f"Lags must be between 1 and {max_lag}."
)
lag_column_names = []
for lag in valid_lags:
lag_column_name = f"{target_column}_lag_{lag}"
result[lag_column_name] = result[target_column].shift(lag)
console.log(f"Created lagged column: [cyan]{lag_column_name}")
lag_column_names.append(lag_column_name)
if drop_nans:
result = result.dropna(subset=lag_column_names)
return result
def remove_stopped_windows(
data: pd.DataFrame,
stopped_process_columns: Dict[str, float],
stopped_process_threshold: float,
time_colname: str,
) -> pd.DataFrame:
"""
Removes time windows from the input DataFrame if the proportion of samples
below a column threshold exceeds the specified limit.
A window is considered "stopped" if *all* specified columns exceed the
stopped sample threshold.
Args:
data: Input DataFrame with process variables and timestamps.
stopped_process_columns: Dict mapping column names to thresholds.
stopped_process_threshold: Proportion threshold (0-1) for marking a
window as stopped.
time_colname: Base name of the timestamp column
(without 'lab_' prefix).
Returns:
A DataFrame with stopped windows removed.
"""
if data.empty:
console.log("[red]Warning: Input data is empty.")
return data
console.log(
"Removing windows where any column exceeds"
+ f" {stopped_process_threshold:.2%} of values below threshold"
)
masks = []
for col, threshold in stopped_process_columns.items():
console.log(
"Evaluating stopped condition for column:"
+ f" [cyan]{col} < {threshold}"
)
below_threshold = data[[col]].lt(threshold)
console.log(
"Counting number of samples below threshold for each window"
)
below_threshold[f"lab_{time_colname}"] = data[f"lab_{time_colname}"]
grouped = below_threshold.groupby(f"lab_{time_colname}")[col].agg(
["sum", "count"]
)
stopped_mask = (
grouped["sum"] / grouped["count"]
) > stopped_process_threshold
console.log(
f"[red]{stopped_mask.sum()} windows marked as stopped by {col}"
)
masks.append(stopped_mask)
# Combine masks across columns: only drop if all agree
combined_mask = pd.concat(masks, axis=1).all(axis=1)
num_removed = combined_mask.sum()
total = combined_mask.shape[0]
console.log(
f"Removing [bold red]{num_removed}[/] out of {total}"
+ f" windows ({num_removed / total:.2%})"
)
to_remove = combined_mask[combined_mask].index
keep_mask = ~data[f"lab_{time_colname}"].isin(to_remove)
return data[keep_mask]
def aggregate_data(
merged_data: pd.DataFrame,
time_colname: str,
target_colname: str,
aggregation_functions: List[str],
) -> pd.DataFrame:
"""
Aggregates a DataFrame by time and target columns using specified
aggregation functions.
Args:
merged_data: Input DataFrame with raw observations.
time_colname: Name of the timestamp column (no 'lab_' prefix).
target_colname: Name of the target/grouping column.
aggregation_functions: List of aggregation functions to apply
(e.g. "mean", "std").
Returns:
Aggregated DataFrame with flattened column names and renamed time
column.
"""
group_by_cols = [f"lab_{time_colname}", target_colname]
aggregated = merged_data.groupby(group_by_cols).agg(aggregation_functions)
# Flatten MultiIndex columns
aggregated.columns = [
"_".join(col) if isinstance(col, tuple) else col
for col in aggregated.columns
] # type: ignore
aggregated = aggregated.reset_index()
console.log(f"[bold green]Aggregated shape: {aggregated.shape}")
return aggregated.rename(columns={f"lab_{time_colname}": time_colname})

View File

@@ -0,0 +1,56 @@
"""
Module with helper functions to read datasets.
"""
import pandas as pd
from openpyxl import load_workbook
from typing import Union
def read_excel_with_colors(
filepath: str, color_columns: list[str], sheet_name: Union[str, int] = 0
) -> pd.DataFrame:
"""
Read an Excel file and extract cell fill colors for specified columns.
Parameters:
filepath (str): Path to the Excel file.
color_columns (List[str]): Column names to extract fill colors from.
sheet_name (str or int): Sheet name or index (default is first sheet).
Returns:
DataFrame: DataFrame with original data and extra color columns.
"""
df = pd.read_excel(filepath, sheet_name=sheet_name)
workbook = load_workbook(filepath)
sheet = (
workbook[sheet_name]
if isinstance(sheet_name, str)
else workbook[workbook.sheetnames[sheet_name]]
)
header = next(sheet.iter_rows(min_row=1, max_row=1, values_only=True))
col_name_to_letter = {
name: chr(65 + idx) for idx, name in enumerate(header)
}
for col_name in color_columns:
if col_name not in df.columns:
raise ValueError(f"Column '{col_name}' not found in Excel file.")
col_letter = col_name_to_letter[col_name]
fill_colors: list[Union[str, None]] = []
for row in range(2, sheet.max_row + 1):
cell = sheet[f"{col_letter}{row}"]
fill = cell.fill
if fill.fill_type == "solid" and fill.fgColor.rgb:
fill_colors.append(fill.fgColor.rgb)
else:
fill_colors.append(None)
df[f"{col_name}_fill_color"] = fill_colors
return df

View File

@@ -0,0 +1,788 @@
"""
Module with scikit-learn transformers for training and inference pipelines.
This module replicates the functionality from the DVC pipeline stages:
- preprocessing.py (stopped process filtering, aggregation)
- filter_columns.py (feature selection, lagged target creation)
"""
import pandas as pd
from typing import Dict, List, Optional
from sklearn.base import BaseEstimator, TransformerMixin
from rich.console import Console
console = Console()
class CourierTrainingTransformer(BaseEstimator, TransformerMixin):
"""
Training transformer that replicates the DVC pipeline preprocessing.
Includes:
1. Remove stopped process windows (training only)
2. Data aggregation
3. Feature selection (learns and applies - dictionary-based only)
4. Create lagged target features
"""
def __init__(
self,
aggregation_functions: List[str] = ["median", "std", "min", "max"],
stopped_process_columns: Dict[str, float] = {
"305-PIT-170": 9,
"305-PIT-175": 9,
},
stopped_process_threshold: float = 0.1,
target_lags: List[int] = [2],
time_colname: str = "timestamp",
target_colname: str = "SiO2_conc",
dictionary_df: Optional[pd.DataFrame] = None,
create_lagged_target: bool = True,
drop_nans: bool = True,
):
"""
Initialize the training transformer.
Args:
aggregation_functions: List of aggregation functions to apply
stopped_process_columns: Dict mapping column names to thresholds
stopped_process_threshold: Proportion threshold for stopped windows
target_lags: List of lag values for target column
time_colname: Name of timestamp column (without 'lab_' prefix)
target_colname: Name of target column
dictionary_df: DataFrame with domain knowledge
(TAG_fill_color column)
create_lagged_target: Whether to create lagged target features
drop_nans: Whether to drop NaNs after creating lags
"""
self.aggregation_functions = aggregation_functions
self.stopped_process_columns = stopped_process_columns
self.stopped_process_threshold = stopped_process_threshold
self.target_lags = target_lags
self.time_colname = time_colname
self.target_colname = target_colname
self.dictionary_df = dictionary_df
self.create_lagged_target = create_lagged_target
self.drop_nans = drop_nans
# Will be learned during fit
self.selected_features_ = None
def _infer_lab_timestamp(self, data: pd.DataFrame) -> pd.DataFrame:
"""
Creates a lab timestamp column by rounding the timestamp up to the next
even hour.
Args:
data: Input DataFrame
Returns:
DataFrame with added lab timestamp column
"""
lab_col_name = f"lab_{self.time_colname}"
if lab_col_name in data.columns:
console.log(
f"[yellow]Lab timestamp column {lab_col_name} already exists, "
+ "skipping inference"
)
return data
# Check if timestamp is a column or the index
if self.time_colname in data.columns:
# Timestamp is a regular column
result = data.copy()
timestamp_col = pd.to_datetime(result[self.time_colname])
elif data.index.name == self.time_colname or (
hasattr(data.index, "names")
and self.time_colname in data.index.names
):
# Timestamp is the index (or part of a MultiIndex)
result = data.copy()
timestamp_col = pd.to_datetime(
result.index.get_level_values(self.time_colname)
if hasattr(result.index, "names")
and len(result.index.names) > 1
else result.index
)
else:
raise ValueError(
f"Timestamp '{self.time_colname}' not found in data columns "
+ f"or index. Available columns: {list(data.columns)}, "
+ f"index name: {data.index.name}"
)
# Round up to next even hour
# Step 1: Floor to the hour to remove minutes/seconds
# Handle both Series (from column) and DatetimeIndex (from index)
if hasattr(timestamp_col, "dt"):
# timestamp_col is a Series
hour_floor = timestamp_col.dt.floor("H")
hour = hour_floor.dt.hour
else:
# timestamp_col is a DatetimeIndex
hour_floor = timestamp_col.floor("H")
hour = hour_floor.hour
# Step 3: Determine if rounding is needed
# - If hour is odd, round up to next even hour
# - If hour is even but original timestamp had minutes/seconds,
# round up to next even hour
# - If hour is even and original timestamp was exactly on the hour,
# keep it
needs_rounding = (hour % 2 == 1) | (timestamp_col != hour_floor)
# Calculate next even hour
next_even_hour = ((hour // 2) + 1) * 2
# Handle case where next even hour >= 24 (next day)
days_to_add = (next_even_hour >= 24).astype(int)
hour_component = next_even_hour % 24
# Create the lab timestamp
if hasattr(timestamp_col, "dt"):
# timestamp_col is a Series
lab_timestamp = hour_floor.where(
~needs_rounding,
hour_floor.dt.floor("D")
+ pd.to_timedelta(days_to_add, unit="D")
+ pd.to_timedelta(hour_component, unit="H"),
)
else:
# timestamp_col is a DatetimeIndex
base_date = hour_floor.floor("D")
next_even_timestamp = (
base_date
+ pd.to_timedelta(days_to_add, unit="D")
+ pd.to_timedelta(hour_component, unit="H")
)
lab_timestamp = pd.Series(
hour_floor.where(~needs_rounding, next_even_timestamp),
index=result.index,
)
result[lab_col_name] = lab_timestamp
console.log(
f"[bold green]Created lab timestamp column: {lab_col_name}"
)
return result
def _remove_stopped_windows(
self,
data: pd.DataFrame,
) -> pd.DataFrame:
"""
Removes time windows where process was stopped.
Replicates remove_stopped_windows from preprocessing.py
"""
if data.empty:
console.log("[red]Warning: Input data is empty.")
return data
console.log(
"Removing windows where any column exceeds"
+ f" {self.stopped_process_threshold:.2%} of values below"
+ " threshold"
)
masks = []
for col, threshold in self.stopped_process_columns.items():
console.log(
"Evaluating stopped condition for column:"
+ f" [cyan]{col} < {threshold}"
)
below_threshold = data[[col]].lt(threshold)
console.log(
"Counting number of samples below threshold for each window"
)
below_threshold[f"lab_{self.time_colname}"] = data[
f"lab_{self.time_colname}"
]
grouped = below_threshold.groupby(f"lab_{self.time_colname}")[
col
].agg(["sum", "count"])
stopped_mask = (
grouped["sum"] / grouped["count"]
) > self.stopped_process_threshold
console.log(
f"[red]{stopped_mask.sum()} windows marked as stopped by {col}"
)
masks.append(stopped_mask)
# Combine masks across columns: only drop if all agree
if not masks:
return data
mask_df = pd.concat(masks, axis=1)
combined_mask = mask_df.all(axis=1)
num_removed = int(combined_mask.sum()) # type: ignore
total = combined_mask.shape[0]
console.log(
f"Removing [bold red]{num_removed}[/] out of {total}"
+ f" windows ({num_removed / total:.2%})"
)
to_remove = combined_mask[combined_mask].index
keep_mask = ~data[f"lab_{self.time_colname}"].isin(to_remove)
filtered_data = data[keep_mask]
return filtered_data # type: ignore
def _aggregate_data(
self,
merged_data: pd.DataFrame,
) -> pd.DataFrame:
"""
Aggregates data by time and target columns.
Replicates aggregate_data from preprocessing.py
"""
group_by_cols = [f"lab_{self.time_colname}", self.target_colname]
aggregated = merged_data.groupby(group_by_cols).agg(
self.aggregation_functions
)
# Flatten MultiIndex columns
aggregated.columns = [
"_".join(col) if isinstance(col, tuple) else col
for col in aggregated.columns
] # type: ignore
aggregated = aggregated.reset_index()
# Rename timestamp column and ensure it's datetime
aggregated = aggregated.rename(
columns={f"lab_{self.time_colname}": self.time_colname}
)
aggregated[self.time_colname] = pd.to_datetime(
aggregated[self.time_colname]
)
console.log(f"[bold green]Aggregated shape: {aggregated.shape}")
return aggregated
def _learn_feature_selection(
self,
data: pd.DataFrame,
) -> List[str]:
"""
Learn which features to keep based on dictionary only.
Replicates domain knowledge filtering from filter_columns.py
"""
if self.dictionary_df is None:
console.log(
"[yellow]Warning: No dictionary data provided. Using all"
+ " features."
)
return [col for col in data.columns if col != self.time_colname]
# Domain knowledge filter - only keep columns with TAG_fill_color
columns_to_keep_dict = self.dictionary_df.loc[
self.dictionary_df["TAG_fill_color"].notna(), "TAG"
].values
# Filter data columns to only those that match dictionary tags
available_columns = [
col for col in data.columns if col != self.time_colname
]
columns_to_keep = [
col
for col in available_columns
if col.split("_")[0] in columns_to_keep_dict
]
console.log(
f"Keeping {len(columns_to_keep)}/{len(available_columns)}"
+ " columns based on dictionary."
)
# Always include target
columns_to_keep.append(self.target_colname)
return columns_to_keep
def _create_lagged_target(
self,
data: pd.DataFrame,
) -> pd.DataFrame:
"""
Creates lagged versions of target column.
Replicates create_lagged_target from preprocessing.py
"""
if data.empty:
console.log("[red]Warning: Input data is empty.")
return data
result = data.copy()
if self.target_colname not in result.columns:
raise ValueError(
f"Target column '{self.target_colname}' not found in data."
)
# Validate lag values
max_lag = len(data) - 1
valid_lags = [lag for lag in self.target_lags if 1 <= lag <= max_lag]
if len(valid_lags) < len(self.target_lags):
invalid_lags = set(self.target_lags) - set(valid_lags)
console.log(
f"[yellow]Warning: Ignoring invalid lag values: {invalid_lags}"
+ f". Lags must be between 1 and {max_lag}."
)
lag_column_names = []
for lag in valid_lags:
lag_column_name = f"{self.target_colname}_lag_{lag}"
result[lag_column_name] = result[self.target_colname].shift(lag)
console.log(f"Created lagged column: [cyan]{lag_column_name}")
lag_column_names.append(lag_column_name)
if self.drop_nans and lag_column_names:
result = result.dropna(subset=lag_column_names)
return result
def fit(self, X: pd.DataFrame, y=None):
"""
Learn feature selection parameters.
Args:
X: Input DataFrame with merged process and quality data
y: Not used
Returns:
self
"""
console.log("[bold blue]Training transformer fit phase")
# Create a copy for processing
data = X.copy()
# Step 0: Infer lab timestamp if needed
console.log("[bold blue]Inferring lab timestamp")
data = self._infer_lab_timestamp(data)
# Step 1: Remove stopped process windows (training only)
console.log("[bold blue]Removing stopped process windows")
data = self._remove_stopped_windows(data)
# Step 2: Aggregate data
console.log("[bold blue]Aggregating data")
data = self._aggregate_data(data)
# Step 3: Learn feature selection
console.log("[bold blue]Learning feature selection")
self.selected_features_ = self._learn_feature_selection(data)
console.log(
f"[bold green]Learned {len(self.selected_features_)}"
+ " features for selection"
)
self._feature_names = self.selected_features_
return self
def transform(self, X: pd.DataFrame) -> pd.DataFrame:
"""
Apply the complete training transformation pipeline.
Args:
X: Input DataFrame with merged process and quality data
Returns:
Transformed DataFrame ready for model training
"""
if self.selected_features_ is None:
raise ValueError("Transformer must be fitted before transform.")
console.log("[bold blue]Training transformer transform phase")
# Create a copy for processing
data = X.copy()
# Step 0: Infer lab timestamp if needed
console.log("[bold blue]Inferring lab timestamp")
data = self._infer_lab_timestamp(data)
# Step 1: Remove stopped process windows (training only)
console.log("[bold blue]Removing stopped process windows")
data = self._remove_stopped_windows(data)
# Step 2: Aggregate data
console.log("[bold blue]Aggregating data")
data = self._aggregate_data(data)
# Step 3: Apply feature selection
console.log("[bold blue]Applying feature selection")
# Set timestamp as index for filtering and ensure it's datetime
data[self.time_colname] = pd.to_datetime(data[self.time_colname])
data = data.set_index(self.time_colname)
data = data[self.selected_features_]
# Step 4: Create lagged target features
if self.create_lagged_target:
console.log("[bold blue]Creating lagged target features")
data = self._create_lagged_target(data)
console.log(f"[bold green]Final training data shape: {data.shape}")
return data
class CourierInferenceTransformer(BaseEstimator, TransformerMixin):
"""
Inference transformer that replicates DVC pipeline preprocessing
without training-specific steps.
Includes:
1. Data aggregation (higher frequency - no grouping by target)
2. Feature selection (applies learned selection)
3. Create lagged target features
Note: Does NOT include stopped process filtering (training only).
"""
def __init__(
self,
selected_features: List[str],
aggregation_functions: List[str] = ["median", "std", "min", "max"],
target_lags: List[int] = [2],
time_colname: str = "timestamp",
target_colname: str = "SiO2_conc",
create_lagged_target: bool = True,
drop_nans: bool = True,
):
"""
Initialize the inference transformer.
Args:
selected_features: Pre-learned list of features to select
aggregation_functions: List of aggregation functions to apply
target_lags: List of lag values for target column
time_colname: Name of timestamp column (without 'lab_' prefix)
target_colname: Name of target column
create_lagged_target: Whether to create lagged target features
drop_nans: Whether to drop NaNs after creating lags
"""
self.selected_features = selected_features
self.aggregation_functions = aggregation_functions
self.target_lags = target_lags
self.time_colname = time_colname
self.target_colname = target_colname
self.create_lagged_target = create_lagged_target
self.drop_nans = drop_nans
def _infer_lab_timestamp(self, data: pd.DataFrame) -> pd.DataFrame:
"""
Creates a lab timestamp column by rounding the timestamp up to the next
even hour.
Args:
data: Input DataFrame
Returns:
DataFrame with added lab timestamp column
"""
lab_col_name = f"lab_{self.time_colname}"
if lab_col_name in data.columns:
console.log(
f"[yellow]Lab timestamp column {lab_col_name} already exists, "
+ "skipping inference"
)
return data
# Check if timestamp is a column or the index
if self.time_colname in data.columns:
# Timestamp is a regular column
result = data.copy()
timestamp_col = pd.to_datetime(result[self.time_colname])
elif data.index.name == self.time_colname or (
hasattr(data.index, "names")
and self.time_colname in data.index.names
):
# Timestamp is the index (or part of a MultiIndex)
result = data.copy()
timestamp_col = pd.to_datetime(
result.index.get_level_values(self.time_colname)
if hasattr(result.index, "names")
and len(result.index.names) > 1
else result.index
)
else:
raise ValueError(
f"Timestamp '{self.time_colname}' not found in data columns "
+ f"or index. Available columns: {list(data.columns)}, "
+ f"index name: {data.index.name}"
)
# Round up to next even hour
# Step 1: Floor to the hour to remove minutes/seconds
# Handle both Series (from column) and DatetimeIndex (from index)
if hasattr(timestamp_col, "dt"):
# timestamp_col is a Series
hour_floor = timestamp_col.dt.floor("H")
hour = hour_floor.dt.hour
else:
# timestamp_col is a DatetimeIndex
hour_floor = timestamp_col.floor("H")
hour = hour_floor.hour
# Step 3: Determine if rounding is needed
# - If hour is odd, round up to next even hour
# - If hour is even but original timestamp had minutes/seconds,
# round up to next even hour
# - If hour is even and original timestamp was exactly on the hour,
# keep it
needs_rounding = (hour % 2 == 1) | (timestamp_col != hour_floor)
# Calculate next even hour
next_even_hour = ((hour // 2) + 1) * 2
# Handle case where next even hour >= 24 (next day)
days_to_add = (next_even_hour >= 24).astype(int)
hour_component = next_even_hour % 24
# Create the lab timestamp
if hasattr(timestamp_col, "dt"):
# timestamp_col is a Series
lab_timestamp = hour_floor.where(
~needs_rounding,
hour_floor.dt.floor("D")
+ pd.to_timedelta(days_to_add, unit="D")
+ pd.to_timedelta(hour_component, unit="H"),
)
else:
# timestamp_col is a DatetimeIndex
base_date = hour_floor.floor("D")
next_even_timestamp = (
base_date
+ pd.to_timedelta(days_to_add, unit="D")
+ pd.to_timedelta(hour_component, unit="H")
)
lab_timestamp = pd.Series(
hour_floor.where(~needs_rounding, next_even_timestamp),
index=result.index,
)
result[lab_col_name] = lab_timestamp
console.log(
f"[bold green]Created lab timestamp column: {lab_col_name}"
)
return result
def _aggregate_data(
self,
merged_data: pd.DataFrame,
) -> pd.DataFrame:
"""
Aggregate data into 2-hour non-overlapping windows.
Each row corresponds to one 2-hour window ending at an even hour.
"""
if merged_data.empty:
console.log("[red]Warning: Input data is empty.")
return merged_data
lab_col = f"lab_{self.time_colname}"
if lab_col not in merged_data.columns:
raise ValueError(
f"Missing '{lab_col}' column. Call _infer_lab_timestamp first."
)
# Get numeric columns only for aggregation
numeric_cols = merged_data.select_dtypes(
include=["number"]
).columns.tolist()
# Remove time and target columns if present
cols_to_remove = [self.time_colname, lab_col, self.target_colname]
for col in cols_to_remove:
if col in numeric_cols:
numeric_cols.remove(col)
groups = merged_data.groupby(lab_col)
if numeric_cols:
aggregated_numeric = groups[numeric_cols].agg(
self.aggregation_functions
)
# Flatten MultiIndex columns: (col, func) -> "col_func"
aggregated_numeric.columns = [
f"{col}_{func}"
for col, func in aggregated_numeric.columns.to_flat_index()
]
else:
# Create empty frame indexed by the 2-hour windows
aggregated_numeric = groups.size().to_frame(name="__rows__")
aggregated_numeric = aggregated_numeric.drop(columns=["__rows__"])
# Add target column as the last non-null value per window
if self.target_colname in merged_data.columns:
target_per_window = groups[self.target_colname].apply(
lambda s: s.dropna().iloc[-1]
if not s.dropna().empty
else None
)
aggregated_numeric[self.target_colname] = target_per_window
# Reset index and rename lab timestamp to main time column
result = aggregated_numeric.reset_index().rename(
columns={lab_col: self.time_colname}
)
# Ensure timestamp is datetime
result[self.time_colname] = pd.to_datetime(result[self.time_colname])
console.log(
f"[bold green]Aggregated to windowed shape: {result.shape}"
)
return result
def _create_lagged_target(
self,
data: pd.DataFrame,
) -> pd.DataFrame:
"""
Creates lagged target column names with target values for inference.
In inference, we assume the data is already properly lagged,
so we just create the expected column names with the target values.
"""
if data.empty:
console.log("[red]Warning: Input data is empty.")
return data
result = data.copy()
if self.target_colname not in result.columns:
raise ValueError(
f"Target column '{self.target_colname}' not found in data."
)
# Create lagged column names with target values (no actual shifting)
lag_column_names = []
# Get target column as a Series to ensure we have exactly one column
target_series = result[self.target_colname]
if isinstance(target_series, pd.DataFrame):
# If we accidentally got a DataFrame, take the first column
target_values = target_series.iloc[:, 0].values
else:
target_values = target_series.values
for lag in self.target_lags:
lag_column_name = f"{self.target_colname}_lag_{lag}"
# Copy target values instead of shifting for inference
result[lag_column_name] = target_values
console.log(f"Created lagged column: [cyan]{lag_column_name}")
lag_column_names.append(lag_column_name)
return result
def fit(self, X: pd.DataFrame, y=None):
"""
No-op for inference transformer (no learning needed).
Args:
X: Input DataFrame
y: Not used
Returns:
self
"""
console.log("[bold blue]Inference transformer fit (no-op)")
return self
def transform(self, X: pd.DataFrame) -> pd.DataFrame:
"""
Apply the inference transformation pipeline.
Args:
X: Input DataFrame with merged process and quality data
Returns:
Transformed DataFrame ready for model inference
"""
console.log("[bold blue]Inference transformer transform phase")
# Create a copy for processing
data = X.copy()
# Step 0: Infer lab timestamp if needed
console.log("[bold blue]Inferring lab timestamp")
data = self._infer_lab_timestamp(data)
# Step 1: Aggregate data (higher frequency - no target grouping)
console.log("[bold blue]Aggregating data")
data = self._aggregate_data(data)
# Step 2: Apply learned feature selection
console.log("[bold blue]Applying learned feature selection")
# Set timestamp as index for filtering and ensure it's datetime
data[self.time_colname] = pd.to_datetime(data[self.time_colname])
data = data.set_index(self.time_colname)
# Filter to selected features (handle missing columns gracefully)
available_features = [
col for col in self.selected_features if col in data.columns
]
missing_features = set(self.selected_features) - set(
available_features
)
if missing_features:
console.log(
"[yellow]Warning: Missing features in inference data:"
+ f" {missing_features}"
)
# Ensure target column is included but avoid duplicates
if self.target_colname not in available_features:
available_features.append(self.target_colname)
data = data[available_features]
# Step 3: Create lagged target features
if self.create_lagged_target:
console.log("[bold blue]Creating lagged target features")
data = self._create_lagged_target(data)
console.log(f"[bold green]Final inference data shape: {data.shape}")
return data
def create_transformers_from_training_transformer(
training_transformer: CourierTrainingTransformer,
) -> tuple[CourierTrainingTransformer, CourierInferenceTransformer]:
"""
Create both training and inference transformers with shared parameters.
Args:
training_transformer: Fitted training transformer
Returns:
Tuple of (training_transformer, inference_transformer)
"""
if training_transformer.selected_features_ is None:
raise ValueError("Training transformer must be fitted first.")
inference_transformer = CourierInferenceTransformer(
selected_features=training_transformer.selected_features_,
aggregation_functions=training_transformer.aggregation_functions,
target_lags=training_transformer.target_lags,
time_colname=training_transformer.time_colname,
target_colname=training_transformer.target_colname,
create_lagged_target=training_transformer.create_lagged_target,
drop_nans=training_transformer.drop_nans,
)
return training_transformer, inference_transformer

View File

@@ -0,0 +1,51 @@
"""
Functions needed to load parameters from params.yaml tracked with DVC
"""
import sys
import os
from typing import Optional
import yaml
from rich.console import Console
console = Console()
def get_params(stage_fn: Optional[str] = None):
"""
Reads parameters for a given DVC stage from params.yaml.
The stage name is inferred from the name of the python file that calls this
function.
Args:
stage_fn (str): Name of the stage. If None, the name of the file
that calls this function is used. Defaults to None.
Returns:
dict with parameters for the stage
Raises:
KeyError: if the stage name is not found in params.yaml
"""
if stage_fn is None:
stage_fn = os.path.basename(sys.argv[0]).replace(".py", "")
try:
params = yaml.safe_load(open("params.yaml"))[stage_fn]
except KeyError as exc:
console.print(f'ERROR: Key "{stage_fn}" not in parameters.yaml.')
raise KeyError(
f"Is the stage file name ({sys.argv[0]}) "
+ "the same as the stage name in params.yaml?"
) from exc
try:
all_params = yaml.safe_load(open("params.yaml"))["all"]
params = {**params, **all_params}
except KeyError:
console.print(
'[orange]WARNING: Key "all" not in parameters.yaml.'
+ "Only returning stage parameters."
)
return params

View File

@@ -0,0 +1,325 @@
"""
Module with functions for wrapping time series models for MLflow.
"""
import os
import tempfile
import pickle
from typing import Optional, Union, Dict, Any, List
import mlflow.pyfunc
import pandas as pd
import numpy as np
from mlflow.models import ModelSignature
from ..models.stacking_time_series import StackingTimeSeriesModel
from ..data.transformers import (
create_transformers_from_training_transformer,
CourierTrainingTransformer,
)
class StackingWrapper(mlflow.pyfunc.PythonModel): # type: ignore
"""
MLflow wrapper for StackingTimeSeriesModel.
Allows the model to be saved and served via MLflow's pyfunc interface.
"""
def __init__(self, model: Optional[StackingTimeSeriesModel] = None):
self.model = model
@property
def _console(self):
from rich.console import Console
return Console()
def load_context(self, context: Any) -> None:
"""Load model from artifact path in MLflow context."""
try:
model_path = context.artifacts["model"]
self.model = StackingTimeSeriesModel.load(model_path)
self._console.log("[green]Model loaded from context[/green]")
except Exception as e:
self._console.print(f"[red]Error loading model: {e}[/red]")
raise
def predict(
self,
context: Any,
model_input: Union[pd.DataFrame, np.ndarray, Dict[str, Any]],
) -> Union[pd.Series, pd.DataFrame, np.ndarray]:
"""Run inference using the wrapped model."""
if self.model is None:
raise ValueError("Model not loaded. Call load_context first.")
# Extract data from input
if isinstance(model_input, dict):
X = model_input.get("data")
if X is None:
raise ValueError("Dict input must contain 'data' key.")
else:
X = model_input
# Ensure DataFrame input (models expect pandas DataFrames)
if not isinstance(X, pd.DataFrame):
raise ValueError("Input must be a pandas DataFrame.")
try:
predictions = self.model.predict(X)
return predictions.to_frame(name=self.model.target_col)
except Exception as e:
self._console.print(f"[red]Prediction failed: {e}[/red]")
raise
def get_model_summary(self) -> str:
"""Return human-readable model summary."""
return self.model.summary() if self.model else "No model loaded"
def store_model(
self,
path: Optional[str] = None,
artifact_path: str = "stacking_model",
signature: Optional[ModelSignature] = None,
pip_requirements: Optional[Union[str, list]] = None,
code_path: Optional[List[str]] = None,
to_disk: bool = False,
) -> None:
"""
Store the model using MLflow pyfunc interface.
Logs to the current MLflow run by default. Optionally saves locally.
Args:
path: Local path to save model (required if to_disk=True)
artifact_path: MLflow artifact path
signature: Optional MLflow model signature
pip_requirements: pip requirements (list or path)
code_path: List of local Python source files/directories to bundle
to_disk: Save locally if True, otherwise logs to MLflow
"""
if self.model is None:
raise ValueError("No model to store.")
with tempfile.TemporaryDirectory() as tmp:
model_artifact = os.path.join(tmp, "stacking_model.pkl")
self.model.save(model_artifact, compression="lzma")
common_args = {
"python_model": self,
"artifacts": {"model": model_artifact},
}
if signature:
common_args["signature"] = signature
if pip_requirements:
common_args["pip_requirements"] = pip_requirements
if code_path:
common_args["code_path"] = code_path
if to_disk:
if not path:
raise ValueError("`path` required for to_disk=True.")
mlflow.pyfunc.save_model(path=path, **common_args)
self._console.log(
f"[blue]Model saved locally to {path}[/blue]"
)
else:
mlflow.pyfunc.log_model(
artifact_path=artifact_path, **common_args
)
self._console.log(
f"[green]Model logged to MLflow at '{artifact_path}'"
)
def __getstate__(self):
state = self.__dict__.copy()
state["model"] = None # avoid double saving
return state
def __setstate__(self, state):
self.__dict__.update(state)
class TransformerWrapper(mlflow.pyfunc.PythonModel): # type: ignore
"""
MLflow wrapper for data transformers.
Allows transformers to be saved and served via MLflow's pyfunc interface.
Supports both training and inference transformers.
"""
def __init__(
self,
transformer: Optional[CourierTrainingTransformer] = None,
):
"""
Initialize the transformer wrapper.
Args:
transformer: The training transformer to wrap
"""
self.training_transformer = transformer
self.inference_transformer = None
@property
def _console(self):
from rich.console import Console
return Console()
def load_context(self, context: Any) -> None:
"""
Load training transformer from artifact path and create
inference transformer.
"""
try:
transformer_path = context.artifacts["transformer"]
with open(transformer_path, "rb") as f:
self.training_transformer = pickle.load(f)
_, self.inference_transformer = (
create_transformers_from_training_transformer(
self.training_transformer
)
)
self._console.log(
"[green]Training transformer loaded and inference transformer "
+ "created from context[/green]"
)
except Exception as e:
self._console.print(f"[red]Error loading transformer: {e}[/red]")
raise
def predict(
self,
context: Any,
model_input: Union[pd.DataFrame, np.ndarray, Dict[str, Any]],
transformer_type: str = "inference",
) -> Union[pd.Series, pd.DataFrame, np.ndarray]:
"""
Transform data using the selected transformer.
Args:
context: MLflow context
model_input: Input data to transform (pandas DataFrame expected)
transformer_type: Either "training" or "inference"
"""
if transformer_type == "training":
transformer = self.training_transformer
elif transformer_type == "inference":
transformer = self.inference_transformer
else:
raise ValueError(
"transformer_type must be 'training' or 'inference'"
)
if transformer is None:
raise ValueError(
f"{transformer_type.title()} transformer not loaded. "
+ "Call load_context first."
)
# Extract data from input
if isinstance(model_input, dict):
X = model_input.get("data")
if X is None:
raise ValueError("Dict input must contain 'data' key.")
else:
X = model_input
# Ensure DataFrame input (transformers expect pandas DataFrames)
if not isinstance(X, pd.DataFrame):
raise ValueError("Input must be a pandas DataFrame.")
try:
# Apply transformer
transformed_data = transformer.transform(X)
return transformed_data
except Exception as e:
self._console.print(f"[red]Transformation failed: {e}[/red]")
raise
def get_transformer_summary(self) -> str:
"""Return human-readable transformer summary."""
if self.training_transformer is None:
return "No training transformer loaded"
training_class = self.training_transformer.__class__.__name__
inference_status = (
"available" if self.inference_transformer else "not created"
)
return (
f"{training_class} (training loaded, inference {inference_status})"
)
def store_transformer(
self,
path: Optional[str] = None,
artifact_path: str = "transformer",
signature: Optional[ModelSignature] = None,
pip_requirements: Optional[Union[str, list]] = None,
code_path: Optional[List[str]] = None,
to_disk: bool = False,
) -> None:
"""
Store the training transformer using MLflow pyfunc interface.
Logs to the current MLflow run by default. Optionally saves locally.
Args:
path: Local path to save transformer (required if to_disk=True)
artifact_path: MLflow artifact path
signature: Optional MLflow model signature
pip_requirements: pip requirements (list or path)
code_path: List of local Python source files/directories to bundle
to_disk: Save locally if True, otherwise logs to MLflow
"""
if self.training_transformer is None:
raise ValueError("No training transformer to store.")
with tempfile.TemporaryDirectory() as tmp:
transformer_artifact = os.path.join(
tmp, "training_transformer.pkl"
)
with open(transformer_artifact, "wb") as f:
pickle.dump(self.training_transformer, f)
common_args = {
"python_model": self,
"artifacts": {"transformer": transformer_artifact},
}
if signature:
common_args["signature"] = signature
if pip_requirements:
common_args["pip_requirements"] = pip_requirements
if code_path:
common_args["code_path"] = code_path
if to_disk:
if not path:
raise ValueError("`path` required for to_disk=True.")
mlflow.pyfunc.save_model(path=path, **common_args)
self._console.log(
"[blue]Training transformer saved locally to "
+ f"{path}[/blue]"
)
else:
mlflow.pyfunc.log_model(
artifact_path=artifact_path, **common_args
)
self._console.log(
"[green]Training transformer logged to MLflow at "
+ f"'{artifact_path}'[/green]"
)
def __getstate__(self):
state = self.__dict__.copy()
state["training_transformer"] = None # avoid double saving
state["inference_transformer"] = None # avoid double saving
return state
def __setstate__(self, state):
self.__dict__.update(state)

View File

@@ -0,0 +1,24 @@
"""
Models package for time series forecasting.
This package provides standardized interfaces and implementations for
various time series forecasting models.
"""
from .base import (
TimeSeriesModel,
UnivariateTimeSeriesModel,
MultivariateTimeSeriesModel,
)
from .factory import create_model, load_model, get_available_models
from .evaluation import timeseries_metrics
__all__ = [
"TimeSeriesModel",
"UnivariateTimeSeriesModel",
"MultivariateTimeSeriesModel",
"create_model",
"load_model",
"get_available_models",
"timeseries_metrics",
]

View File

@@ -0,0 +1,389 @@
"""
ARIMA univariate time series forecasting model implementation.
"""
from typing import Optional, Tuple
import numpy as np
import pandas as pd
from statsmodels.tsa.arima.model import ARIMA, ARIMAResults
from rich.console import Console
from .base import UnivariateTimeSeriesModel, ensure_fitted
console = Console()
class ARIMAModel(UnivariateTimeSeriesModel):
"""ARIMA model for univariate time series forecasting.
This class implements an ARIMA model for forecasting univariate time
series data. It provides methods for fitting the model, making predictions,
forecasting future values, and updating the model with new data.
Attributes:
order (Tuple[int, int, int]): The (p, d, q) order of the ARIMA model.
model_ (Optional[ARIMA]): The ARIMA model instance.
result_ (Optional[ARIMAResults]): The fitted ARIMA model results.
training_series_ (Optional[pd.Series]): The training data used to fit
the model.
"""
def __init__(
self,
order: Tuple[int, int, int] = (1, 0, 0),
name: Optional[str] = None,
time_col: str = "ds",
target_col: str = "y",
random_seed: int = 42,
forecast_horizon: int = 2,
) -> None:
"""Initializes the ARIMAModel with specified parameters.
Args:
order (Tuple[int, int, int]): The (p, d, q) order of the ARIMA
model.
name (Optional[str]): The name of the model.
time_col (str): The name of the time column in the input data.
target_col (str): The name of the target column in the input data.
random_seed (int): The random seed for reproducibility.
forecast_horizon (int): The number of steps to forecast ahead.
"""
super().__init__(
name=name,
time_col=time_col,
target_col=target_col,
random_seed=random_seed,
)
self.order: Tuple[int, int, int] = order
self.model_: Optional[ARIMA] = None
self.result_: Optional[ARIMAResults] = None
self.training_series_: pd.Series = pd.Series(dtype=float)
self.observed_series_: pd.Series = pd.Series(dtype=float)
self.backtest_predictions_: Optional[pd.Series] = None
self.forecast_horizon: int = forecast_horizon
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:
"""
Fits the ARIMA model to the provided training data.
Args:
y: The target time series data.
X: Optional exogenous variables.
X_val: Validation feature matrix (not used for ARIMA).
y_val: Validation target series (not used for ARIMA).
"""
y_array: np.ndarray = self._validate_y(y)
self.training_series_ = y.copy()
self.observed_series_ = y.copy()
self.model_ = ARIMA(y_array, order=self.order)
self.result_ = self.model_.fit()
@ensure_fitted
def predict(self, X: Optional[pd.DataFrame] = None) -> pd.Series:
"""
Generates in-sample predictions from the fitted ARIMA model.
After this method is called, if X is provided, the model will be
updated with the new data, but the coefficients will not be refit.
This is useful for generating predictions on new data without
retraining the model.
Args:
X (Optional[pd.DataFrame]): Optional dataframe with future
measurements of y for in-sample predictions.
If None, the model will predict on the observed data
(observed_series_).
Returns:
pd.Series: The in-sample predictions.
Raises:
ValueError: If the model has not been fitted yet.
"""
if self.result_ is None:
raise ValueError("Model is not fitted.")
if X is None:
fitted_values = self.result_.fittedvalues
if fitted_values is None:
raise ValueError("Fitted values are None")
return pd.Series(
fitted_values,
index=self.training_series_.index[: len(fitted_values)],
name=self.target_col,
)
# Validate the input data
target_series = (
X[self.target_col] if self.target_col in X else X.iloc[:, 0]
)
if not isinstance(target_series, pd.Series):
target_series = pd.Series(target_series, index=X.index)
X_validated = self._validate_y(target_series)
# Update the model with the validated data without refitting
self.update(pd.Series(X_validated, index=X.index), refit=False)
if self.result_ is None:
raise ValueError("Model result is None after update")
fitted_values = self.result_.fittedvalues
if fitted_values is None:
raise ValueError("Fitted values are None")
return_series = pd.Series(
fitted_values[-len(X) :], index=X.index, name=self.target_col
)
return return_series
@ensure_fitted
def forecast(self, forecast_horizon: int) -> np.ndarray:
"""Generates out-of-sample forecasts from the fitted ARIMA model.
TODO: Change return to include index of the forecasted values.
Args:
forecast_horizon (int): The number of steps to forecast ahead.
Returns:
np.ndarray: The out-of-sample forecasts.
Raises:
ValueError: If the model has not been fitted yet.
"""
if self.result_ is None:
raise ValueError("Model is not fitted.")
return self.result_.forecast(steps=forecast_horizon)
@ensure_fitted
def backtest(
self,
y: pd.Series,
X: Optional[pd.DataFrame] = None,
retrain_every: int = 50,
reuse_previous_execution: bool = False,
) -> pd.Series:
"""
Perform comprehensive backtesting with periodic retraining.
This method implements walk-forward validation with periodic
retraining, providing robust evaluation of model performance in
production-like scenarios. Uses 1-step ahead forecasting by default.
Args:
y: Target time series for backtesting
X: Unused (included for base class compatibility)
retrain_every: Number of steps between model retraining
reuse_previous_execution: Whether to reuse previous backtest results
Returns:
Series of backtested predictions indexed by timestamp
Raises:
ValueError: If parameters are invalid or data is insufficient
RuntimeError: If backtesting fails
"""
if self.result_ is None:
raise ValueError("Model is not fitted")
if self.training_series_ is None:
raise ValueError("No training series found")
# Handle reuse of previous execution
if reuse_previous_execution and self.backtest_predictions_ is not None:
expected_index = y.index
if (
len(self.backtest_predictions_) == len(expected_index)
and (self.backtest_predictions_.index == expected_index).all()
):
console.log(
"[yellow]Reusing previous backtest results[/yellow]"
)
return self.backtest_predictions_
else:
console.log(
"[yellow]Previous results incompatible, running new"
+ " backtest[/yellow]"
)
try:
console.log(
f"[blue]Starting ARIMA backtest with {self.forecast_horizon}"
+ "-step forecasting...[/blue]"
)
# Validate and prepare data
y_sorted = y.sort_index()
# Check for overlapping data
training_series = self.training_series_
if any(t in training_series.index for t in y_sorted.index):
console.print(
"[yellow]Warning: Backtest data overlaps with training"
+ " data[/yellow]"
)
# Initialize backtesting
predictions = []
# Start with training data
current_series = training_series.copy()
total_steps = len(y_sorted)
console.log(
f"[blue]Running {total_steps} backtest steps with retraining"
+ f" every {retrain_every} steps...[/blue]"
)
# Create initial model state
current_model = ARIMA(current_series.values, order=self.order)
current_result = current_model.fit()
# Perform walk-forward validation
for i, (timestamp, actual_value) in enumerate(y_sorted.items()):
if i % 50 == 0 and i > 0: # Progress logging
console.log(
f"[blue]Backtest progress: {i}/{len(y_sorted)}[/blue]"
)
try:
# Check if we need to retrain
if i % retrain_every == 0 and i > 0:
console.log(
f"[blue]Retraining model at step {i}[/blue]"
)
current_model = ARIMA(
current_series.values, order=self.order
)
current_result = current_model.fit()
# Generate forecast_horizon-step ahead forecast
forecast = current_result.forecast(
steps=self.forecast_horizon
)[0]
predictions.append((timestamp, forecast))
# Update the series with actual observed value
current_series = pd.concat(
[
current_series,
pd.Series([actual_value], index=[timestamp]),
]
)
# For ARIMA, we can extend the model without full refit
if i % retrain_every != 0:
try:
current_result = current_result.extend(
[actual_value], refit=False
)
except Exception:
# If extend fails, do a quick refit
current_model = ARIMA(
current_series.values, order=self.order
)
current_result = current_model.fit()
except Exception as step_error:
console.print(
f"[yellow]Error at step {i}: {step_error}, using"
+ " NaN[/yellow]"
)
predictions.append((timestamp, np.nan))
# Still update the series for continuity
current_series = pd.concat(
[
current_series,
pd.Series([actual_value], index=[timestamp]),
]
)
# Create results series
if predictions:
pred_index, pred_values = zip(*predictions)
self.backtest_predictions_ = pd.Series(
pred_values,
index=pd.Index(pred_index),
name=f"{self.target_col}_backtest",
)
else:
self.backtest_predictions_ = pd.Series(
dtype=float, name=f"{self.target_col}_backtest"
)
console.log(
"[green]ARIMA backtest completed: "
+ f"{len(self.backtest_predictions_)} predictions[/green]"
)
return self.backtest_predictions_
except Exception as e:
console.print(f"[red]ARIMA backtest failed: {e}[/red]")
raise RuntimeError(f"Failed to perform backtest: {e}") from e
@ensure_fitted
def summary(self) -> str:
"""Generates a summary of the fitted ARIMA model.
Returns:
str: The summary of the fitted model.
Raises:
ValueError: If the model has not been fitted yet.
"""
if self.result_ is None:
raise ValueError("Model is not fitted.")
return str(self.result_.summary())
@ensure_fitted
def update(self, new_data: pd.Series, refit: bool = True) -> None:
"""Updates the ARIMA model with new observed data.
This method allows for two modes of updating the model:
1. **Refitting**: The model is retrained on the combined dataset
(original training data + new data).
2. **Incremental Update**: The model is updated using the new data
without retraining, preserving the original model parameters.
Args:
new_data (pd.Series): New observed values to update the model with.
refit (bool, optional): If True, the model is retrained on the
combined dataset. Defaults to True.
Raises:
TypeError: If `new_data` is not a pandas Series.
ValueError: If `new_data` is empty or if the model has not been
fitted yet.
"""
if not isinstance(new_data, pd.Series):
raise TypeError("new_data must be a pandas Series.")
if new_data.empty:
raise ValueError("new_data is empty.")
if refit:
self.training_series_ = pd.concat(
[self.training_series_, new_data]
)
self.observed_series_ = self.training_series_.copy()
y_array = self._validate_y(self.training_series_)
self.model_ = ARIMA(y_array, order=self.order)
self.result_ = self.model_.fit()
else:
self.observed_series_ = pd.concat(
[self.observed_series_, new_data]
)
y_array = self._validate_y(self.observed_series_)
if self.result_ is None:
raise ValueError("Model is not fitted.")
self.result_ = self.result_.apply(y_array, refit=False)
if self.result_ is not None:
self.model_ = self.result_.model

View File

@@ -0,0 +1,824 @@
"""
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)

View File

@@ -0,0 +1,510 @@
"""
CatBoost implementation for multivariate time series forecasting.
TODO:
- Implement support for categorical features.
"""
from typing import Optional, List, Union, Tuple, Dict, Any, cast
import pandas as pd
from catboost import CatBoostRegressor, CatBoostClassifier, Pool
from rich.console import Console
import numpy as np
from .base import MultivariateTimeSeriesModel, ensure_fitted
console = Console()
class CatBoostTimeSeriesModel(MultivariateTimeSeriesModel):
"""
CatBoost implementation for multivariate time series forecasting.
This class wraps the CatBoost models with additional functionality for
time series forecasting, following the MultivariateTimeSeriesModel
interface. Supports both regression and classification tasks with
comprehensive error handling and type safety.
"""
def __init__(
self,
name: Optional[str] = None,
learning_task: str = "regression",
differentiate_target: bool = False,
n_lags: int = 0,
iterations: int = 1000,
learning_rate: float = 0.1,
depth: int = 6,
loss_function: Optional[str] = None,
bins: Optional[List[float]] = None,
random_seed: int = 42,
time_col: str = "ds",
target_col: str = "y",
verbose: bool = False,
) -> None:
"""
Initialize the CatBoost time series model.
Args:
name: Optional identifier for the model
learning_task: Type of learning task, either 'regression',
'multiclass' or 'binary'.
differentiate_target: Whether to differentiate the target
series before fitting the model.
n_lags: Number of lagged target values included as features.
These lags are expected to already be present in the same
dataset as the exogenous features.
iterations: Number of boosting iterations
learning_rate: Learning rate for the model
depth: Depth of the tree
loss_function: Loss function to optimize
bins: Optional list of bin edges for multiclass classification
random_seed: Random seed for reproducibility
time_col: Name of the time column
target_col: Name of the target column
verbose: Whether to enable verbose output
"""
super().__init__(
name=name,
time_col=time_col,
target_col=target_col,
random_seed=random_seed,
n_lags=n_lags,
learning_task=learning_task,
differentiate_target=differentiate_target,
bins=bins,
)
self.iterations = iterations
self.learning_rate = learning_rate
self.depth = depth
self.verbose = verbose
# Set default loss function based on learning task
self.loss_function = self._get_default_loss_function(loss_function)
# Initialize model state
self.model_: Optional[Union[CatBoostRegressor, CatBoostClassifier]] = (
None
)
self.training_series_: Optional[pd.Series] = None
self.X_train_: Optional[pd.DataFrame] = None
self.backtest_predictions_: Optional[pd.Series] = None
if self.verbose:
console.log(
"[green]Initialized CatBoostTimeSeriesModel:"
+ f" {self.summary()}[/green]"
)
def _create_model(self) -> Union[CatBoostRegressor, CatBoostClassifier]:
"""Creates a new instance of CatBoost model with current parameters.
Returns:
A new CatBoost model instance (Regressor or Classifier).
"""
base_params = {
"iterations": self.iterations,
"learning_rate": self.learning_rate,
"depth": self.depth,
"loss_function": self.loss_function,
"random_seed": self.random_seed,
"verbose": self.verbose,
}
if self.learning_task in ["binary", "multiclass"]:
return CatBoostClassifier(
auto_class_weights="Balanced", **base_params
)
else:
return CatBoostRegressor(**base_params)
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 for CatBoost model with optional validation data.
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)
"""
if X is None or not isinstance(X, pd.DataFrame):
raise ValueError("Feature matrix X must be a non-empty DataFrame.")
y_processed, X_processed, y_val_processed, X_val_processed = (
self._preprocess_data(y, X, X_val, y_val)
)
# Ensure X is not None after preprocessing
if X_processed is None:
raise ValueError(
"Feature matrix X cannot be None after preprocessing."
)
X_array, y_array = self._validate_X_y(X_processed, y_processed)
self.training_series_ = y_processed.copy()
self.X_train_ = X_processed.copy()
self.model_ = self._create_model()
eval_set = None
if X_val_processed is not None and y_val_processed is not None:
X_val_array, y_val_array = self._validate_X_y(
X_val_processed, y_val_processed
)
eval_set = Pool(data=X_val_array, label=y_val_array)
train_pool = Pool(data=X_array, label=y_array)
if self.verbose:
console.log(
f"[blue]Training CatBoost model for {self.iterations}"
+ " iterations...[/blue]"
)
self.model_.fit(train_pool, eval_set=eval_set)
if self.verbose:
console.log(
"[green]CatBoost model training completed successfully[/green]"
)
@ensure_fitted
def predict(self, X: pd.DataFrame) -> pd.Series:
"""
Generate predictions using the fitted CatBoost model.
Args:
X: The feature matrix for prediction
Returns:
pd.Series: Predicted values
"""
if self.model_ is None:
raise ValueError("Model is not fitted yet.")
# Create a copy to avoid modifying the original DataFrame
X_pred = X.copy()
if self.selected_features_ is not None:
X_pred = cast(pd.DataFrame, X_pred[self.selected_features_])
X_array = self._validate_X(X_pred)
predictions = self.model_.predict(X_array)
# Convert predictions to numpy array if needed
if hasattr(predictions, "squeeze"):
predictions = predictions.squeeze()
elif isinstance(predictions, list):
predictions = np.array(predictions)
return_series = pd.Series(
predictions,
index=X.index,
name=self.target_col,
)
return return_series
@ensure_fitted
def backtest(
self,
y: pd.Series,
X: Optional[pd.DataFrame] = None,
retrain_every: int = 50,
reuse_previous_execution: bool = False,
) -> pd.Series:
"""
Performs backtesting (walk-forward validation) with periodic
retraining.
This method simulates a production scenario by iterating through a test
set, making a one-step-ahead prediction, and then retraining the model
periodically with the newly available data.
Args:
X: DataFrame with features for the backtesting period.
y: Series with the true target values for the backtesting period.
retrain_every: The frequency of retraining. The model will be
retrained every `retrain_every` steps.
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:
A series of backtested predictions, indexed by the backtest data's
index.
"""
if self.model_ is None:
raise ValueError("Model is not fitted yet.")
if self.training_series_ is None:
raise ValueError("Training series is not set.")
if self.X_train_ is None:
raise ValueError("Training feature matrix is not set.")
if X is None or not isinstance(X, pd.DataFrame):
raise ValueError("Feature matrix X must be a non-empty DataFrame.")
if reuse_previous_execution:
if self.backtest_predictions_ is None:
raise ValueError("No previous execution found.")
if (self.backtest_predictions_.shape[0] != y.shape[0]) or (
not (self.backtest_predictions_.index == y.index).all()
):
raise ValueError(
"Previous execution index does not match y index."
)
return self.backtest_predictions_
# Prepare
total_steps = len(X)
predictions = []
current_model = self.model_
y_history = self.training_series_.copy()
X_history = self.X_train_.copy()
# Iterate in chunks instead of single steps
for start in range(0, total_steps, retrain_every):
end = min(start + retrain_every, total_steps)
# Batch prediction for current chunk
X_chunk = X.iloc[start:end].copy()
if self.selected_features_:
X_chunk = X_chunk[self.selected_features_]
X_array = self._validate_X(X_chunk)
preds = current_model.predict(X_array)
# Handle different prediction formats
if hasattr(preds, "squeeze"):
preds = preds.squeeze()
if preds.ndim == 0: # single point
preds = [preds]
predictions.extend(preds)
# Update training history
y_chunk = y.iloc[start:end]
y_history = pd.concat([y_history, y_chunk])
X_history = pd.concat([X_history, X_chunk])
# Retrain the model for next chunk (if needed)
if end < total_steps:
if self.verbose:
console.print(
f"[cyan]Backtesting: Retraining at step {end}..."
)
current_model = self._create_model()
(y_fit, X_fit, _, _) = (
self._preprocess_data(y_history, X_history)
)
# Ensure X_fit is not None after preprocessing
if X_fit is None:
raise ValueError(
"Feature matrix cannot be None after preprocessing."
)
X_fit_array, y_fit_array = self._validate_X_y(X_fit, y_fit)
train_pool = Pool(data=X_fit_array, label=y_fit_array)
current_model.fit(train_pool)
# Store backtest predictions for potential reuse
self.backtest_predictions_ = pd.Series(
predictions, index=X.index, name=f"{self.target_col}_pred"
)
return self.backtest_predictions_
def select_features(
self,
X: pd.DataFrame,
y: pd.Series,
X_val: Optional[pd.DataFrame] = None,
y_val: Optional[pd.Series] = None,
features_to_select: Optional[int] = None,
algorithm: str = "RecursiveByShapValues",
steps: int = 1,
verbose: bool = False,
) -> List[str]:
"""Identify and select the most important features.
Uses CatBoost's built-in feature selection capabilities to determine
feature importance and select the most relevant features.
Args:
X: The feature matrix
y: The target series
X_val: Optional validation feature matrix
y_val: Optional validation target series
features_to_select: Number of features to select. If None,
will select half of the features.
algorithm: Feature selection algorithm. One of:
'RecursiveByShapValues', 'RecursiveByPredictionValuesChange'
steps: How many times a full model will be trained.
More steps give more accurate results.
verbose: Whether to print progress
Returns:
List[str]: List of selected feature names
"""
(y_processed, X_processed, y_val_processed, X_val_processed) = (
self._preprocess_data(y, X, X_val, y_val)
)
# Validate input data
if X_processed is None or not isinstance(X_processed, pd.DataFrame):
raise ValueError("Feature matrix X must be a non-empty DataFrame.")
X_array, y_array = self._validate_X_y(X_processed, y_processed)
# Set default number of features to select if not specified
if features_to_select is None:
features_to_select = X_processed.shape[1] // 2
# Create and prepare model
temp_model = self._create_model()
train_pool = Pool(data=X_array, label=y_array)
# Prepare validation data if provided
eval_set = None
if X_val_processed is not None and y_val_processed is not None:
X_val_array, y_val_array = self._validate_X_y(
X_val_processed, y_val_processed
)
eval_set = Pool(data=X_val_array, label=y_val_array)
# Perform feature selection
selected_features = temp_model.select_features(
train_pool,
eval_set=eval_set,
features_for_select=list(range(X_processed.shape[1])),
num_features_to_select=features_to_select,
algorithm=algorithm,
steps=steps,
logging_level="Verbose" if verbose else "Silent",
train_final_model=False,
)
# Map feature indices to feature names with proper type casting
selected_feature_names: List[str] = [
str(X_processed.columns[idx])
for idx in selected_features["selected_features"]
]
self.selected_features_ = selected_feature_names
self.feature_names_in_ = selected_feature_names
self.n_features_in_ = len(selected_feature_names)
return selected_feature_names
@classmethod
def tune_hyperparameters(
cls,
y: pd.Series,
X: Optional[pd.DataFrame] = None,
X_val: Optional[pd.DataFrame] = None,
y_val: Optional[pd.Series] = None,
selected_features: Optional[List[str]] = None,
param_grid: Optional[Dict[str, Any]] = None,
n_trials: int = 10,
early_stopping_rounds: Optional[int] = 50,
random_seed: int = 42,
**kwargs,
) -> Tuple[Dict[str, Any], "CatBoostTimeSeriesModel"]:
"""
Tune hyperparameters for the CatBoost model.
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)
selected_features: List of features to use for tuning
param_grid: Dictionary of hyperparameters to search
n_trials: Number of trials for hyperparameter tuning
early_stopping_rounds: Number of rounds for early stopping
random_seed: Random seed for reproducibility
**kwargs: Additional keyword arguments for model initialization
Returns:
Tuple[Dict[str, Any], CatBoostTimeSeriesModel]: Best hyperparameters
and fitted model
"""
if n_trials <= 0:
raise ValueError("n_trials must be a positive integer.")
# Create a temporary model instance to use its preprocessing method
temp_model = cls(
learning_task=kwargs.get("learning_task", "regression"),
differentiate_target=kwargs.get("differentiate_target", False),
bins=kwargs.get("bins", None),
random_seed=random_seed,
)
if selected_features is not None:
temp_model.selected_features_ = selected_features
(y_processed, X_processed, _, _) = (
temp_model._preprocess_data(y, X, X_val, y_val)
)
if kwargs.get("learning_task", "regression") == "classification":
search_model = CatBoostClassifier(
random_seed=random_seed,
logging_level="Silent",
early_stopping_rounds=early_stopping_rounds,
loss_function=kwargs.get("loss_function", "Logloss"),
class_weights="Balanced",
)
else:
search_model = CatBoostRegressor(
random_seed=random_seed,
logging_level="Silent",
early_stopping_rounds=early_stopping_rounds,
loss_function=kwargs.get("loss_function", "RMSE"),
)
if param_grid is None:
param_grid = {
"iterations": [100, 500, 1000, 2000],
"learning_rate": [0.01, 0.05, 0.1, 0.2],
"depth": [4, 6, 8],
}
if X_processed is None:
raise ValueError(
"Feature matrix cannot be None after preprocessing."
)
train_pool = Pool(data=X_processed, label=y_processed)
results = search_model.randomized_search(
param_grid,
X=train_pool,
n_iter=n_trials,
verbose=False,
refit=False,
)
best_params = results["params"]
best_model = cls(
iterations=best_params["iterations"],
learning_rate=best_params["learning_rate"],
depth=best_params["depth"],
random_seed=random_seed,
time_col=kwargs.get("time_col", "ds"),
target_col=kwargs.get("target_col", "y"),
n_lags=kwargs.get("n_lags", 0),
name=kwargs.get("name", None),
loss_function=kwargs.get("loss_function", None),
learning_task=kwargs.get("learning_task", "regression"),
bins=kwargs.get("bins", None),
differentiate_target=kwargs.get("differentiate_target", False),
)
if selected_features is not None:
best_model.selected_features_ = selected_features
best_model.fit(y, X, X_val, y_val)
return best_params, best_model

View File

@@ -0,0 +1,92 @@
"""
Module with functions of timeseries evaluation.
"""
from typing import Optional, List, Sequence
import numpy as np
import pandas as pd
from rich.console import Console
from sklearn.metrics import (
mean_absolute_error,
mean_squared_error,
accuracy_score,
f1_score,
confusion_matrix,
)
console = Console()
def timeseries_metrics(
y_pred: Sequence[float], y_true: Sequence[float]
) -> dict[str, float]:
"""
Compute MAE, MSE, and trend capture for time series predictions.
Parameters:
y_pred (ArrayLike): Predicted values.
y_true (ArrayLike): Ground truth values.
Returns:
dict[str, float]: Dictionary with MAE, MSE, and trend_capture.
"""
if len(y_true) != len(y_pred):
raise ValueError("y_true and y_pred must have the same length")
if len(y_true) == 0:
raise ValueError("y_true and y_pred must not be empty")
y_true_np = np.asarray(y_true)
y_pred_np = np.asarray(y_pred)
mae = mean_absolute_error(y_true_np, y_pred_np)
mse = mean_squared_error(y_true_np, y_pred_np)
# Compute directional trend: 1 if up, 0 if down or flat
if len(y_true) == 1:
return {"MAE": mae, "MSE": mse, "trend_capture": 1.0}
true_trend = np.diff(y_true_np) > 0
pred_trend = np.diff(y_pred_np) > 0
trend_capture = np.mean(true_trend == pred_trend)
return {"MAE": mae, "MSE": mse, "trend_capture": trend_capture}
def timeseries_classification_metrics(
y_pred: Sequence[float],
y_true: Sequence[float],
bins: Optional[List] = None,
) -> dict[str, float]:
"""
Compute accuracy for classification predictions.
Parameters:
y_pred (ArrayLike): Predicted values.
y_true (ArrayLike): Ground truth values.
bins (List[int], optional): Bin edges for categorizing predictions.
Returns:
dict[str, float]: Dictionary with accuracy.
"""
if bins is not None:
console.log(
f"Using bins for classification: {bins}"
)
y_true_binned = pd.cut(y_true, bins=bins, labels=False)
else:
console.log("No bins provided, using default classification (y > 0).")
y_true_binned = (y_true > 0).astype(int)
console.log(
f"y_true_binned: {y_true_binned.value_counts()}"
)
console.log(
f"y_pred: {pd.Series(y_pred).value_counts()}"
)
y_pred = np.array(y_pred).astype(int)
acc = accuracy_score(y_true_binned, y_pred)
f1 = f1_score(y_true_binned, y_pred, average="weighted")
# Calculate the confusion matrix
cm = confusion_matrix(y_true_binned, y_pred)
return {"accuracy": acc, "f1_score": f1, "confusion_matrix": cm}

View File

@@ -0,0 +1,140 @@
"""
Model factory for creating, loading, and discovering time series models.
This module provides centralized utility functions to handle different time
series model implementations based on a string identifier. It uses a
central registry (`SUPPORTED_MODELS`) that maps model type strings (e.g.,
'arima') to their corresponding model classes (e.g., ARIMAModel). This
approach allows for easy extension and decouples model instantiation logic
from the code that uses the models.
Key Functions:
create_model: Creates a new instance of a specified model type by looking
up the type string in the `SUPPORTED_MODELS` registry and
passing keyword arguments to the retrieved model class's
constructor.
load_model: Loads a previously saved model instance from disk. It uses
the provided model type string to find the correct class in
the registry and then calls that class's `.load()` classmethod.
get_available_models: Returns a dictionary listing the registered model
types (keys in `SUPPORTED_MODELS`) and their
descriptions, automatically derived from the model
class docstrings.
Extensibility:
Adding support for a new model involves the following steps:
1. Ensure the new model class (e.g., `MyNewModel`) inherits from the
appropriate base class (e.g., `TimeSeriesModel`) and implements all
required abstract methods.
2. Ensure the new model class has a `.load()` classmethod compatible
with the `save()` method in the base `Model` class (if loading is
to be supported via this factory).
3. Import the new model class into this factory module.
4. Add an entry to the `SUPPORTED_MODELS` dictionary, mapping a unique,
lowercase string identifier to the model class itself:
`SUPPORTED_MODELS = {..., "mynewmodel": MyNewModel}`
Once added to the registry, the model can be created and loaded via the
factory functions, and it will automatically appear in the output of
`get_available_models()`.
"""
from typing import Dict, Type
import logging
from .base import TimeSeriesModel
from .arima import ARIMAModel
from .neural_prophet_model import NeuralProphetModel
from .catboost_time_series import CatBoostTimeSeriesModel
from .linear_regression_time_series import ElasticNetTimeSeriesModel
from .stacking_time_series import StackingTimeSeriesModel
# from .prophet import ProphetModel # Example for future
# *** Central registry of supported models
SUPPORTED_MODELS: Dict[str, Type[TimeSeriesModel]] = {
"arima": ARIMAModel,
"neuralprophet": NeuralProphetModel,
"catboost": CatBoostTimeSeriesModel,
"elasticnet": ElasticNetTimeSeriesModel,
"stacking": StackingTimeSeriesModel,
# "prophet": ProphetModel, # Add new models here
}
def create_model(model_type: str, **kwargs) -> TimeSeriesModel:
"""
Create a new model instance of the specified type using a registry.
Args:
model_type: Type of model to create (case-insensitive).
**kwargs: Model-specific parameters passed to its constructor.
Returns:
New model instance inheriting from TimeSeriesModel.
Raises:
ValueError: If the model type is not supported or kwargs are invalid.
"""
model_type = model_type.lower()
model_class = SUPPORTED_MODELS.get(model_type)
if model_class:
try:
instance = model_class(**kwargs)
return instance
except TypeError as e:
logging.error(f"Kwargs issue for {model_type}: {kwargs}")
raise ValueError(
f"Invalid parameters for model type '{model_type}'. Error: {e}"
) from e
else:
supported_list = ", ".join(f"'{k}'" for k in SUPPORTED_MODELS.keys())
raise ValueError(
f"Unsupported model type: '{model_type}'. "
f"Currently supported models are: {supported_list}."
)
def load_model(path: str, model_type: str) -> TimeSeriesModel:
"""
Load a model from disk using a registry.
Args:
path: Path to the saved model.
model_type: Expected type of model to load (case-insensitive).
Returns:
Loaded model instance.
Raises:
ValueError: If the model type is not supported.
# Other errors might come from the underlying .load() method
"""
model_type = model_type.lower()
model_class = SUPPORTED_MODELS.get(model_type)
if model_class:
return model_class.load(path)
else:
supported_list = ", ".join(f"'{k}'" for k in SUPPORTED_MODELS.keys())
raise ValueError(
f"Unsupported model type: '{model_type}'. "
f"Currently supported models are: {supported_list}."
)
def get_available_models() -> Dict[str, str]:
"""
Dynamically get a dictionary of available model types and their
descriptions from the SUPPORTED_MODELS registry and class docstrings.
Returns:
Dictionary mapping model type names to descriptions.
"""
available = {
type_name: (
model_class.__doc__.strip().splitlines()[0]
if model_class.__doc__ else "No description available."
)
for type_name, model_class in SUPPORTED_MODELS.items()
}
return available

View File

@@ -0,0 +1,303 @@
"""
ElasticNet implementation for multivariate time series forecasting.
"""
from typing import Optional, Sequence
import pandas as pd
from sklearn.linear_model import ElasticNet
from sklearn.impute import SimpleImputer
from rich.console import Console
from .base import MultivariateTimeSeriesModel, ensure_fitted
console = Console()
class ElasticNetTimeSeriesModel(MultivariateTimeSeriesModel):
"""
ElasticNet implementation for multivariate time series forecasting.
This class wraps the scikit-learn ElasticNet model with additional
functionality for time series forecasting, following the
MultivariateTimeSeriesModel interface. It's suitable for regression
tasks where features might be correlated.
"""
def __init__(
self,
name: Optional[str] = None,
n_lags: int = 1,
alpha: float = 1.0,
l1_ratio: float = 0.5,
fit_intercept: bool = True,
max_iter: int = 1000,
tol: float = 1e-4,
random_seed: int = 42,
time_col: str = "ds",
target_col: str = "y",
differentiate_target: bool = False,
bins: Optional[list] = None,
learning_task: Optional[str] = None,
) -> None:
"""
Initialize the ElasticNet time series model.
Args:
name: Optional identifier for the model.
n_lags: Number of lagged target values to include as inputs.
alpha: Constant that multiplies the penalty terms.
l1_ratio: The ElasticNet mixing parameter (0 <= l1_ratio <= 1).
For l1_ratio = 0, it's L2 penalty (Ridge).
For l1_ratio = 1, it's L1 penalty (Lasso).
fit_intercept: Whether to calculate the intercept for this model.
max_iter: Maximum number of iterations.
tol: Tolerance for stopping criteria.
random_seed: Random seed for reproducibility.
time_col: Name of the time column.
target_col: Name of the target column.
differentiate_target: Whether to apply differencing to make series
stationary.
bins: Bin edges for multiclass classification target
transformation.
learning_task: Type of learning task ('regression', 'binary',
'multiclass').
"""
super().__init__(
name=name,
time_col=time_col,
target_col=target_col,
random_seed=random_seed,
n_lags=n_lags,
differentiate_target=differentiate_target,
bins=bins,
learning_task=learning_task,
)
self.alpha = alpha
self.l1_ratio = l1_ratio
self.fit_intercept = fit_intercept
self.max_iter = max_iter
self.tol = tol
self.model_: Optional[ElasticNet] = None
self.training_series_: Optional[pd.Series] = None
self.imputer_: Optional[SimpleImputer] = None
self.X_train_: Optional[pd.DataFrame] = None
def _create_model(self) -> ElasticNet:
"""Creates a new instance of ElasticNet with current parameters.
Returns:
A new ElasticNet model instance.
"""
return ElasticNet(
alpha=self.alpha,
l1_ratio=self.l1_ratio,
fit_intercept=self.fit_intercept,
max_iter=self.max_iter,
tol=self.tol,
random_state=self.random_seed,
)
def _impute_missing_values(self, X: pd.DataFrame) -> pd.DataFrame:
"""
Handle missing values in the feature matrix using median imputation.
If the imputer is not fitted, it will be fitted on the data.
Args:
X: The feature matrix potentially containing missing values.
Returns:
pd.DataFrame: The feature matrix with imputed values.
"""
if self.imputer_ is None:
self.imputer_ = SimpleImputer(
strategy="median", copy=True, add_indicator=False
)
# Fit the imputer and transform the data
imputed_values = self.imputer_.fit_transform(X)
else:
# Use the fitted imputer to transform new data
imputed_values = self.imputer_.transform(X)
# Convert back to DataFrame with original index and column names
return pd.DataFrame(imputed_values, index=X.index, columns=X.columns)
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 for the ElasticNet model.
Args:
y: The target time series data.
X: The feature matrix (including exogenous features).
X_val: Validation feature matrix (ignored).
y_val: Validation target series (ignored).
"""
# Use base class preprocessing
y_processed, X_processed, _, _ = self._preprocess_data(
y, X, X_val, y_val
)
if X_processed is None or not isinstance(X_processed, pd.DataFrame):
raise ValueError("Feature matrix X must be a non-empty DataFrame.")
# First impute missing values in X
X_imputed = self._impute_missing_values(X_processed)
# Validate X and y after imputation
X_array, y_array = self._validate_X_y(
X_imputed, y_processed, allow_nan=False
)
self.training_series_ = y_processed.copy()
self.model_ = self._create_model()
self.model_.fit(X_array, y_array)
self.X_train_ = X_processed.copy()
@ensure_fitted
def predict(self, X: Optional[pd.DataFrame] = None) -> Sequence:
"""
Generate predictions using the fitted ElasticNet model.
Args:
X: The feature matrix for prediction.
Returns:
pd.Series: Predicted values with the original index.
"""
if X is None:
raise ValueError("Feature matrix X is required for prediction.")
if self.model_ is None:
raise ValueError("Model is not fitted yet.")
if self.training_series_ is None:
raise ValueError("Training series is not available.")
if self.imputer_ is None:
raise ValueError("Imputer is not fitted yet.")
# If target column is present, drop it
X_pred = X.copy()
if self.target_col in X_pred.columns:
X_pred = X_pred.drop(columns=[self.target_col])
# For prediction, we need to reconstruct lagged features
# This is a simplified approach - in practice, you'd need
# the historical target values to create proper lags
# Handle missing values using fitted imputer
X_processed = self._impute_missing_values(X_pred)
# Validate X after imputation
X_array = self._validate_X(X_processed, allow_nan=False)
predictions = self.model_.predict(X_array)
return pd.Series(
predictions, index=X_processed.index, name=self.target_col
)
@ensure_fitted
def backtest(
self,
y: pd.Series,
X: Optional[pd.DataFrame] = None,
retrain_every: int = 50,
reuse_previous_execution: bool = False,
) -> pd.Series:
"""
Performs backtesting (walk-forward validation) with periodic
retraining.
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.
"""
if self.model_ is None:
raise ValueError("Model is not fitted yet.")
if self.training_series_ is None:
raise ValueError("Training series is not set.")
if self.X_train_ is None:
raise ValueError("Training feature matrix is not set.")
if X is None or not isinstance(X, pd.DataFrame):
raise ValueError("Feature matrix X must be a non-empty DataFrame.")
total_steps = len(X)
predictions = []
current_model = self.model_
y_history = self.training_series_.copy()
X_history = self.X_train_.copy()
# Iterate in chunks instead of single steps
for start in range(0, total_steps, retrain_every):
end = min(start + retrain_every, total_steps)
# Batch prediction for current chunk
X_chunk = X.iloc[start:end].copy()
if self.selected_features_:
X_chunk = X_chunk[self.selected_features_]
X_imputed = self._impute_missing_values(X_chunk)
X_array = self._validate_X(X_imputed, allow_nan=False)
preds = current_model.predict(X_array).squeeze()
if preds.ndim == 0: # single point
preds = [preds]
predictions.extend(preds)
# Update training history
y_chunk = y.iloc[start:end]
y_history = pd.concat([y_history, y_chunk])
X_history = pd.concat([X_history, X_chunk])
# Retrain the model for next chunk (if needed)
if end < total_steps:
console.print(
f"[cyan]Backtesting: Retraining at step {end}...[/cyan]"
)
current_model = self._create_model()
y_fit, X_fit, *_ = self._preprocess_data(y_history, X_history)
X_fit_imputed = self._impute_missing_values(X_fit)
X_fit_array, y_fit_array = self._validate_X_y(
X_fit_imputed, y_fit, allow_nan=False
)
current_model.fit(X_fit_array, y_fit_array)
return pd.Series(
predictions, index=X.index, name=f"{self.target_col}_pred"
)
@ensure_fitted
def feature_importance(self) -> Optional[pd.DataFrame]:
"""
Returns feature importance based on model coefficients.
Returns:
A DataFrame with feature names and their corresponding
coefficients (importance scores), or None if no features.
"""
if self.model_ is None or self.feature_names_in_ is None:
return None
importance = self.model_.coef_
feature_importance_df = pd.DataFrame(
{"Feature": self.feature_names_in_, "Importance": importance}
)
feature_importance_df = feature_importance_df.sort_values(
by="Importance", key=abs, ascending=False
).reset_index(drop=True)
return feature_importance_df

View File

@@ -0,0 +1,888 @@
"""
NeuralProphet implementation for univariate time series forecasting.
This module provides a comprehensive wrapper around the NeuralProphet library,
implementing enterprise-level features including robust error handling,
parameter validation, type safety, and integration with the base model
architecture.
"""
import os
from typing import Optional, cast, Tuple, Dict, Any
import numpy as np
import pandas as pd
import torch
from neuralprophet import NeuralProphet
from rich.console import Console
from .base import UnivariateTimeSeriesModel, ensure_fitted
console = Console()
# Configure PyTorch for optimal performance
torch.set_num_threads(os.cpu_count() or 1)
class NeuralProphetModel(UnivariateTimeSeriesModel):
"""
Enterprise-grade NeuralProphet implementation for univariate
time series forecasting.
This class provides a robust wrapper around the NeuralProphet model with
comprehensive error handling, parameter validation, and integration with
the base model architecture. It includes features like automatic data
validation, performance monitoring, and enterprise-level logging.
Key Features:
- Comprehensive parameter validation
- Robust error handling with detailed diagnostics
- Memory-efficient data processing
- Integration with base model utilities
- Performance monitoring and logging
- Support for various seasonality patterns
- Flexible forecasting capabilities
Example:
>>> model = NeuralProphetModel(
... n_lags=7,
... n_forecasts=3,
... epochs=50,
... weekly_seasonality=True
... )
>>> model.fit(y_train)
>>> predictions = model.predict()
>>> future_forecast = model.forecast(forecast_horizon=3)
"""
# Class constants for validation
VALID_SEASONALITY_MODES = {"additive", "multiplicative"}
VALID_LOSS_FUNCTIONS = {"Huber", "MSE", "MAE"}
VALID_NORMALIZE_OPTIONS = {"auto", "soft", "off", "minmax"}
MIN_EPOCHS = 1
MAX_EPOCHS = 10000
MIN_N_LAGS = 0
MAX_N_LAGS = 365
MIN_N_FORECASTS = 1
MAX_N_FORECASTS = 365
def __init__(
self,
name: Optional[str] = None,
n_lags: int = 1,
n_forecasts: int = 2,
weekly_seasonality: bool = True,
daily_seasonality: bool = True,
yearly_seasonality: bool = False,
seasonality_mode: str = "additive",
epochs: int = 100,
learning_rate: Optional[float] = None,
batch_size: Optional[int] = None,
loss_func: str = "Huber",
normalize: str = "auto",
impute_missing: bool = True,
drop_missing: bool = False,
time_col: str = "ds",
target_col: str = "y",
random_seed: int = 42,
):
"""
Initialize the NeuralProphet time series model with comprehensive
validation.
Args:
name: Optional identifier for the model
n_lags: Number of lagged target values to include as inputs (0-365)
n_forecasts: Number of steps ahead to forecast (1-365)
weekly_seasonality: Whether to include weekly seasonality
daily_seasonality: Whether to include daily seasonality
yearly_seasonality: Whether to include yearly seasonality
seasonality_mode: Type of seasonality ('additive' or
'multiplicative')
epochs: Number of training epochs (1-10000)
learning_rate: Learning rate for optimizer (auto if None)
batch_size: Training batch size (auto if None)
loss_func: Loss function ('Huber', 'MSE', 'MAE')
normalize: Normalization type ('auto', 'soft', 'off', 'minmax')
impute_missing: Whether to automatically impute missing values
drop_missing: Whether to drop missing values in training data
time_col: Name of the time column
target_col: Name of the target column
random_seed: Random seed for reproducibility
Raises:
ValueError: If any parameters are invalid
TypeError: If parameters have incorrect types
"""
self._validate_and_set_parameters(
n_lags=n_lags,
n_forecasts=n_forecasts,
seasonality_mode=seasonality_mode,
epochs=epochs,
learning_rate=learning_rate,
batch_size=batch_size,
loss_func=loss_func,
normalize=normalize,
weekly_seasonality=weekly_seasonality,
daily_seasonality=daily_seasonality,
yearly_seasonality=yearly_seasonality,
impute_missing=impute_missing,
drop_missing=drop_missing,
)
super().__init__(
name=name,
time_col=time_col,
target_col=target_col,
random_seed=random_seed,
n_lags=n_lags,
)
self.forecast_horizon = n_forecasts
# Initialize model state
self.model_: Optional[NeuralProphet] = None
self.backtest_predictions_: Optional[pd.Series] = None
self._training_metrics: Dict[str, float] = {}
console.log(
f"[green]Initialized NeuralProphetModel: {self.summary()}[/green]"
)
def _validate_and_set_parameters(
self,
n_lags: int,
n_forecasts: int,
seasonality_mode: str,
epochs: int,
learning_rate: Optional[float],
batch_size: Optional[int],
loss_func: str,
normalize: str,
weekly_seasonality: bool,
daily_seasonality: bool,
yearly_seasonality: bool,
impute_missing: bool,
drop_missing: bool,
) -> None:
"""Validate and set model parameters with comprehensive checks."""
# Validate integer parameters
if not (self.MIN_N_LAGS <= n_lags <= self.MAX_N_LAGS):
raise ValueError(
f"n_lags must be between {self.MIN_N_LAGS} and "
+ f"{self.MAX_N_LAGS}, got {n_lags}"
)
if not (self.MIN_N_FORECASTS <= n_forecasts <= self.MAX_N_FORECASTS):
raise ValueError(
f"n_forecasts must be between {self.MIN_N_FORECASTS} and "
+ f"{self.MAX_N_FORECASTS}, got {n_forecasts}"
)
if not (self.MIN_EPOCHS <= epochs <= self.MAX_EPOCHS):
raise ValueError(
f"epochs must be between {self.MIN_EPOCHS} and "
+ f"{self.MAX_EPOCHS}, got {epochs}"
)
# Validate string parameters
if seasonality_mode not in self.VALID_SEASONALITY_MODES:
raise ValueError(
"seasonality_mode must be one of "
+ f"{self.VALID_SEASONALITY_MODES}, got {seasonality_mode}"
)
if loss_func not in self.VALID_LOSS_FUNCTIONS:
raise ValueError(
"loss_func must be one of "
+ f"{self.VALID_LOSS_FUNCTIONS}, got {loss_func}"
)
if normalize not in self.VALID_NORMALIZE_OPTIONS:
raise ValueError(
"normalize must be one of "
+ f"{self.VALID_NORMALIZE_OPTIONS}, got {normalize}"
)
# Validate optional float parameters
if learning_rate is not None:
if (
not isinstance(learning_rate, (int, float))
or learning_rate <= 0
):
raise ValueError(
"learning_rate must be a positive number, "
+ f"got {learning_rate}"
)
if batch_size is not None:
if not isinstance(batch_size, int) or batch_size <= 0:
raise ValueError(
"batch_size must be a positive integer, "
+ f"got {batch_size}"
)
# Validate boolean parameters
for param_name, param_value in [
("weekly_seasonality", weekly_seasonality),
("daily_seasonality", daily_seasonality),
("yearly_seasonality", yearly_seasonality),
("impute_missing", impute_missing),
("drop_missing", drop_missing),
]:
if not isinstance(param_value, bool):
raise TypeError(
f"{param_name} must be a boolean, "
+ f"got {type(param_value)}"
)
# Set validated parameters
self.n_forecasts = n_forecasts
self.weekly_seasonality = weekly_seasonality
self.daily_seasonality = daily_seasonality
self.yearly_seasonality = yearly_seasonality
self.seasonality_mode = seasonality_mode
self.epochs = epochs
self.learning_rate = learning_rate
self.batch_size = batch_size
self.loss_func = loss_func
self.normalize = normalize
self.impute_missing = impute_missing
self.drop_missing = drop_missing
def _create_model(self) -> NeuralProphet:
"""
Create a new NeuralProphet instance with validated parameters.
Returns:
A new NeuralProphet model instance configured with current
parameters.
Raises:
RuntimeError: If model creation fails
"""
try:
model_params = {
"n_lags": self.n_lags,
"n_forecasts": self.n_forecasts,
"weekly_seasonality": self.weekly_seasonality,
"daily_seasonality": self.daily_seasonality,
"yearly_seasonality": self.yearly_seasonality,
"seasonality_mode": self.seasonality_mode,
"loss_func": self.loss_func,
"normalize": self.normalize,
"impute_missing": self.impute_missing,
"drop_missing": self.drop_missing,
"impute_rolling": 1000000,
"impute_linear": 100000,
}
# Add optional parameters if specified
if self.learning_rate is not None:
model_params["learning_rate"] = self.learning_rate
if self.batch_size is not None:
model_params["batch_size"] = self.batch_size
console.log(
f"[blue]Creating NeuralProphet with params: "
f"{model_params}[/blue]"
)
return NeuralProphet(**model_params)
except Exception as e:
raise RuntimeError(
f"Failed to create NeuralProphet model: {e}"
) from e
def _validate_and_prepare_data(
self, y: pd.Series
) -> Tuple[pd.DataFrame, pd.Series]:
"""
Validate and prepare time series data for NeuralProphet.
Args:
y: Input time series data
Returns:
Tuple of (prepared_dataframe, validated_series)
Raises:
ValueError: If data validation fails
"""
try:
# Validate target series
y_array = self._validate_y(y)
# Ensure datetime index
if not pd.api.types.is_datetime64_any_dtype(y.index):
try:
y_datetime = y.copy()
y_datetime.index = pd.to_datetime(y.index)
console.log("[yellow]Converted index to datetime[/yellow]")
except Exception as e:
raise ValueError(
"y's index must be a DateTime index or convertible "
f"to DateTime. Conversion failed: {e}"
) from e
else:
y_datetime = y.copy()
# Check for minimum data requirements
if len(y_datetime) < max(self.n_lags + 1, 10):
raise ValueError(
"Insufficient data: need at least "
+ f"{max(self.n_lags + 1, 10)} observations, "
+ f"got {len(y_datetime)}"
)
# Create NeuralProphet format DataFrame
df = pd.DataFrame({"ds": y_datetime.index, "y": y_array})
# Validate for missing values if not configured to handle them
if not self.impute_missing and bool(df["y"].isna().any()):
raise ValueError(
"Data contains missing values but impute_missing=False. "
"Either set impute_missing=True or clean the data."
)
console.log(
f"[green]Data validation successful: {len(df)} "
f"observations[/green]"
)
return df, y_datetime
except Exception as e:
console.print(f"[red]Data validation failed: {e}[/red]")
raise
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 for NeuralProphet model with enhanced error
handling.
Args:
y: The target time series data with DateTime index
X: Optional DataFrame (unused for univariate model)
X_val: Validation features (unused for NeuralProphet)
y_val: Validation target (unused for NeuralProphet)
Raises:
ValueError: If data validation fails
RuntimeError: If model fitting fails
"""
try:
console.log("[blue]Starting NeuralProphet model fitting...[/blue]")
df, y_datetime = self._validate_and_prepare_data(y)
self.training_series_ = y_datetime
self.model_ = self._create_model()
console.log(
f"[blue]Training model for {self.epochs} epochs...[/blue]"
)
fit_result = self.model_.fit(df, epochs=self.epochs)
# Store training metrics if available
if hasattr(fit_result, "losses") and fit_result is not None:
losses = getattr(fit_result, "losses", None)
if losses:
self._training_metrics = {
"final_loss": float(losses[-1]),
"epochs_trained": len(losses),
}
console.log(
f"[green]Model fitting completed successfully. "
f"Metrics: {self._training_metrics}[/green]"
)
except Exception as e:
console.print(f"[red]Model fitting failed: {e}[/red]")
# Reset model state on failure
self.model_ = None
self.training_series_ = None
raise RuntimeError(
f"Failed to fit NeuralProphet model: {e}"
) from e
def _prepare_prediction_data(
self, X: Optional[pd.DataFrame] = None
) -> Tuple[pd.DataFrame, pd.Series]:
"""
Prepare data for prediction with comprehensive validation.
Args:
X: Optional DataFrame containing prediction data
Returns:
Tuple of (prepared_dataframe, prediction_index)
Raises:
ValueError: If data preparation fails
"""
if self.training_series_ is None:
raise ValueError("No training series available")
training_series = cast(pd.Series, self.training_series_)
if X is None:
# Predict on training data
df = pd.DataFrame(
{
"ds": training_series.index,
"y": training_series.to_numpy(),
}
)
return df, pd.Series(training_series.index)
# Handle various input formats for X
try:
ds_values, y_values = self._extract_time_and_target_from_X(X)
# Ensure datetime format
if not pd.api.types.is_datetime64_any_dtype(ds_values):
ds_values = pd.to_datetime(ds_values)
df = pd.DataFrame(
{
"ds": ds_values,
"y": y_values,
}
)
return df, ds_values
except Exception as e:
raise ValueError(
f"Failed to prepare prediction data: {e}"
) from e
def _extract_time_and_target_from_X(
self, X: pd.DataFrame
) -> Tuple[pd.Series, pd.Series]:
"""
Extract time and target columns from input DataFrame.
Args:
X: Input DataFrame
Returns:
Tuple of (time_series, target_series)
Raises:
ValueError: If extraction fails
"""
# Scenario 1: Explicit time and target columns
if self.time_col in X.columns and self.target_col in X.columns:
return (
X[self.time_col],
self._validate_y(X[self.target_col])
)
# Scenario 2: DateTime index
elif pd.api.types.is_datetime64_any_dtype(X.index):
ds_values = pd.Series(X.index, name=self.time_col)
if self.target_col in X.columns:
# DateTime index with explicit target column
return (
ds_values,
self._validate_y(X[self.target_col])
)
elif X.shape[1] == 1:
# DateTime index with single data column
return ds_values, self._validate_y(
X.iloc[:, 0].rename(self.target_col)
)
elif X.shape[1] == 0:
# Only index, no columns - forecast scenario
y_values = pd.Series(
np.nan, index=X.index, name=self.target_col
)
return ds_values, y_values
else:
raise ValueError(
"X has DateTime index but cannot identify target column. "
+ f"Expected '{self.target_col}' or single column. "
+ f"Found: {X.columns.tolist()}"
)
else:
raise ValueError(
"Cannot determine time and target from X. "
+ f"Provide columns '{self.time_col}' and '{self.target_col}' "
+ "or use DateTime index."
)
@ensure_fitted
def predict(self, X: Optional[pd.DataFrame] = None) -> pd.Series:
"""
Generate in-sample predictions with enhanced error handling.
Args:
X: Optional DataFrame containing timestamps and target values.
If None, predicts on training data.
Returns:
Series of predictions indexed by timestamp
Raises:
ValueError: If model is not fitted or prediction fails
RuntimeError: If prediction computation fails
"""
if self.model_ is None:
raise ValueError("Model is not fitted")
try:
console.log("[blue]Generating predictions...[/blue]")
# Prepare prediction data
df, predictions_index = self._prepare_prediction_data(X)
# Get training context for lagged features
training_series = cast(pd.Series, self.training_series_)
past_values = pd.DataFrame(
{
"ds": training_series.index,
"y": training_series.to_numpy(),
}
).iloc[-self.n_lags :, :]
# Combine past and prediction data
combined_df = pd.concat([past_values, df], ignore_index=True)
combined_df = (
combined_df.sort_values(by="ds")
.reset_index(drop=True)
.drop_duplicates(subset="ds", keep="last")
)
# Generate forecast
forecast = self.model_.predict(combined_df)
# Handle different forecast column formats
forecast_col = f"yhat{self.n_forecasts}"
if forecast_col not in forecast.columns:
forecast = self.model_.get_last_forecast(
forecast, include_previous_forecasts=self.n_forecasts
)
# Extract predictions for requested indices
forecast = forecast.set_index("ds")
predictions = forecast.loc[predictions_index, forecast_col]
console.log(
f"[green]Generated {len(predictions)} predictions[/green]"
)
return predictions
except Exception as e:
console.print(f"[red]Prediction failed: {e}[/red]")
raise RuntimeError(f"Failed to generate predictions: {e}") from e
@ensure_fitted
def forecast(self, forecast_horizon: int) -> np.ndarray:
"""
Generate future forecasts with comprehensive validation.
Args:
forecast_horizon: Number of steps to forecast ahead
(1 to n_forecasts)
Returns:
Array of forecasted values, indexed by the forecast horizon
Raises:
ValueError: If forecast_horizon is invalid or model not fitted
RuntimeError: If forecast generation fails
"""
if self.model_ is None:
raise ValueError("Model is not fitted")
if not (1 <= forecast_horizon <= self.n_forecasts):
raise ValueError(
"forecast_horizon must be between 1 and "
+ f"{self.n_forecasts}, got {forecast_horizon}"
)
try:
console.log(
f"[blue]Generating {forecast_horizon}-step forecast...[/blue]"
)
training_series = cast(pd.Series, self.training_series_)
# Create future dataframe
future_df = self.model_.make_future_dataframe(
df=pd.DataFrame(
{
"ds": training_series.index,
"y": training_series.to_numpy(),
}
),
periods=forecast_horizon,
)
# Generate forecasts
forecast = self.model_.predict(future_df)
# Extract forecasted values for each horizon
forecasted_values = np.empty(forecast_horizon)
for i in range(forecast_horizon):
col_name = f"yhat{i + 1}"
if col_name in forecast.columns:
values = forecast[col_name].dropna()
if len(values) > 0:
forecasted_values[i] = values.iloc[0]
else:
forecasted_values[i] = np.nan
else:
forecasted_values[i] = np.nan
console.log(
f"[green]Generated forecast: {len(forecasted_values)}"
+ " values[/green]"
)
return forecasted_values
except Exception as e:
console.print(f"[red]Forecast generation failed: {e}[/red]")
raise RuntimeError(f"Failed to generate forecast: {e}") from e
@ensure_fitted
def backtest(
self,
y: pd.Series,
X: Optional[pd.DataFrame] = None,
retrain_every: int = 50,
reuse_previous_execution: bool = False,
) -> pd.Series:
"""
Perform comprehensive backtesting with enhanced monitoring.
This method implements walk-forward validation with periodic
retraining, providing robust evaluation of model performance in
production-like scenarios.
Args:
y: Target time series for backtesting (must have DateTime index)
X: Unused (included for base class compatibility)
retrain_every: Unused (model retrains at each step)
reuse_previous_execution: Whether to reuse previous backtest
results
Returns:
Series of backtested predictions indexed by timestamp
Raises:
ValueError: If parameters are invalid or data is insufficient
RuntimeError: If backtesting fails
"""
if self.model_ is None:
raise ValueError("Model is not fitted")
if self.training_series_ is None:
raise ValueError("No training series found")
if not (1 <= self.forecast_horizon <= self.n_forecasts):
raise ValueError(
"forecast_horizon must be between 1 and "
+ f"{self.n_forecasts}, got {self.forecast_horizon}"
)
# Handle reuse of previous execution
if reuse_previous_execution and self.backtest_predictions_ is not None:
expected_index = y.iloc[self.forecast_horizon:].index
if (
len(self.backtest_predictions_) == len(expected_index)
and (self.backtest_predictions_.index == expected_index).all()
):
console.log(
"[yellow]Reusing previous backtest results[/yellow]"
)
return self.backtest_predictions_
else:
console.log(
"[yellow]Previous results incompatible, running new"
+ " backtest[/yellow]"
)
try:
console.log(
f"[blue]Starting backtest with {self.forecast_horizon}-step"
+ " horizon...[/blue]"
)
# Validate and prepare data
y_sorted = y.sort_index()
# Check for overlapping data
training_series = cast(pd.Series, self.training_series_)
if any(t in training_series.index for t in y_sorted.index):
console.print(
"[yellow]Warning: Backtest data overlaps with training"
+ " data[/yellow]"
)
# Initialize backtesting
predictions = []
training_base = pd.DataFrame(
{
"ds": training_series.index,
"y": training_series.values,
}
)
timestamps = y_sorted.index
total_steps = len(timestamps) - self.forecast_horizon
console.log(
f"[blue]Running {total_steps} backtest steps...[/blue]"
)
# Perform walk-forward validation
for i in range(self.forecast_horizon, len(timestamps)):
if i % 50 == 0: # Progress logging
console.log(
f"[blue]Backtest progress: {i}/{len(timestamps)}[/blue]"
)
t = timestamps[i]
t_minus_h = timestamps[i - self.forecast_horizon]
# Prepare training data up to t - h
history = y_sorted.loc[:t_minus_h]
train_df = (
pd.concat(
[
training_base,
pd.DataFrame(
{"ds": history.index, "y": history.values}
),
],
ignore_index=True,
)
.drop_duplicates(subset="ds")
.sort_values("ds")
)
# Check minimum data requirement
if len(train_df) < max(self.n_lags + 1, 10):
console.print(
f"[yellow]Insufficient data at step {i}, "
+ "skipping[/yellow]"
)
predictions.append((t, np.nan))
continue
try:
# Retrain model
model = self._create_model()
model.fit(train_df, epochs=self.epochs)
# Generate forecast
mask = train_df["ds"] <= t_minus_h
future_df = model.make_future_dataframe(
df=train_df.loc[mask], periods=self.forecast_horizon
)
forecast = model.predict(future_df, decompose=False)
# Extract prediction
forecast_col = f"yhat{self.forecast_horizon}"
prediction_rows = forecast[forecast["ds"] == t]
if (
len(prediction_rows) > 0
and forecast_col in forecast.columns
):
prediction = prediction_rows[forecast_col].iloc[0]
else:
prediction = np.nan
predictions.append((t, prediction))
except Exception as step_error:
console.print(
f"[yellow]Error at step {i}: {step_error}[/yellow]"
)
predictions.append((t, np.nan))
# Create results series
if predictions:
pred_index, pred_values = zip(*predictions)
self.backtest_predictions_ = pd.Series(
pred_values,
index=pd.Index(pred_index),
name=f"yhat{self.forecast_horizon}",
)
else:
self.backtest_predictions_ = pd.Series(
dtype=float, name=f"yhat{self.forecast_horizon}"
)
console.log(
f"[green]Backtest completed: {len(self.backtest_predictions_)}"
+ " predictions[/green]"
)
return self.backtest_predictions_
except Exception as e:
console.print(f"[red]Backtest failed: {e}[/red]")
raise RuntimeError(f"Failed to perform backtest: {e}") from e
def get_params_dict(self) -> Dict[str, Any]:
"""Get comprehensive model parameters for logging/serialization."""
base_params = super().get_params_dict()
neural_prophet_params = {
"n_forecasts": self.n_forecasts,
"weekly_seasonality": self.weekly_seasonality,
"daily_seasonality": self.daily_seasonality,
"yearly_seasonality": self.yearly_seasonality,
"seasonality_mode": self.seasonality_mode,
"epochs": self.epochs,
"learning_rate": self.learning_rate,
"batch_size": self.batch_size,
"loss_func": self.loss_func,
"normalize": self.normalize,
"impute_missing": self.impute_missing,
"drop_missing": self.drop_missing,
"training_metrics": self._training_metrics,
}
return {**base_params, **neural_prophet_params}
def summary(self) -> str:
"""Generate comprehensive model summary."""
fitted_status = (
"✓ Fitted" if self.__sklearn_is_fitted__() else "✗ Not fitted"
)
seasonality_features = []
if self.weekly_seasonality:
seasonality_features.append("Weekly")
if self.daily_seasonality:
seasonality_features.append("Daily")
if self.yearly_seasonality:
seasonality_features.append("Yearly")
seasonality_str = (
", ".join(seasonality_features) if seasonality_features else "None"
)
summary_lines = [
f"Model: {self.__class__.__name__}",
f"Status: {fitted_status}",
f"Lags: {self.n_lags}, Forecasts: {self.n_forecasts}",
f"Seasonality: {seasonality_str} ({self.seasonality_mode})",
f"Training: {self.epochs} epochs, {self.loss_func} loss",
f"Data Handling: Impute={self.impute_missing}, "
f"Drop={self.drop_missing}",
]
if self._training_metrics:
metrics_str = ", ".join(
f"{k}={v:.4f}" for k, v in self._training_metrics.items()
)
summary_lines.append(f"Metrics: {metrics_str}")
return "\n".join(summary_lines)

View File

@@ -0,0 +1,695 @@
"""
Stacking implementation for time series forecasting.
This module provides a stacking regressor implementation that combines multiple
time series models' predictions using a meta-model. The stacking approach
helps improve prediction accuracy by combining the strengths of different
base models through a learned meta-model.
Key features:
- **Model Stacking**: Combines predictions from multiple base models
- **Time-Series Aware**: Uses proper time-based cross-validation
- **Meta-Model Learning**: Learns optimal combination weights
- **Comprehensive Error Handling**: Robust error handling and validation
- **Rich Logging**: Colored console output for better debugging
The stacking model loads pre-trained base models and uses their predictions
as features for training a meta-model (CatBoost by default).
"""
from typing import Optional, List, Dict, Any, Union
import pandas as pd
import numpy as np
from catboost import CatBoostRegressor, CatBoostClassifier, Pool
from rich.console import Console
import gzip
import pickle
import lzma
from .base import (
MultivariateTimeSeriesModel,
ensure_fitted,
TimeSeriesModel,
)
console = Console()
class StackingTimeSeriesModel(MultivariateTimeSeriesModel):
"""
Stacking implementation for time series forecasting.
This class implements stacking of multiple base models, using their
predictions as features for a meta-model. It handles time-based
cross-validation to generate out-of-fold predictions for training.
The model supports both regression and classification tasks through
the meta-model configuration.
Attributes:
base_models_: List of loaded base models
model_: The trained meta-model (CatBoost)
training_series_: Copy of training target data
base_predictions_train_: Base model predictions on training data
backtest_predictions_: Stored backtest predictions for reuse
Example:
>>> stacking_model = StackingTimeSeriesModel(
... base_model_paths=["model1.pkl", "model2.pkl"],
... base_model_types=["catboost", "elasticnet"]
... )
>>> stacking_model.fit(y=target_series, X=feature_matrix)
>>> predictions = stacking_model.predict(X=test_features)
"""
def __init__(
self,
base_model_paths: List[str],
base_model_types: List[str],
name: Optional[str] = None,
learning_task: str = "regression",
retrain_every: int = 100,
meta_iterations: int = 1000,
meta_learning_rate: float = 0.1,
meta_depth: int = 6,
early_stopping_rounds: Optional[int] = None,
meta_loss_function: Optional[str] = None,
time_col: str = "ds",
target_col: str = "y",
random_seed: int = 42,
verbose: bool = False,
differentiate_target: bool = False,
bins: Optional[List[float]] = None,
use_predict_for_training: bool = True,
) -> None:
"""
Initialize the stacking model.
Args:
base_model_paths: Paths to saved base models
base_model_types: Types of base models (must match order of paths)
name: Optional identifier for the model
learning_task: Type of learning task ('regression', 'binary',
'multiclass')
retrain_every: Frequency of retraining during backtesting
meta_iterations: Number of iterations for meta-model
meta_learning_rate: Learning rate for meta-model
meta_depth: Tree depth for meta-model
early_stopping_rounds: Early stopping rounds for meta-model
meta_loss_function: Loss function for meta-model
time_col: Name of time column
target_col: Name of target column
random_seed: Random seed
verbose: Whether to print verbose logging
differentiate_target: Whether to differentiate the target series
bins: Bin edges for multiclass classification
use_predict_for_training: If True, use predict() instead of
backtest() for generating base model predictions during
training. This is much faster but may lead to overfitting
since the meta-model trains on in-sample predictions.
"""
super().__init__(
name=name,
time_col=time_col,
target_col=target_col,
random_seed=random_seed,
learning_task=learning_task,
differentiate_target=differentiate_target,
bins=bins,
)
# Validate inputs
if not base_model_paths:
raise ValueError("base_model_paths cannot be empty")
if not base_model_types:
raise ValueError("base_model_types cannot be empty")
if len(base_model_paths) != len(base_model_types):
raise ValueError(
"base_model_paths and base_model_types must have same length"
)
self.retrain_every = retrain_every
self.base_model_paths = base_model_paths
self.base_model_types = base_model_types
self.meta_iterations = meta_iterations
self.meta_learning_rate = meta_learning_rate
self.meta_depth = meta_depth
self.meta_loss_function = self._get_default_loss_function(
meta_loss_function
)
self.early_stopping_rounds = early_stopping_rounds
self.verbose = verbose
self.use_predict_for_training = use_predict_for_training
# Will be set during fit
self.base_models_: List[TimeSeriesModel] = []
self.model_: Optional[
Union[CatBoostRegressor, CatBoostClassifier]
] = None
self.training_series_: Optional[pd.Series] = None
self.base_predictions_train_: Optional[pd.DataFrame] = None
self.backtest_predictions_: Optional[pd.Series] = None
self._load_base_models()
if self.verbose:
console.log(
"[green]Initialized StackingTimeSeriesModel: "
+ f"{self.summary()}[/green]"
)
def _load_base_models(self) -> None:
"""Load all base models from their saved paths."""
from .factory import load_model
self.base_models_ = []
for model_path, model_type in zip(
self.base_model_paths, self.base_model_types
):
try:
model = load_model(model_path, model_type)
self.base_models_.append(model)
if self.verbose:
console.log(
f"[blue]Loaded {model_type} model from "
+ f"{model_path}[/blue]"
)
except Exception as e:
console.print(
f"[red]Error loading model from {model_path}: {e}[/red]"
)
raise ValueError(f"Failed to load model: {model_path}") from e
if self.verbose:
console.log(
f"[green]Successfully loaded {len(self.base_models_)} "
+ "base models[/green]"
)
def _create_meta_model(
self,
) -> Union[CatBoostRegressor, CatBoostClassifier]:
"""Creates a new instance of meta-model.
Returns:
A new CatBoost model instance (Regressor or Classifier).
"""
base_params = {
"iterations": self.meta_iterations,
"learning_rate": self.meta_learning_rate,
"depth": self.meta_depth,
"loss_function": self.meta_loss_function,
"early_stopping_rounds": self.early_stopping_rounds,
"random_seed": self.random_seed,
"verbose": self.verbose,
}
if self.learning_task in ["binary", "multiclass"]:
return CatBoostClassifier(
auto_class_weights="Balanced", **base_params
)
else:
return CatBoostRegressor(**base_params)
def _get_base_predictions(
self, X: Optional[pd.DataFrame], y: Optional[pd.Series] = None
) -> pd.DataFrame:
"""Get predictions from all base models.
Args:
X: Feature matrix
y: Target series (optional, used for validation)
Returns:
DataFrame with predictions in model_0, model_1, etc columns
"""
if X is None:
raise ValueError("Feature matrix X cannot be None")
all_preds = []
for i, model in enumerate(self.base_models_):
try:
preds = model.predict(X)
model_name = f"model_{model.name or i}"
all_preds.append(
pd.Series(preds, name=model_name, index=X.index)
)
if self.verbose:
console.log(
f"[cyan]Generated predictions from {model_name}[/cyan]"
)
except Exception as e:
console.print(
f"[red]Error getting predictions from model {i}: {e}[/red]"
)
raise
return pd.concat(all_preds, axis=1)
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 for stacking model.
Get predictions from base models on training data, then train
meta-model on those predictions.
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)
"""
if X is None:
raise ValueError("Feature matrix X must be provided for stacking")
# Preprocess data using parent class method
y_processed, X_processed, y_val_processed, X_val_processed = (
self._preprocess_data(y, X, X_val, y_val)
)
if X_processed is None:
raise ValueError(
"Feature matrix X cannot be None after preprocessing"
)
if self.verbose:
method_name = (
"predictions"
if self.use_predict_for_training
else "backtest predictions"
)
console.log(
f"[blue]Generating base model {method_name} "
+ "for stacking...[/blue]"
)
# Get base predictions - use either predict or backtest based on
# setting
if self.use_predict_for_training:
# Fast approach: use direct predictions (may overfit)
meta_features = self._get_base_predictions(X_processed)
# Align with target data
common_index = meta_features.index.intersection(y_processed.index)
meta_features = meta_features.loc[common_index]
y_aligned = y_processed.loc[common_index]
if self.verbose:
console.log(
"[yellow]Warning: Using predict() for training may "
+ "lead to overfitting since meta-model trains on "
+ "in-sample predictions[/yellow]"
)
else:
# Robust approach: use backtesting to avoid overfitting
all_preds = []
for i, model in enumerate(self.base_models_):
try:
preds = model.backtest(
y_processed,
X_processed,
retrain_every=self.retrain_every,
)
model_name = f"model_{model.name or i}"
all_preds.append(pd.Series(preds, name=model_name))
if self.verbose:
console.log(
"[cyan]Generated backtest predictions from "
+ f"{model_name}[/cyan]"
)
except Exception as e:
console.print(
f"[red]Error during backtesting for model {i}: "
+ f"{e}[/red]"
)
raise
meta_features = pd.concat(all_preds, axis=1)
# Align with target data (backtest might have different length)
common_index = meta_features.index.intersection(y_processed.index)
meta_features = meta_features.loc[common_index]
y_aligned = y_processed.loc[common_index]
if self.verbose:
console.log(
f"[blue]Training meta-model with {len(meta_features)} "
+ f"samples and {meta_features.shape[1]} base model "
+ "features[/blue]"
)
meta_X, meta_y = self._validate_X_y(meta_features, y_aligned)
train_pool = Pool(data=meta_X, label=meta_y)
# Prepare validation data if provided
eval_set = None
if X_val_processed is not None and y_val_processed is not None:
val_predictions = self._get_base_predictions(
X_val_processed, y_val_processed
)
val_X, val_y = self._validate_X_y(val_predictions, y_val_processed)
eval_set = Pool(data=val_X, label=val_y)
if self.verbose:
console.log(
"[blue]Using validation set with "
+ f"{len(val_predictions)} samples[/blue]"
)
# Create and train meta-model
self.model_ = self._create_meta_model()
self.model_.fit(train_pool, eval_set=eval_set)
# Store training data
self.training_series_ = y_processed.copy()
self.base_predictions_train_ = meta_features.copy()
if self.verbose:
console.log(
"[green]Meta-model training completed successfully[/green]"
)
@ensure_fitted
def predict(self, X: pd.DataFrame) -> pd.Series:
"""
Generate predictions using the stacking model.
Args:
X: Feature matrix for prediction
Returns:
Series containing predictions
"""
if self.model_ is None:
raise ValueError("Model has not been fitted yet")
base_predictions = self._get_base_predictions(X)
X_array = self._validate_X(base_predictions)
predictions = self.model_.predict(X_array)
# Convert predictions to numpy array if needed
if hasattr(predictions, "squeeze"):
predictions = predictions.squeeze()
elif isinstance(predictions, list):
predictions = np.array(predictions)
return pd.Series(predictions, index=X.index, name=self.target_col)
@ensure_fitted
def feature_importance(self) -> Optional[pd.DataFrame]:
"""
Returns feature importance from the meta-model.
Returns:
DataFrame with feature names and their importance scores,
or None if not available.
"""
if self.model_ is None or not hasattr(
self.model_, "feature_importances_"
):
return None
if self.base_predictions_train_ is None:
return None
importances = self.model_.feature_importances_
feature_names = self.base_predictions_train_.columns
return pd.DataFrame(
{
"feature": feature_names,
"importance": importances,
}
).sort_values("importance", ascending=False)
@ensure_fitted
def backtest(
self,
y: pd.Series,
X: Optional[pd.DataFrame] = None,
retrain_every: int = 50,
reuse_previous_execution: bool = False,
) -> pd.Series:
"""
Performs backtesting (walk-forward validation) with periodic
retraining.
This method simulates a production scenario by iterating through a test
set, making a one-step-ahead prediction, and then retraining the model
periodically with the newly available data.
Args:
y: Series with the true target values for the backtesting period
X: DataFrame with features for the backtesting period
retrain_every: The frequency of retraining. The model will be
retrained every `retrain_every` steps
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:
A series of backtested predictions, indexed by the backtest data's
index
"""
if self.model_ is None:
raise ValueError("Model is not fitted yet")
if self.training_series_ is None:
raise ValueError("Training series is not set")
if X is None:
raise ValueError("Feature matrix X must be provided")
if reuse_previous_execution:
if self.backtest_predictions_ is None:
raise ValueError("No previous execution found")
if (self.backtest_predictions_.shape[0] != y.shape[0]) or (
not (self.backtest_predictions_.index == y.index).all()
):
raise ValueError(
"Previous execution index does not match y index"
)
return self.backtest_predictions_
if self.verbose:
console.log(
f"[blue]Starting backtest with {len(y)} samples, "
+ f"retraining every {retrain_every} steps[/blue]"
)
# Get base model predictions for the entire backtest period
all_base_preds = []
for i, model in enumerate(self.base_models_):
try:
preds = model.backtest(
y,
X,
retrain_every=retrain_every,
reuse_previous_execution=reuse_previous_execution,
)
model_name = f"model_{model.name or i}"
all_base_preds.append(pd.Series(preds, name=model_name))
if self.verbose:
console.log(
f"[cyan]Completed backtest for {model_name}[/cyan]"
)
except Exception as e:
console.print(
f"[red]Error during backtest for model {i}: {e}[/red]"
)
raise
meta_features = pd.concat(all_base_preds, axis=1)
# Generate meta-model predictions
predictions = self.model_.predict(self._validate_X(meta_features))
# Store backtest predictions for potential reuse
self.backtest_predictions_ = pd.Series(
predictions,
index=meta_features.index,
name=f"{self.target_col}_pred",
)
if self.verbose:
console.log(
"[green]Backtest completed: "
+ f"{len(self.backtest_predictions_)} predictions "
+ "generated[/green]"
)
return self.backtest_predictions_
def get_base_model_names(self) -> List[str]:
"""Get names of all base models.
Returns:
List of base model names
"""
return [
model.name or f"model_{i}"
for i, model in enumerate(self.base_models_)
]
def get_params_dict(self) -> Dict[str, Any]:
"""Get model parameters as dictionary for logging/serialization."""
base_params = super().get_params_dict()
stacking_params = {
"base_model_paths": self.base_model_paths,
"base_model_types": self.base_model_types,
"retrain_every": self.retrain_every,
"meta_iterations": self.meta_iterations,
"meta_learning_rate": self.meta_learning_rate,
"meta_depth": self.meta_depth,
"meta_loss_function": self.meta_loss_function,
"early_stopping_rounds": self.early_stopping_rounds,
"num_base_models": len(self.base_models_),
"use_predict_for_training": self.use_predict_for_training,
}
return {**base_params, **stacking_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"Base Models: {params.get('num_base_models', 0)}",
f"Task: {params.get('learning_task', 'regression')}",
f"Meta Loss: {params.get('meta_loss_function', 'RMSE')}",
]
return "\n".join(summary_lines)
def _optimize_base_models_for_storage(self) -> None:
"""
Optimizes base models for storage by removing unnecessary data.
This can significantly reduce pickle size, especially for neural
models.
"""
if self.verbose:
console.log("[blue]Optimizing base models for storage...[/blue]")
for i, model in enumerate(self.base_models_):
try:
# For neuralprophet models, remove training history and
# large artifacts
model_attr = getattr(model, "model", None)
if model_attr is not None and hasattr(model_attr, "trainer"):
trainer = getattr(model_attr, "trainer", None)
if trainer is not None:
# Remove trainer which contains training logs and
# can be very large
if hasattr(trainer, "logged_metrics"):
setattr(trainer, "logged_metrics", {})
if hasattr(trainer, "progress_bar_metrics"):
setattr(trainer, "progress_bar_metrics", {})
if hasattr(trainer, "callback_metrics"):
setattr(trainer, "callback_metrics", {})
# For any model with training history
if hasattr(model, "training_history_"):
setattr(model, "training_history_", None)
if hasattr(model, "validation_history_"):
setattr(model, "validation_history_", None)
# Remove cached predictions if they exist
if hasattr(model, "_cached_predictions"):
setattr(model, "_cached_predictions", None)
if self.verbose:
model_name = getattr(model, "name", f"model_{i}")
console.log(
f"[cyan]Optimized {model_name} for storage[/cyan]"
)
except Exception as e:
if self.verbose:
console.log(
f"[yellow]Warning: Could not optimize model {i}: "
+ f"{e}[/yellow]"
)
def save(self, path: str, compression: str = "gzip") -> None:
"""
Saves model to disk using compression to reduce file size.
Args:
path: File path to save to
compression: Compression method ('gzip', 'lzma', or 'none')
- 'gzip': Fast compression, ~60-80% size reduction
- 'lzma': Better compression, ~70-90% size reduction, slower
- 'none': No compression
"""
if self.verbose:
console.log(
f"[blue]Saving stacking model with {compression} "
+ f"compression to {path}[/blue]"
)
# Optimize base models for storage first
self._optimize_base_models_for_storage()
if compression == "lzma":
# LZMA provides better compression but is slower
with lzma.open(path, "wb", preset=9) as f:
pickle.dump(self, f, protocol=pickle.HIGHEST_PROTOCOL)
elif compression == "gzip":
# Gzip is faster with good compression
with gzip.open(path, "wb", compresslevel=9) as f:
pickle.dump(self, f, protocol=pickle.HIGHEST_PROTOCOL)
else:
# No compression
with open(path, "wb") as f:
pickle.dump(self, f, protocol=pickle.HIGHEST_PROTOCOL)
if self.verbose:
console.log(
f"[green]Saved compressed stacking model to {path}[/green]"
)
@classmethod
def load(cls, path: str) -> "StackingTimeSeriesModel":
"""
Loads model from disk with automatic format detection.
Supports both compressed formats and legacy joblib format.
"""
import joblib
# Try different formats in order of preference
loading_methods = [
("lzma", lambda p: lzma.open(p, "rb")),
("gzip", lambda p: gzip.open(p, "rb")),
("pickle", lambda p: open(p, "rb")),
("joblib", None), # Special case for joblib
]
for format_name, open_func in loading_methods:
try:
if format_name == "joblib":
return joblib.load(path)
else:
with open_func(path) as f:
return pickle.load(f)
except (
lzma.LZMAError,
gzip.BadGzipFile,
OSError,
pickle.UnpicklingError,
ValueError,
):
continue
raise ValueError(
f"Could not load model from {path} - unknown or corrupted format"
)

View File

@@ -0,0 +1,11 @@
channels:
- conda-forge
dependencies:
- python=3.10.16
- pip<=25.0
- pip:
- mlflow==2.7.1
- pandas
- numpy
- scikit-learn
name: mlflow-env

View File

@@ -0,0 +1,7 @@
python: 3.10.16
build_dependencies:
- pip==25.0
- setuptools==79.0.0
- wheel==0.45.1
dependencies:
- -r requirements.txt

View File

@@ -0,0 +1,4 @@
mlflow==2.7.1
pandas
numpy
scikit-learn