306 lines
12 KiB
Python
306 lines
12 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 _check_model_features(self, columns: pd.Index, model_features: list) -> None:
|
|
"""Raise when columns lack any feature the regressor was fitted with."""
|
|
missing_model_features = [c for c in model_features if c not in columns]
|
|
if missing_model_features:
|
|
raise ValueError(
|
|
'Prediction data is missing required transformed feature columns: '
|
|
+ ', '.join(missing_model_features)
|
|
)
|
|
|
|
def _apply_clipping(self, predictions: np.ndarray) -> None:
|
|
"""Replace in place predictions above max with Q3 and below min with Q1."""
|
|
for i in range(len(predictions)):
|
|
# Each limit is optional on its own (SIENTIAPDE-2219)
|
|
if 'max' in self.clipping and predictions[i] > self.clipping['max']:
|
|
predictions[i] = self.q3_target
|
|
elif 'min' in self.clipping and predictions[i] < self.clipping['min']:
|
|
predictions[i] = self.q1_target
|
|
|
|
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 _fittable_rows(
|
|
self, features: pd.DataFrame, target: pd.Series
|
|
) -> tuple[pd.DataFrame, pd.Series]:
|
|
"""
|
|
Keep only what LinearRegression can fit: inf becomes NaN, all-null columns are dropped
|
|
(degree 1 only) and rows with a NaN feature or a NaN target are removed.
|
|
|
|
Raises:
|
|
ValueError: If no rows are left.
|
|
"""
|
|
# Handle infinite values
|
|
features = features.astype(float).replace([np.inf, -np.inf], np.nan)
|
|
|
|
# Drop columns with all null values. With degree > 1 they are kept (PolynomialFeatures
|
|
# needs every variable column), so the row mask below leaves no rows and fit raises.
|
|
if self.degree == 1:
|
|
all_null_cols = features.columns[features.isnull().all()].tolist()
|
|
features = features.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.')
|
|
|
|
# Rows with NaN in any feature or in the target cannot be fitted (SIENTIAPDE-2215).
|
|
# Masked before PolynomialFeatures, which rejects NaN.
|
|
mask = features.notna().all(axis=1) & target.notna()
|
|
features = features[mask]
|
|
target = target[mask]
|
|
if features.empty:
|
|
raise ValueError(
|
|
'No training rows left after removing rows with NaN features or NaN target.'
|
|
)
|
|
if self.verbose:
|
|
print(f'X_train shape after removing NaN rows: {features.shape}')
|
|
print(f'y_train shape after removing NaN rows: {target.shape}')
|
|
return features, target
|
|
|
|
def fit(self, input_data: pd.DataFrame) -> 'LinearRegressionModel':
|
|
"""
|
|
Function to fit the model. Rows with NaN features or a NaN target are left out of the fit.
|
|
|
|
Args:
|
|
input_data (pandas.DataFrame): The data used to fit the Linear Regression model
|
|
|
|
Returns:
|
|
LinearRegressionModel: The prediction model object
|
|
|
|
Raises:
|
|
ValueError: If target_variable is unset, a feature column is missing, or no
|
|
training rows are left after removing rows with NaN features or NaN target.
|
|
KeyError: If the target column is missing from input_data.
|
|
"""
|
|
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}')
|
|
|
|
X_train, y_train = self._fittable_rows(X_train, y_train)
|
|
|
|
# Create polynomial features
|
|
if self.degree > 1:
|
|
X_train = self.create_poly_features(X_train, fit=True)
|
|
|
|
self.q1_target = float(y_train.quantile(0.25))
|
|
self.q3_target = float(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: One prediction per input row; NaN where any model input is NaN
|
|
"""
|
|
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)
|
|
)
|
|
|
|
# astype(float) first: laborious sends object columns with None, and replace() on
|
|
# object dtype raises a pandas downcasting FutureWarning on every call.
|
|
X_test = input_data[self.variable_columns].astype(float).replace([np.inf, -np.inf], np.nan)
|
|
if self.verbose:
|
|
print(f'X_test shape: {X_test.shape}')
|
|
|
|
model_features = list(self.regr.feature_names_in_)
|
|
# Rows with any NaN model input get a NaN prediction, so the output keeps the input
|
|
# length (SIENTIAPDE-2215). With degree > 1 the mask uses the raw variable columns,
|
|
# because PolynomialFeatures rejects NaN.
|
|
if self.degree > 1:
|
|
row_mask = X_test.notna().all(axis=1)
|
|
else:
|
|
self._check_model_features(X_test.columns, model_features)
|
|
row_mask = X_test[model_features].notna().all(axis=1)
|
|
|
|
y_pred = np.full(len(input_data), np.nan)
|
|
if not row_mask.any():
|
|
return y_pred
|
|
|
|
x_valid = X_test[row_mask]
|
|
if self.degree > 1:
|
|
x_valid = self.create_poly_features(x_valid, fit=False)
|
|
self._check_model_features(x_valid.columns, model_features)
|
|
if self.verbose:
|
|
print(f'X_test shape after removing NaN rows: {x_valid.shape}')
|
|
|
|
valid_pred = self.regr.predict(x_valid[model_features])
|
|
if self.clipping:
|
|
self._apply_clipping(valid_pred)
|
|
|
|
y_pred[row_mask.to_numpy()] = valid_pred
|
|
return y_pred
|