266 lines
10 KiB
Python
266 lines
10 KiB
Python
import numpy as np
|
|
import pandas as pd
|
|
from sklearn.base import BaseEstimator, TransformerMixin
|
|
from sklearn.linear_model import LinearRegression
|
|
from sklearn.preprocessing import PolynomialFeatures
|
|
|
|
|
|
class LinearRegressionModel(BaseEstimator, TransformerMixin):
|
|
def __init__(
|
|
self,
|
|
target_variable: str = '',
|
|
variable_columns: list | None = None,
|
|
model_params: dict | None = None,
|
|
clipping_max: float | None = None,
|
|
clipping_min: float | None = None,
|
|
weights: dict | None = None,
|
|
degree: int = 1,
|
|
interaction_only: bool = False,
|
|
verbose: bool = False,
|
|
):
|
|
"""
|
|
Linear Regression Model for Time Series Analysis.
|
|
|
|
Args:
|
|
target_variable (str): The target variable name
|
|
variable_columns (list): The input columns names in a list
|
|
model_params (dict): The parameters used for training the model \\
|
|
clipping (dict): The lower and upper limits for the target variable to be clipped \\
|
|
*Format: {'min': min_value, 'max': max_value}*
|
|
weights (dict): The weights for the Linear Regression model \\
|
|
*Format: {'variable_name': weight}*
|
|
degree (int): The degree of the polynomial features
|
|
interaction_only (bool): Whether to include interaction terms only
|
|
verbose (bool): Whether to print verbose output
|
|
|
|
Returns:
|
|
LinearRegressionModel: The prediction model object
|
|
"""
|
|
self.target_variable: str = target_variable
|
|
self.variable_columns: list = variable_columns if variable_columns else []
|
|
self.model_params: dict = model_params if model_params else {}
|
|
self.regr = LinearRegression()
|
|
self.q1_target: float | None = None
|
|
self.q3_target: float | None = None
|
|
self.weights: dict = weights if weights else {}
|
|
self.degree: int = degree
|
|
self.interaction_only: bool = interaction_only
|
|
self.poly: PolynomialFeatures | None = None
|
|
self.verbose: bool = verbose
|
|
|
|
self.clipping: dict = {}
|
|
|
|
if clipping_max is not None:
|
|
self.clipping['max'] = clipping_max
|
|
if clipping_min is not None:
|
|
self.clipping['min'] = clipping_min
|
|
|
|
def get_regressor(self) -> LinearRegression:
|
|
"""
|
|
Get LinearRegression object
|
|
|
|
Returns:
|
|
LinearRegression objetc
|
|
"""
|
|
return self.regr
|
|
|
|
def create_poly_features(self, input_data: pd.DataFrame, fit: bool = False) -> pd.DataFrame:
|
|
"""
|
|
Function to create polynomial features
|
|
|
|
Args:
|
|
input_data (pandas.DataFrame): The data used to create the polynomial features
|
|
fit (bool): Whether to fit the polynomial features creator or just transform
|
|
|
|
Returns:
|
|
pandas.DataFrame: The data with the polynomial features
|
|
"""
|
|
original_index = input_data.index
|
|
original_shape = input_data.shape
|
|
|
|
# Create polynomial features
|
|
interaction_only = self.interaction_only
|
|
|
|
degree = self.degree
|
|
if fit:
|
|
if self.verbose:
|
|
print(f'Fitting polynomial features with degree {degree}')
|
|
self.poly = PolynomialFeatures(
|
|
degree=degree, interaction_only=interaction_only, include_bias=False
|
|
)
|
|
input_data = self.poly.fit_transform(input_data)
|
|
else:
|
|
if self.verbose:
|
|
print(f'Transforming polynomial features with degree {degree}')
|
|
if self.poly is not None:
|
|
input_data = self.poly.transform(input_data)
|
|
else:
|
|
raise ValueError(
|
|
'Polynomial transformer is not fitted. Call fit() before predict() when degree > 1.'
|
|
)
|
|
|
|
# Create new columns names
|
|
modified_columns = self.poly.get_feature_names_out(self.variable_columns)
|
|
|
|
input_data = pd.DataFrame(input_data, columns=modified_columns, index=original_index)
|
|
|
|
if self.verbose:
|
|
print(f'Data shape before transformation: {original_shape}')
|
|
print(f'Data shape after transformation: {input_data.shape}')
|
|
|
|
return input_data
|
|
|
|
def fit(self, input_data: pd.DataFrame) -> 'LinearRegressionModel':
|
|
"""
|
|
Function to fit the model
|
|
|
|
Args:
|
|
input_data (pandas.DataFrame): The data used to fit the Linear Regression model
|
|
|
|
Returns:
|
|
LinearRegressionModel: The prediction model object
|
|
"""
|
|
if self.verbose:
|
|
# Display the header
|
|
text = 'INITIATING LINEAR REGRESSION MODEL FIT'
|
|
print('\n' + '-' * len(text))
|
|
print(text)
|
|
print('-' * len(text) + '\n')
|
|
|
|
# Basic validations and inference
|
|
if not self.target_variable:
|
|
raise ValueError("'target_variable' must be set before calling fit().")
|
|
|
|
if not self.variable_columns:
|
|
# Infer all columns except target
|
|
self.variable_columns = [c for c in input_data.columns if c != self.target_variable]
|
|
|
|
# Validate required columns exist
|
|
missing_features = [c for c in self.variable_columns if c not in input_data.columns]
|
|
if missing_features:
|
|
raise ValueError(
|
|
'Training data is missing required feature columns: ' + ', '.join(missing_features)
|
|
)
|
|
|
|
if self.target_variable not in input_data.columns:
|
|
# Let pandas raise KeyError with the column name to satisfy tests too
|
|
raise KeyError(repr(self.target_variable))
|
|
|
|
X_train = input_data[self.variable_columns]
|
|
y_train = input_data[self.target_variable]
|
|
|
|
if self.verbose:
|
|
print(f'X_train shape: {X_train.shape}')
|
|
print(f'y_train shape: {y_train.shape}')
|
|
|
|
# Create polynomial features
|
|
if self.degree > 1:
|
|
X_train = self.create_poly_features(X_train, fit=True)
|
|
|
|
# Handle infinite values
|
|
X_train = X_train.replace([np.inf, -np.inf], np.nan)
|
|
|
|
# Drop columns with all null values
|
|
all_null_cols = X_train.columns[X_train.isnull().all()].tolist()
|
|
X_train = X_train.drop(columns=all_null_cols)
|
|
|
|
if self.verbose and all_null_cols:
|
|
print(f'{len(all_null_cols)} columns with all null values were dropped.')
|
|
|
|
# Get mask of rows with any NaN values
|
|
mask = X_train.notna().all(axis=1)
|
|
X_train = X_train[mask]
|
|
y_train = y_train[mask]
|
|
if self.verbose:
|
|
print(f'X_train shape after removing NaN rows: {X_train.shape}')
|
|
print(f'y_train shape after removing NaN rows: {y_train.shape}')
|
|
|
|
self.q1_target = y_train.quantile(0.25)
|
|
self.q3_target = y_train.quantile(0.75)
|
|
|
|
# Fit the model
|
|
self.regr.fit(X_train, y_train)
|
|
if self.verbose:
|
|
print('Model training completed successfully.')
|
|
|
|
# Get the weights using actual trained feature names (handles polynomial features)
|
|
# ravel() so coef_ is 1D when y was a column (e.g. DataFrame), avoiding scalar conversion errors
|
|
round_coef = np.round(self.regr.coef_, 3).ravel()
|
|
round_intercept = np.round(self.regr.intercept_, 3).ravel()
|
|
feature_names = list(getattr(self.regr, 'feature_names_in_', X_train.columns))
|
|
|
|
# Save the weights (strict=True to ensure feature_names and coefficients align)
|
|
weights = dict(zip(feature_names, [float(c) for c in round_coef], strict=True))
|
|
weights = dict(sorted(weights.items(), key=lambda item: abs(item[1]), reverse=True))
|
|
weights = {'Bias': float(round_intercept.flat[0]), **weights}
|
|
self.weights = weights
|
|
|
|
if self.verbose:
|
|
print('\nModel Weights:')
|
|
for key, value in weights.items():
|
|
print(f'{key}: {value}')
|
|
|
|
return self
|
|
|
|
def predict(self, input_data: pd.DataFrame) -> np.ndarray:
|
|
"""
|
|
Function to predict the target variable.
|
|
If clipping is True, the predictions are clipped based on the target variable quartiles.
|
|
|
|
Args:
|
|
input_data (pandas.DataFrame): The data used to predict the target variable
|
|
|
|
Returns:
|
|
numpy.ndarray: The predicted target variable
|
|
"""
|
|
if self.verbose:
|
|
# Display the header
|
|
text = 'INITIATING LINEAR REGRESSION MODEL PREDICTION'
|
|
print('\n' + '-' * len(text))
|
|
print(text)
|
|
print('-' * len(text) + '\n')
|
|
|
|
# Validate required columns exist for prediction
|
|
missing_features = [c for c in self.variable_columns if c not in input_data.columns]
|
|
if missing_features:
|
|
raise ValueError(
|
|
'Prediction data is missing required feature columns: '
|
|
+ ', '.join(missing_features)
|
|
)
|
|
|
|
X_test = input_data[self.variable_columns]
|
|
if self.verbose:
|
|
print(f'X_test shape: {X_test.shape}')
|
|
|
|
# Create polynomial features
|
|
if self.degree > 1:
|
|
X_test = self.create_poly_features(X_test, fit=False)
|
|
X_test = X_test.replace([np.inf, -np.inf], np.nan)
|
|
|
|
model_features = self.regr.feature_names_in_
|
|
# Ensure required model features exist
|
|
missing_model_features = [c for c in model_features if c not in X_test.columns]
|
|
if missing_model_features:
|
|
raise ValueError(
|
|
'Prediction data is missing required transformed feature columns: '
|
|
+ ', '.join(missing_model_features)
|
|
)
|
|
|
|
# Get mask of rows with any NaN values
|
|
mask = X_test.notna().all(axis=1)
|
|
X_test = X_test[mask]
|
|
if self.verbose:
|
|
print(f'X_test shape after removing NaN rows: {X_test.shape}')
|
|
|
|
# For simplicity and to keep behavior consistent with prior tests, drop NaN rows before prediction
|
|
y_pred = self.regr.predict(X_test[model_features])
|
|
|
|
if self.clipping:
|
|
for i in range(len(y_pred)):
|
|
if y_pred[i] > self.clipping['max']:
|
|
y_pred[i] = self.q3_target
|
|
elif y_pred[i] < self.clipping['min']:
|
|
y_pred[i] = self.q1_target
|
|
|
|
return y_pred
|