Initial commit
This commit is contained in:
vitor-aignosi
265
models/linear_regression/linear_regression_model.py
Normal file
265
models/linear_regression/linear_regression_model.py
Normal file
@@ -0,0 +1,265 @@
|
||||
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
|
||||
Reference in New Issue
Block a user