SIENTIAPDE-1430: Implement advanced model training capabilities and enhanced data preprocessing. This includes support for Polynomial Regression with configurable degree and interaction terms, flexible per-variable lag configurations, and new data filtering options by date range and removed intervals. Comprehensive business validations are now enforced for all parameters, and MLflow logging has been extended to capture these detailed configurations. Additionally, Reduced Coulomb Energy (RCE) metrics are added for drift detection, with a new changelog documenting all pipeline parameter updates.

PIPELINE_PARAMS_CHANGELOG.md

@@ -0,0 +1,224 @@
+# Changelog de Parâmetros do Pipeline de Treinamento
+
+Este documento descreve as alterações nos parâmetros de entrada do pipeline Temporal para treinamento de modelos.
+
+## Resumo das Alterações
+
+### Parâmetros ALTERADOS (Breaking Changes)
+
+| Parâmetro | Tipo Anterior | Tipo Novo | Descrição |
+|-----------|---------------|-----------|-----------|
+| `lag_train` | `int` | `dict[str, int]` | Agora é um dicionário com lag por variável. Ex: `{"var1": 2, "var2": 3}` |
+| `lag_val` | `int` | `dict[str, int]` | Agora é um dicionário com lag por variável. Ex: `{"var1": 1, "var2": 1}` |
+
+### Parâmetros NOVOS (Obrigatórios)
+
+| Parâmetro | Tipo | Descrição | Valores Válidos |
+|-----------|------|-----------|-----------------|
+| `model_name` | `str` | Nome do tipo de modelo | `"Linear Regression"`, `"Polynomial Regression"` |
+| `degree` | `int` | Grau do polinômio (1 = linear) | `>= 1` |
+| `interaction_only` | `bool` | Apenas termos de interação para polinomial | `true`, `false` |
+| `nan_treatment` | `str` | Tratamento de valores NaN | `"drop"`, `"linear interpolation"`, `"fill linear"` |
+| `scaler_name` | `str` | Nome do scaler a usar | `"Standard Scaler"`, `"None"` |
+
+### Parâmetros NOVOS (Opcionais)
+
+| Parâmetro | Tipo | Descrição | Default |
+|-----------|------|-----------|---------|
+| `start_date` | `str \| null` | Data inicial para filtrar dados (formato: `"YYYY-MM-DD HH:MM:SS"`) | `null` |
+| `end_date` | `str \| null` | Data final para filtrar dados (formato: `"YYYY-MM-DD HH:MM:SS"`) | `null` |
+| `support_filters` | `dict \| null` | Filtros customizados por variável | `{}` |
+
+---
+
+## Exemplo de Input Completo
+
+### Formato ANTERIOR (não funciona mais):
+
+```json
+{
+  "experiment_run_id": 123,
+  "target_variable": "temperatura",
+  "variable_columns": ["pressao", "umidade", "velocidade"],
+  "lag_train": 2,
+  "lag_val": 1,
+  "rem_static_win": true,
+  "low_lim": {"pressao": 0, "umidade": 0, "velocidade": 0},
+  "upp_lim": {"pressao": 100, "umidade": 100, "velocidade": 50},
+  "window": 0,
+  "use_scaler": true,
+  "include_ar": false,
+  "bucket_name": "training-data",
+  "file_name": "dataset.csv",
+  "line_separator": ";",
+  "decimal_separator": ",",
+  "train_size": 80,
+  "shuffle": false,
+  "experiment_name": "modelo-temperatura",
+  "removed_intervals": []
+}
+```
+
+### Formato NOVO (obrigatório):
+
+```json
+{
+  "experiment_run_id": 123,
+  "target_variable": "temperatura",
+  "variable_columns": ["pressao", "umidade", "velocidade"],
+  
+  "lag_train": {
+    "pressao": 2,
+    "umidade": 2,
+    "velocidade": 2
+  },
+  "lag_val": {
+    "pressao": 1,
+    "umidade": 1,
+    "velocidade": 1
+  },
+  
+  "rem_static_win": true,
+  "low_lim": {"pressao": 0, "umidade": 0, "velocidade": 0},
+  "upp_lim": {"pressao": 100, "umidade": 100, "velocidade": 50},
+  "window": 0,
+  "use_scaler": true,
+  "include_ar": false,
+  "bucket_name": "training-data",
+  "file_name": "dataset.csv",
+  "line_separator": ";",
+  "decimal_separator": ",",
+  "train_size": 80,
+  "shuffle": false,
+  "experiment_name": "modelo-temperatura",
+  "removed_intervals": [],
+  
+  "model_name": "Linear Regression",
+  "degree": 1,
+  "interaction_only": false,
+  "nan_treatment": "drop",
+  "scaler_name": "Standard Scaler",
+  
+  "start_date": null,
+  "end_date": null,
+  "support_filters": {}
+}
+```
+
+---
+
+## Exemplo para Polynomial Regression
+
+```json
+{
+  "experiment_run_id": 124,
+  "target_variable": "temperatura",
+  "variable_columns": ["pressao", "umidade"],
+  
+  "lag_train": {
+    "pressao": 0,
+    "umidade": 0
+  },
+  "lag_val": {
+    "pressao": 0,
+    "umidade": 0
+  },
+  
+  "rem_static_win": false,
+  "low_lim": {"pressao": 0, "umidade": 0},
+  "upp_lim": {"pressao": 100, "umidade": 100},
+  "window": 0,
+  "use_scaler": true,
+  "include_ar": false,
+  "bucket_name": "training-data",
+  "file_name": "dataset.csv",
+  "line_separator": ";",
+  "decimal_separator": ",",
+  "train_size": 80,
+  "shuffle": false,
+  "experiment_name": "modelo-polinomial",
+  "removed_intervals": [],
+  
+  "model_name": "Polynomial Regression",
+  "degree": 2,
+  "interaction_only": false,
+  "nan_treatment": "linear interpolation",
+  "scaler_name": "Standard Scaler",
+  
+  "start_date": "2024-01-01 00:00:00",
+  "end_date": "2024-12-31 23:59:59",
+  "support_filters": {}
+}
+```
+
+---
+
+## Exemplo com Intervalos Removidos
+
+```json
+{
+  "removed_intervals": [
+    ["2024-03-01 00:00:00", "2024-03-15 23:59:59"],
+    ["2024-06-01 00:00:00", "2024-06-30 23:59:59"]
+  ]
+}
+```
+
+---
+
+## Parâmetros Logados no MLflow
+
+Os seguintes parâmetros são agora logados no MLflow:
+
+| Parâmetro MLflow | Descrição |
+|------------------|-----------|
+| `model_name` | Nome do modelo (`Linear Regression` ou `Polynomial Regression`) |
+| `models_params` | `{"degree": int, "interaction_only": bool}` |
+| `target_variable` | Variável alvo |
+| `input_variables` | Lista de variáveis de entrada |
+| `nan_treatment` | Tratamento de NaN |
+| `lag_train` | Dicionário de lags para treino |
+| `lag_transform` | Dicionário de lags para transformação |
+| `static_threshold` | Threshold para janelas estáticas (1 ou null) |
+| `lower_limits` | Limites inferiores por variável |
+| `upper_limits` | Limites superiores por variável |
+| `scaler_name` | Nome do scaler |
+| `scaler_params` | Parâmetros do scaler (mean, variance) |
+| `include_ar` | Se inclui variável autoregressiva |
+| `train_size` | Proporção de treino (0.0 - 1.0) |
+| `test_size` | Proporção de teste (0.0 - 1.0) |
+| `start_date` | Data inicial (ou null) |
+| `end_date` | Data final (ou null) |
+| `removed_intervals` | Lista de intervalos removidos |
+| `retrain` | Sempre `false` para novos modelos |
+| `support_filters` | Filtros customizados |
+
+---
+
+## Validações de Negócio
+
+O sistema valida automaticamente:
+
+1. **`train_size`**: Deve estar entre 10 e 100
+2. **`variable_columns`**: Não pode estar vazio
+3. **`lag_train` / `lag_val`**: Todos os valores devem ser >= 0
+4. **`window`**: Deve ser >= 0
+5. **`degree`**: Deve ser >= 1
+6. **`nan_treatment`**: Deve ser `"drop"`, `"linear interpolation"` ou `"fill linear"`
+7. **`scaler_name`**: Deve ser `"Standard Scaler"` ou `"None"`
+8. **`model_name`**: Deve ser `"Linear Regression"` ou `"Polynomial Regression"`
+9. **`low_lim` / `upp_lim`**: Devem ter as mesmas chaves, e `low_lim[var] < upp_lim[var]`
+10. **`bucket_name` / `file_name` / `experiment_name`**: Não podem estar vazios
+11. **`degree` vs `model_name`**: Se `model_name` = "Polynomial Regression", `degree` deve ser >= 2; se "Linear Regression", `degree` deve ser = 1
+12. **`removed_intervals`**: Cada elemento deve ser lista/tupla com pelo menos 2 elementos (start, end)
+13. **`target_variable`**: Não pode estar vazio
+
+---
+
+## Arquivos Modificados
+
+- `model_manager/sientia/models.py` - `DataPreprocessor` e `LinearRegressionModel`
+- `model_manager/sientia/metrics.py` - Funções RCE adicionadas
+- `model_manager/utils/models/train_model_params.py` - Novos parâmetros
+- `model_manager/utils/repository/training_repository.py` - Uso dos novos parâmetros
+- `model_manager/utils/repository/model_repository.py` - Logging no MLflow

model_manager/sientia/metrics.py

@@ -26,3 +26,123 @@ def r2(real_data: pd.Series, predictions: pd.Series) -> float:
     Calculates the R2 score between the real data and the predictions.
     """
     return round(r2_score(real_data.astype(np.float64), predictions.astype(np.float64)), 2)
+
+
+def silverman_radius(data: np.ndarray) -> float:
+    """
+    Calculate the Silverman bandwidth (radius) for a given dataset.
+
+    Args:
+        data (np.ndarray): Input data (1D array)
+
+    Returns:
+        float: Silverman bandwidth (radius)
+    """
+    n = len(data)
+    sigma = np.std(data)
+    iqr = np.percentile(data, 75) - np.percentile(data, 25)
+    radius = 0.9 * min(sigma, iqr / 1.34) * n ** (-1 / 5)
+    return radius
+
+
+def rce_train(training_set: pd.DataFrame, radius: float) -> pd.DataFrame:
+    """
+    Get the Reduced Coulomb Energy (RCE) prototypes.
+
+    Args:
+        training_set (pd.DataFrame): The training set
+        radius (float): The radius of the RCE prototypes
+
+    Returns:
+        pd.DataFrame: The RCE prototypes
+    """
+    train_vectors = training_set.values
+
+    # Vectorized distance computation for the radius calculation
+    diff_vectors = train_vectors[:, np.newaxis] - train_vectors[np.newaxis, :]
+    distances = np.linalg.norm(diff_vectors, axis=-1)
+
+    # Non-parametric radius: Silverman Radius
+    radius = silverman_radius(distances.flatten())
+
+    # Initialize prototypes with the first vector
+    prototypes = [train_vectors[0]]
+
+    for vector in train_vectors[1:]:
+        # Vectorized distance check between current vector and all prototypes
+        distances_to_prototypes = np.linalg.norm(prototypes - vector, axis=1)
+
+        # If no prototype is close, add the current vector as a new prototype
+        if np.all(distances_to_prototypes > radius):
+            prototypes.append(vector)
+
+    return pd.DataFrame(prototypes)
+
+
+def rce_test(test_set: pd.DataFrame, prototypes: pd.DataFrame) -> pd.Series:
+    """
+    Get the signed Reduced Coulomb Energy (RCE) predictions.
+
+    Args:
+        test_set (pd.DataFrame): The test set
+        prototypes (pd.DataFrame): The RCE prototypes
+
+    Returns:
+        pd.Series: The signed distances to the closest prototype for each test vector
+    """
+    test_vectors = test_set.values
+    prototype_vectors = prototypes.values
+
+    # Vectorized computation of distances between test vectors and all prototypes
+    diff_vectors = test_vectors[:, np.newaxis] - prototype_vectors[np.newaxis, :]
+    distances = np.linalg.norm(diff_vectors, axis=-1)
+
+    # Find the closest prototype for each test vector
+    min_distances = np.min(distances, axis=1)
+    closest_prototypes = prototype_vectors[np.argmin(distances, axis=1)]
+
+    # Compute the signed distance for each test vector
+    signed_distances = np.sqrt(min_distances**2) * np.sign(
+        np.mean(test_vectors - closest_prototypes, axis=1)
+    )
+
+    return pd.Series(signed_distances)
+
+
+def rce_drift(reference_data: pd.DataFrame, real_data: pd.DataFrame, column: str) -> pd.Series:
+    """
+    Detect drift using the Reduced Coulomb Energy (RCE) method.
+
+    Args:
+        reference_data (pd.DataFrame): The reference data
+        real_data (pd.DataFrame): The real data
+        column (str): The target column to be analyzed. 'target' or 'prediction'
+
+    Returns:
+        pd.Series: Normalized drift distances
+    """
+    common_columns = list(set(reference_data.columns).intersection(real_data.columns))
+    reference_data = reference_data[common_columns]
+    real_data = real_data[common_columns]
+
+    # Get prototypes
+    if column == 'target':
+        prototypes = rce_train(reference_data.drop(columns=['prediction']), 0.1)
+    else:
+        prototypes = rce_train(reference_data.drop(columns=['target']), 0.1)
+
+    # Distances to prototypes
+    if column == 'target':
+        distances_train = rce_test(reference_data.drop(columns=['prediction']), prototypes)
+        distances_test = rce_test(real_data.drop(columns=['prediction']), prototypes)
+    else:
+        distances_train = rce_test(reference_data.drop(columns=['target']), prototypes)
+        distances_test = rce_test(real_data.drop(columns=['target']), prototypes)
+
+    # Find the maximum absolute distance in the training set
+    max_abs_distance = max(abs(distances_train.max()), abs(distances_train.min()))
+
+    # Normalize while preserving sign
+    distances = distances_test / max_abs_distance
+
+    return distances

model_manager/sientia/models.py

@@ -6,19 +6,24 @@ from sientia_do.operations.df_preprocessor import create_features, limit_dataset
 from sientia_do.timeseries.analyzer import TimeSeriesDiscontinuityAnalyzer
 from sklearn.base import BaseEstimator, TransformerMixin
 from sklearn.linear_model import LinearRegression
-from sklearn.preprocessing import StandardScaler
+from sklearn.preprocessing import PolynomialFeatures, StandardScaler
 
 DISCONTINUITY_TREATMENT = 'Discontinuity Treatment'
 LAG_SELECTION = 'Lag Selection'
+RANGE_SELECTION = 'Range Selection & Data Removal'
 STATIC_WINDOW_REMOVAL = 'Static Window Removal'
 DEFINE_VARIABLES_LIMITS = 'Define Variables Limits'
 NORMALIZATION = 'Normalization'
+FEATURE_CREATION = 'Feature Creation'
+LAG_CREATION = 'Lag Creation'
 
 
 class LinearRegressionModel(BaseEstimator, TransformerMixin):
     """
     Linear Regression Model for Time Series Analysis.
 
+    Supports both simple linear regression and polynomial regression.
+
     Thread-safety: This class is NOT thread-safe during fit() operations.
     Do not call fit() on the same instance from multiple threads simultaneously.
     After fitting, predict() is thread-safe for read-only operations.
@@ -36,6 +41,9 @@ class LinearRegressionModel(BaseEstimator, TransformerMixin):
         model_params: dict[str, Any] | None = None,
         clipping: dict[str, float] | None = None,
         weights: dict[str, float] | None = None,
+        degree: int = 1,
+        interaction_only: bool = False,
+        verbose: bool = False,
     ):
         """
         Linear Regression Model for Time Series Analysis
@@ -48,6 +56,9 @@ class LinearRegressionModel(BaseEstimator, TransformerMixin):
                 *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 (1 = linear, >1 = polynomial)
+            interaction_only (bool): If True, only interaction features are produced
+            verbose (bool): If True, print verbose output during fitting
 
         Returns:
             LinearRegressionModel: The prediction model object
@@ -60,6 +71,44 @@ class LinearRegressionModel(BaseEstimator, TransformerMixin):
         self.q1_target: float | None = None
         self.q3_target: float | None = None
         self.weights: dict[str, float] | None = weights
+        self.degree: int = degree
+        self.interaction_only: bool = interaction_only
+        self.verbose: bool = verbose
+        self.poly: PolynomialFeatures | None = None
+        self.poly_feature_names: list[str] | None = None
+
+    def create_poly_features(self, input_data: pd.DataFrame, fit: bool = False) -> pd.DataFrame:
+        """
+        Create polynomial features from input data.
+
+        Args:
+            input_data (pd.DataFrame): Input data with feature columns
+            fit (bool): If True, fit the PolynomialFeatures transformer
+
+        Returns:
+            pd.DataFrame: DataFrame with polynomial features
+        """
+        if self.degree <= 1:
+            return input_data
+
+        if fit:
+            self.poly = PolynomialFeatures(
+                degree=self.degree,
+                interaction_only=self.interaction_only,
+                include_bias=False,
+            )
+            poly_features = self.poly.fit_transform(input_data)
+            self.poly_feature_names = list(self.poly.get_feature_names_out(input_data.columns))
+        else:
+            if self.poly is None:
+                raise ValueError('PolynomialFeatures not fitted. Call fit() first.')
+            poly_features = self.poly.transform(input_data)
+
+        return pd.DataFrame(
+            poly_features,
+            columns=self.poly_feature_names,
+            index=input_data.index,
+        )
 
     def fit(self, input_data: pd.DataFrame) -> 'LinearRegressionModel':
         """
@@ -71,13 +120,52 @@ class LinearRegressionModel(BaseEstimator, TransformerMixin):
         Returns:
             LinearRegressionModel: The prediction model object
         """
-        assert self.variable_columns is not None, 'variable_columns must be set before fitting'
+        if not self.target_variable:
-        X_train = input_data[self.variable_columns]
+            raise ValueError('target_variable must be set before fitting')
-        y_train = input_data[self.target_variable]
+
+        # Infer variable_columns if not provided
+        if self.variable_columns is None:
+            self.variable_columns = [
+                col for col in input_data.columns if col != self.target_variable
+            ]
+
+        # Validate columns exist
+        missing_cols = [col for col in self.variable_columns if col not in input_data.columns]
+        if missing_cols:
+            raise ValueError(f'Columns not found in input data: {missing_cols}')
+
+        if self.target_variable not in input_data.columns:
+            raise ValueError(f'Target variable {self.target_variable} not found in input data')
+
+        X_train = input_data[self.variable_columns].copy()
+        y_train = input_data[self.target_variable].copy()
+
+        # Handle infinite values
+        X_train = X_train.replace([np.inf, -np.inf], np.nan)
+        y_train = y_train.replace([np.inf, -np.inf], np.nan)
+
+        # Remove rows with NaN
+        valid_mask = ~(X_train.isna().any(axis=1) | y_train.isna())
+        X_train = X_train[valid_mask]
+        y_train = y_train[valid_mask]
+
+        # Remove columns with all NaN values
+        cols_to_drop = X_train.columns[X_train.isna().all()].tolist()
+        if cols_to_drop:
+            if self.verbose:
+                print(f'Dropping columns with all NaN values: {cols_to_drop}')
+            X_train = X_train.drop(columns=cols_to_drop)
+            self.variable_columns = [c for c in self.variable_columns if c not in cols_to_drop]
 
         self.q1_target = y_train.quantile(0.25)
         self.q3_target = y_train.quantile(0.75)
 
+        # Apply polynomial features if degree > 1
+        if self.degree > 1:
+            X_train = self.create_poly_features(X_train, fit=True)
+            if self.verbose and self.poly_feature_names is not None:
+                print(f'Created {len(self.poly_feature_names)} polynomial features')
+
         # Fit the model
         self.regr.fit(X_train, y_train)
 
@@ -86,11 +174,16 @@ class LinearRegressionModel(BaseEstimator, TransformerMixin):
         round_intercept = np.round(self.regr.intercept_, 3)
 
         # Save the weights
-        weights = dict(zip(self.variable_columns, [float(c) for c in round_coef], strict=True))
+        feature_names = self.poly_feature_names if self.degree > 1 else self.variable_columns
+        assert feature_names is not None, 'feature_names should be set at this point'
+        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), **weights}
         self.weights = weights
 
+        if self.verbose and feature_names is not None:
+            print(f'Model fitted with {len(feature_names)} features')
+
         return self
 
     def predict(self, input_data: pd.DataFrame) -> np.ndarray:
@@ -104,7 +197,16 @@ class LinearRegressionModel(BaseEstimator, TransformerMixin):
         Returns:
             numpy.ndarray: The predicted target variable
         """
-        X_test = input_data[self.variable_columns]
+        assert self.variable_columns is not None, 'variable_columns must be set before predict'
+        X_test: pd.DataFrame = input_data[self.variable_columns].copy()
+
+        # Handle infinite values
+        X_test = X_test.replace([np.inf, -np.inf], np.nan)
+
+        # Apply polynomial features if degree > 1
+        if self.degree > 1:
+            X_test = self.create_poly_features(X_test, fit=False)
+
         y_pred = self.regr.predict(X_test)
 
         if self.clipping:
@@ -116,6 +218,15 @@ class LinearRegressionModel(BaseEstimator, TransformerMixin):
 
         return y_pred
 
+    def get_regressor(self) -> LinearRegression:
+        """
+        Get the underlying LinearRegression model.
+
+        Returns:
+            LinearRegression: The sklearn LinearRegression model
+        """
+        return self.regr
+
 
 class DataPreprocessor(BaseEstimator, TransformerMixin):
     """
@@ -141,6 +252,9 @@ class DataPreprocessor(BaseEstimator, TransformerMixin):
         nan_treatment: str | None = None,
         lag_train: dict[str, int] | None = None,
         lag_transform: dict[str, int] | None = None,
+        start_date: str | None = None,
+        end_date: str | None = None,
+        removed_intervals: list[tuple[str, str]] | None = None,
         static_threshold: int | None = None,
         low_lim: dict[str, float] | None = None,
         upp_lim: dict[str, float] | None = None,
@@ -152,6 +266,7 @@ class DataPreprocessor(BaseEstimator, TransformerMixin):
         cross_operations: list[str] | None = None,
         created_lags: dict[str, int] | None = None,
         steps_order: list[str] | None = None,
+        verbose: bool = False,
     ):
         """
         Data Preprocessor for Time Series Analysis
@@ -161,11 +276,15 @@ class DataPreprocessor(BaseEstimator, TransformerMixin):
             target_variable (str): The target variable name
             input_columns (list): The input columns names in a list
             nan_treatment (str): The treatment for missing values \\
-                *Options: 'drop', 'fill linear'*
+                *Options: 'drop', 'fill linear', 'linear interpolation'*
             lag_train (dict): The lags for each variable to be applyed during training \\
                 *Format: {'variable_name': lag}*
             lag_transform (dict): The lags for each variable to be applyed during transformation \\
                 *Format: {'variable_name': lag}*
+            start_date (str): The start date for filtering data (format: 'YYYY-MM-DD HH:MM:SS')
+            end_date (str): The end date for filtering data (format: 'YYYY-MM-DD HH:MM:SS')
+            removed_intervals (list): List of tuples with intervals to remove from data \\
+                *Format: [('start_date', 'end_date'), ...]*
             static_threshold (int): The number of repeated values to be considered as static
             low_lim (dict): The lower limits for each variable \\
                 *Format: {'variable_name': limit}*
@@ -188,11 +307,13 @@ class DataPreprocessor(BaseEstimator, TransformerMixin):
             steps_order (list): The order of the steps to be executed in the pipeline \\
                 *Options for list:  'Discontinuity Treatment',
                                     'Lag Selection',
+                                    'Range Selection & Data Removal',
                                     'Static Window Removal',
                                     'Define Variables Limits',
                                     'Normalization',
                                     'Feature Creation',
                                     'Lag Creation'*
+            verbose (bool): If True, print verbose output during preprocessing
 
         Returns:
             DataPreprocessor: The data preprocessor object
@@ -203,6 +324,9 @@ class DataPreprocessor(BaseEstimator, TransformerMixin):
         self.nan_treatment = nan_treatment
         self.lag_train = lag_train if lag_train else {}
         self.lag_transform = lag_transform if lag_transform else {}
+        self.start_date = start_date
+        self.end_date = end_date
+        self.removed_intervals = removed_intervals if removed_intervals else []
         self.ar_var = ar_var
         self.self_operations = self_operations
         self.cross_operations = cross_operations
@@ -214,6 +338,8 @@ class DataPreprocessor(BaseEstimator, TransformerMixin):
         self.scaler_name = scaler_name
         self.scaler_params = scaler_params
         self.feature_names_order: list[str] = []  # Initialize to avoid AttributeError
+        self.verbose = verbose
+        self._fitted_feature_order: list[str] | None = None  # Track feature order after fit
 
         if self.scaler_name == 'Standard Scaler':
             self.scaler = StandardScaler()
@@ -226,11 +352,12 @@ class DataPreprocessor(BaseEstimator, TransformerMixin):
         possible_steps = [
             DISCONTINUITY_TREATMENT,
             LAG_SELECTION,
+            RANGE_SELECTION,
             STATIC_WINDOW_REMOVAL,
             DEFINE_VARIABLES_LIMITS,
             NORMALIZATION,
-            'Feature Creation',
+            FEATURE_CREATION,
-            'Lag Creation',
+            LAG_CREATION,
         ]
         self.steps_order = steps_order or possible_steps
         for step in possible_steps:
@@ -321,7 +448,61 @@ class DataPreprocessor(BaseEstimator, TransformerMixin):
             pandas.DataFrame: The treated data
         """
         if self.nan_treatment:
-            input_data = treat_nan(input_data, self.nan_treatment)
+            # Map 'linear interpolation' to 'fill linear' for compatibility
+            treatment = self.nan_treatment
+            if treatment == 'linear interpolation':
+                treatment = 'fill linear'
+            input_data = treat_nan(input_data, treatment)
+            if self.verbose:
+                print(f'Applied NaN treatment: {self.nan_treatment}')
+        return input_data
+
+    def range_selection(self, input_data: pd.DataFrame) -> pd.DataFrame:
+        """
+        Filter data by date range and remove specified intervals.
+
+        Args:
+            input_data (pandas.DataFrame): The input data with datetime index
+
+        Returns:
+            pandas.DataFrame: The filtered data
+        """
+        # Filter by start_date and end_date
+        if self.start_date:
+            try:
+                start = pd.to_datetime(self.start_date)
+                input_data = input_data[input_data.index >= start]
+                if self.verbose:
+                    print(f'Filtered data from start_date: {self.start_date}')
+            except (ValueError, TypeError):
+                pass  # Invalid date format, skip filtering
+
+        if self.end_date:
+            try:
+                end = pd.to_datetime(self.end_date)
+                input_data = input_data[input_data.index <= end]
+                if self.verbose:
+                    print(f'Filtered data to end_date: {self.end_date}')
+            except (ValueError, TypeError):
+                pass  # Invalid date format, skip filtering
+
+        # Remove specified intervals
+        if self.removed_intervals:
+            for interval in self.removed_intervals:
+                if len(interval) >= 2:
+                    try:
+                        interval_start = pd.to_datetime(interval[0])
+                        interval_end = pd.to_datetime(interval[1])
+                        mask = ~(
+                            (input_data.index >= interval_start)
+                            & (input_data.index <= interval_end)
+                        )
+                        input_data = input_data[mask]
+                        if self.verbose:
+                            print(f'Removed interval: {interval[0]} to {interval[1]}')
+                    except (ValueError, TypeError):
+                        pass  # Invalid date format, skip this interval
+
         return input_data
 
     def lag_selection(self, input_data: pd.DataFrame, lag_dict: dict) -> pd.DataFrame:
@@ -469,6 +650,10 @@ class DataPreprocessor(BaseEstimator, TransformerMixin):
             if step == LAG_SELECTION:
                 data_treat = self.lag_selection(data_treat, self.lag_train)
 
+            # Range Selection & Data Removal
+            if step == RANGE_SELECTION:
+                data_treat = self.range_selection(data_treat)
+
             # Static Window Treatment
             if step == STATIC_WINDOW_REMOVAL:
                 data_treat = self.treat_static_windows(data_treat)
@@ -493,6 +678,9 @@ class DataPreprocessor(BaseEstimator, TransformerMixin):
                         'variance': round(variance, 3),
                     }
 
+        # Store fitted feature order for predict method
+        self._fitted_feature_order = list(existing_columns)
+
         return self
 
     def transform(self, x: pd.DataFrame) -> pd.DataFrame:
@@ -525,6 +713,10 @@ class DataPreprocessor(BaseEstimator, TransformerMixin):
             if step == LAG_SELECTION:
                 data_treat = self.lag_selection(data_treat, self.lag_transform)
 
+            # Range Selection & Data Removal (typically skipped in transform)
+            if step == RANGE_SELECTION:
+                data_treat = self.range_selection(data_treat)
+
             # Static Window Treatment
             if step == STATIC_WINDOW_REMOVAL:
                 data_treat = self.treat_static_windows(data_treat)
@@ -540,11 +732,11 @@ class DataPreprocessor(BaseEstimator, TransformerMixin):
                 data_treat[feature_cols] = self.scaler.transform(data_treat[feature_cols])
 
             # Feature Creation
-            if step == 'Feature Creation':
+            if step == FEATURE_CREATION:
                 data_treat = self.create_features(data_treat)
 
             # Lag Creation
-            if step == 'Lag Creation':
+            if step == LAG_CREATION:
                 # Autoregressive Variable
                 if self.input_columns is not None and self.ar_var in self.input_columns:
                     data_treat = self.create_ar(data_treat)
@@ -553,3 +745,33 @@ class DataPreprocessor(BaseEstimator, TransformerMixin):
                 data_treat = self.create_lags(data_treat)
 
         return data_treat
+
+    def predict(self, x: pd.DataFrame) -> pd.DataFrame:
+        """
+        Transform data for prediction (removes target variable).
+
+        This method is a wrapper around transform() that:
+        1. Transforms the input data
+        2. Removes the target variable column
+        3. Ensures features are in the same order as during fit
+
+        Args:
+            x (pandas.DataFrame): The input data
+
+        Returns:
+            pandas.DataFrame: The transformed data without target variable,
+                            with features in the same order as during fit
+        """
+        data_treat = self.transform(x)
+
+        # Remove target variable if present
+        if self.target_variable in data_treat.columns:
+            data_treat = data_treat.drop(columns=self.target_variable)
+
+        # Ensure features are in the same order as during fit
+        if self._fitted_feature_order is not None:
+            # Filter to only include columns that exist in both
+            available_cols = [c for c in self._fitted_feature_order if c in data_treat.columns]
+            data_treat = data_treat[available_cols]
+
+        return data_treat

model_manager/utils/models/train_model_params.py

@@ -17,8 +17,8 @@ class TrainModelParams:
 
     Attributes:
         variable_columns (list[str]): List of variable column names to use as features.
-        lag_train (int): Number of lags to apply during training phase.
+        lag_train (dict[str, int]): Dictionary of lags per variable for training phase.
-        lag_val (int): Number of lags to apply during validation phase.
+        lag_val (dict[str, int]): Dictionary of lags per variable for validation phase.
         target_variable (str): Name of the target variable to predict.
         rem_static_win (bool): Whether to remove static windows from data.
         low_lim (dict[str, float]): Dictionary of lower limits for each variable.
@@ -35,11 +35,19 @@ class TrainModelParams:
         experiment_run_id (int): Unique identifier for the experiment run.
         experiment_name (str): Name of the experiment for tracking.
         removed_intervals (list): List of time intervals to remove from the data.
+        model_name (str): Name of the model type ('Linear Regression' or 'Polynomial Regression').
+        degree (int): Degree of polynomial features (1 for linear, >1 for polynomial).
+        interaction_only (bool): If True, only interaction features are produced for polynomial.
+        nan_treatment (str): Treatment for NaN values ('drop' or 'linear interpolation').
+        start_date (str | None): Start date for filtering data.
+        end_date (str | None): End date for filtering data.
+        scaler_name (str): Name of the scaler to use ('Standard Scaler' or 'None').
+        support_filters (dict): Custom support filters per variable.
     """
 
     variable_columns: list[str]
-    lag_train: int
+    lag_train: dict[str, int]
-    lag_val: int
+    lag_val: dict[str, int]
     target_variable: str
     rem_static_win: bool
     low_lim: dict[str, float]
@@ -56,6 +64,14 @@ class TrainModelParams:
     experiment_run_id: int
     experiment_name: str
     removed_intervals: list
+    model_name: str
+    degree: int
+    interaction_only: bool
+    nan_treatment: str
+    start_date: str | None
+    end_date: str | None
+    scaler_name: str
+    support_filters: dict
 
     @classmethod
     def from_dict(cls, data: dict[str, Any]) -> 'TrainModelParams':
@@ -83,8 +99,8 @@ class TrainModelParams:
             variable_columns=cls._check_none(
                 data.get('variable_columns'), list, 'variable_columns'
             ),
-            lag_train=cls._check_none(data.get('lag_train'), int, 'lag_train'),
+            lag_train=cls._check_none(data.get('lag_train'), dict, 'lag_train'),
-            lag_val=cls._check_none(data.get('lag_val'), int, 'lag_val'),
+            lag_val=cls._check_none(data.get('lag_val'), dict, 'lag_val'),
             target_variable=cls._check_none(data.get('target_variable'), str, 'target_variable'),
             rem_static_win=cls._check_none(data.get('rem_static_win'), bool, 'rem_static_win'),
             low_lim=cls._check_none(data.get('low_lim'), dict, 'low_lim'),
@@ -107,6 +123,17 @@ class TrainModelParams:
             removed_intervals=cls._check_type(
                 data.get('removed_intervals'), list, 'removed_intervals'
             ),
+            model_name=cls._check_none(data.get('model_name'), str, 'model_name'),
+            degree=cls._check_none(data.get('degree'), int, 'degree'),
+            interaction_only=cls._check_none(
+                data.get('interaction_only'), bool, 'interaction_only'
+            ),
+            nan_treatment=cls._check_none(data.get('nan_treatment'), str, 'nan_treatment'),
+            start_date=cls._check_type(data.get('start_date'), str, 'start_date'),
+            end_date=cls._check_type(data.get('end_date'), str, 'end_date'),
+            scaler_name=cls._check_none(data.get('scaler_name'), str, 'scaler_name'),
+            support_filters=cls._check_type(data.get('support_filters'), dict, 'support_filters')
+            or {},
         )
 
     @staticmethod
@@ -172,25 +199,81 @@ class TrainModelParams:
         Raises:
             ValueError: If any business rule is violated
         """
-        # Validate train_size range (10-100%)
+        self._validate_numeric_ranges()
+        self._validate_model_params()
+        self._validate_intervals_and_dates()
+        self._validate_limits()
+        self._validate_required_strings()
+
+    def _validate_numeric_ranges(self) -> None:
+        """Validate numeric parameters are within acceptable ranges."""
         if not 10 <= self.train_size <= 100:
             raise ValueError(f'train_size must be between 10 and 100, got {self.train_size}')
 
-        # Validate variable_columns is not empty
         if not self.variable_columns:
             raise ValueError('variable_columns cannot be empty')
 
-        # Validate positive integers
+        for var, lag in self.lag_train.items():
-        if self.lag_train < 0:
+            if lag < 0:
-            raise ValueError(f'lag_train must be positive, got {self.lag_train}')
+                raise ValueError(f'lag_train for {var} must be non-negative, got {lag}')
 
-        if self.lag_val < 0:
+        for var, lag in self.lag_val.items():
-            raise ValueError(f'lag_val must be positive, got {self.lag_val}')
+            if lag < 0:
+                raise ValueError(f'lag_val for {var} must be non-negative, got {lag}')
 
         if self.window < 0:
-            raise ValueError(f'window must be positive, got {self.window}')
+            raise ValueError(f'window must be non-negative, got {self.window}')
 
-        # Validate low_lim and upp_lim consistency
+    def _validate_model_params(self) -> None:
+        """Validate model-related parameters."""
+        if self.degree < 1:
+            raise ValueError(f'degree must be at least 1, got {self.degree}')
+
+        valid_nan_treatments = ['drop', 'linear interpolation', 'fill linear']
+        if self.nan_treatment not in valid_nan_treatments:
+            raise ValueError(
+                f'nan_treatment must be one of {valid_nan_treatments}, got {self.nan_treatment}'
+            )
+
+        valid_scalers = ['Standard Scaler', 'None']
+        if self.scaler_name not in valid_scalers:
+            raise ValueError(f'scaler_name must be one of {valid_scalers}, got {self.scaler_name}')
+
+        valid_models = ['Linear Regression', 'Polynomial Regression']
+        if self.model_name not in valid_models:
+            raise ValueError(f'model_name must be one of {valid_models}, got {self.model_name}')
+
+        if self.model_name == 'Polynomial Regression' and self.degree < 2:
+            raise ValueError(
+                f'degree must be at least 2 for Polynomial Regression, got {self.degree}'
+            )
+
+        if self.model_name == 'Linear Regression' and self.degree != 1:
+            raise ValueError(f'degree must be 1 for Linear Regression, got {self.degree}')
+
+    def _validate_intervals_and_dates(self) -> None:
+        """Validate removed_intervals format and date parameters."""
+        if self.removed_intervals:
+            for i, interval in enumerate(self.removed_intervals):
+                if not isinstance(interval, (list, tuple)):
+                    raise ValueError(
+                        f'removed_intervals[{i}] must be a list or tuple, '
+                        f'got {type(interval).__name__}'
+                    )
+                if len(interval) < 2:
+                    raise ValueError(
+                        f'removed_intervals[{i}] must have at least 2 elements (start, end), '
+                        f'got {len(interval)}'
+                    )
+
+        if self.start_date is not None and not isinstance(self.start_date, str):
+            raise TypeError(f'start_date must be a string, got {type(self.start_date).__name__}')
+
+        if self.end_date is not None and not isinstance(self.end_date, str):
+            raise TypeError(f'end_date must be a string, got {type(self.end_date).__name__}')
+
+    def _validate_limits(self) -> None:
+        """Validate low_lim and upp_lim consistency."""
         if set(self.low_lim.keys()) != set(self.upp_lim.keys()):
             raise ValueError(
                 f'low_lim and upp_lim must have the same keys. '
@@ -198,7 +281,6 @@ class TrainModelParams:
                 f'upp_lim keys: {set(self.upp_lim.keys())}'
             )
 
-        # Validate that low_lim < upp_lim for each variable
         for var in self.low_lim:
             if self.low_lim[var] >= self.upp_lim[var]:
                 raise ValueError(
@@ -206,13 +288,16 @@ class TrainModelParams:
                     f'Got low_lim={self.low_lim[var]}, upp_lim={self.upp_lim[var]}'
                 )
 
-        # Validate bucket_name and file_name are not empty
+    def _validate_required_strings(self) -> None:
+        """Validate required string fields are not empty."""
+        if not self.target_variable.strip():
+            raise ValueError('target_variable cannot be empty or whitespace')
+
         if not self.bucket_name.strip():
             raise ValueError('bucket_name cannot be empty or whitespace')
 
         if not self.file_name.strip():
             raise ValueError('file_name cannot be empty or whitespace')
 
-        # Validate experiment_name is not empty
         if not self.experiment_name.strip():
             raise ValueError('experiment_name cannot be empty or whitespace')

model_manager/utils/repository/model_repository.py

@@ -183,8 +183,6 @@ class ModelRepository:
             raise ValueError(error_msg)
 
         # Prepare parameters
-        train_test_split = f'{data.params.train_size}-{100 - data.params.train_size}'
-
         interval_strs = [
             (str(interval[0]), str(interval[1]))
             for interval in (data.params.removed_intervals or [])
@@ -197,19 +195,31 @@ class ModelRepository:
             run_name=data.run_name, description=data.params.experiment_name
         ):
             # Log model parameters
-            self.model_serving.log_param('model_type', 'Linear Regression')
+            self.model_serving.log_param('model_name', data.params.model_name)
+            self.model_serving.log_param(
+                'models_params',
+                {'degree': data.params.degree, 'interaction_only': data.params.interaction_only},
+            )
             self.model_serving.log_param('target_variable', data.params.target_variable)
             self.model_serving.log_param('input_variables', data.params.variable_columns)
+            self.model_serving.log_param('nan_treatment', data.params.nan_treatment)
             self.model_serving.log_param('lag_train', data.params.lag_train)
-            self.model_serving.log_param('lag_val', data.params.lag_val)
+            self.model_serving.log_param('lag_transform', data.params.lag_val)
-            self.model_serving.log_param('ma', data.params.window)
+            self.model_serving.log_param(
-            self.model_serving.log_param('low_lim', data.params.low_lim)
+                'static_threshold', 1 if data.params.rem_static_win else None
-            self.model_serving.log_param('upp_lim', data.params.upp_lim)
+            )
-            self.model_serving.log_param('normalized', data.scaler_dict)
+            self.model_serving.log_param('lower_limits', data.params.low_lim)
-            self.model_serving.log_param('ar', data.params.include_ar)
+            self.model_serving.log_param('upper_limits', data.params.upp_lim)
-            self.model_serving.log_param('Train_test_split', train_test_split)
+            self.model_serving.log_param('scaler_name', data.params.scaler_name)
-            self.model_serving.log_param('Removed_intervals', interval_strs)
+            self.model_serving.log_param('scaler_params', data.scaler_dict)
-            self.model_serving.log_param('Retrain', False)
+            self.model_serving.log_param('include_ar', data.params.include_ar)
+            self.model_serving.log_param('train_size', round(data.params.train_size / 100, 2))
+            self.model_serving.log_param('test_size', round(1 - (data.params.train_size / 100), 2))
+            self.model_serving.log_param('start_date', data.params.start_date)
+            self.model_serving.log_param('end_date', data.params.end_date)
+            self.model_serving.log_param('removed_intervals', interval_strs)
+            self.model_serving.log_param('retrain', False)
+            self.model_serving.log_param('support_filters', data.params.support_filters)
 
             # Log evaluation metrics
             self.model_serving.log_metric('MSE', data.mse_val)

model_manager/utils/repository/training_repository.py

@@ -88,6 +88,8 @@ class TrainingRepository:
         regr = LinearRegressionModel(
             target_variable=params.target_variable,
             variable_columns=params.variable_columns,
+            degree=params.degree,
+            interaction_only=params.interaction_only,
         )
 
         regr.fit(data_train)
@@ -236,20 +238,27 @@ class TrainingRepository:
         Returns:
             DataPreprocessor: Configured preprocessor ready for fitting
         """
-        # Create lag dictionaries for each variable
+        # Convert removed_intervals to list of tuples if needed
-        lag_train_dict = dict.fromkeys(params.variable_columns, params.lag_train)
+        removed_intervals = None
-        lag_val_dict = dict.fromkeys(params.variable_columns, params.lag_val)
+        if params.removed_intervals:
+            removed_intervals = [
+                (interval[0], interval[1]) if isinstance(interval, (list, tuple)) else interval
+                for interval in params.removed_intervals
+            ]
 
         return DataPreprocessor(
             target_variable=params.target_variable,
             input_columns=params.variable_columns,
-            lag_train=lag_train_dict,
+            nan_treatment=params.nan_treatment,
-            lag_transform=lag_val_dict,
+            lag_train=params.lag_train,
+            lag_transform=params.lag_val,
+            start_date=params.start_date,
+            end_date=params.end_date,
+            removed_intervals=removed_intervals,
             static_threshold=1 if params.rem_static_win else None,
             low_lim=params.low_lim,
             upp_lim=params.upp_lim,
-            window=params.window,
+            scaler_name=params.scaler_name,
-            scaler_name='Standard Scaler' if params.use_scaler else 'None',
             scaler_params={} if params.use_scaler else None,
             ar_var=params.target_variable if params.include_ar else None,
         )
@@ -279,10 +288,17 @@ class TrainingRepository:
         coefficients = regr.regr.coef_
         intercept = regr.regr.intercept_
 
+        # Get feature names - for polynomial models, use poly_feature_names
+        if params.degree > 1 and regr.poly_feature_names:
+            feature_names = regr.poly_feature_names
+        else:
+            feature_names = params.variable_columns
+
         # Create coefficients dictionary
         coefficients_dict = {}
-        for i, var in enumerate(params.variable_columns):
+        for i, var in enumerate(feature_names):
-            coefficients_dict[var] = float(coefficients[i])
+            if i < len(coefficients):
+                coefficients_dict[var] = float(coefficients[i])
 
         # Create equation string
         equation_parts = [f'{coef:.6f} * {var}' for var, coef in coefficients_dict.items()]
@@ -300,5 +316,8 @@ class TrainingRepository:
             'intercept': float(intercept),
             'equation_string': equation_string,
             'latex_equation': latex_equation,
-            'model_type': 'Linear Regression',
+            'model_type': params.model_name,
+            'degree': params.degree,
+            'interaction_only': params.interaction_only,
+            'original_features': params.variable_columns,
         }

tests/sientia/test_models.py

@@ -81,13 +81,15 @@ def test_linear_regression_model_fit():
 
 
 def test_linear_regression_model_fit_without_variable_columns():
-    """Test LinearRegressionModel fit raises AssertionError without variable_columns."""
+    """Test LinearRegressionModel fit infers variable_columns when not set."""
     model = LinearRegressionModel(target_variable='target')
 
     data = pd.DataFrame({'var1': [1, 2, 3], 'target': [3, 5, 7]})
 
-    with raises(AssertionError, match='variable_columns must be set before fitting'):
+    # Model should infer variable_columns from data (all columns except target)
-        model.fit(data)
+    result = model.fit(data)
+    assert result is model
+    assert model.variable_columns == ['var1']
 
 
 def test_linear_regression_model_predict_without_clipping():
@@ -196,7 +198,7 @@ def test_data_preprocessor_init_with_custom_steps_order():
 
     assert 'Normalization' in preprocessor.steps_order
     assert 'Feature Creation' in preprocessor.steps_order
-    assert len(preprocessor.steps_order) == 7
+    assert len(preprocessor.steps_order) == 8  # Now includes RANGE_SELECTION step
 
 
 def test_data_preprocessor_get_scaler():

tests/utils/models/test_train_model_params.py

@@ -8,8 +8,8 @@ def valid_train_params_dict():
     """Create a valid dictionary for TrainModelParams."""
     return {
         'variable_columns': ['var1', 'var2'],
-        'lag_train': 5,
+        'lag_train': {'var1': 5, 'var2': 5},
-        'lag_val': 3,
+        'lag_val': {'var1': 3, 'var2': 3},
         'target_variable': 'target',
         'rem_static_win': True,
         'low_lim': {'var1': 0.0, 'var2': 1.0},
@@ -26,6 +26,14 @@ def valid_train_params_dict():
         'experiment_run_id': 1,
         'experiment_name': 'test_experiment',
         'removed_intervals': [],
+        'model_name': 'Linear Regression',
+        'degree': 1,
+        'interaction_only': False,
+        'nan_treatment': 'drop',
+        'start_date': None,
+        'end_date': None,
+        'scaler_name': 'Standard Scaler',
+        'support_filters': {},
     }
 
 
@@ -36,8 +44,8 @@ def test_train_model_params_from_dict_success(valid_train_params_dict):
     params = TrainModelParams.from_dict(valid_train_params_dict)
 
     assert params.variable_columns == ['var1', 'var2']
-    assert params.lag_train == 5
+    assert params.lag_train == {'var1': 5, 'var2': 5}
-    assert params.lag_val == 3
+    assert params.lag_val == {'var1': 3, 'var2': 3}
     assert params.target_variable == 'target'
     assert params.rem_static_win is True
     assert params.low_lim == {'var1': 0.0, 'var2': 1.0}
@@ -118,9 +126,9 @@ def test_train_model_params_from_dict_wrong_type(valid_train_params_dict):
     """Test from_dict raises TypeError when field has wrong type."""
     from model_manager.utils.models.train_model_params import TrainModelParams
 
-    valid_train_params_dict['lag_train'] = 'not_an_int'
+    valid_train_params_dict['lag_train'] = 'not_a_dict'
 
-    with pytest.raises(TypeError, match='lag_train must be of type int, but got str'):
+    with pytest.raises(TypeError, match='lag_train must be of type dict, but got str'):
         TrainModelParams.from_dict(valid_train_params_dict)
 
 
@@ -179,10 +187,10 @@ def test_validate_business_rules_negative_lag_train(valid_train_params_dict):
     """Test validate_business_rules raises error when lag_train is negative."""
     from model_manager.utils.models.train_model_params import TrainModelParams
 
-    valid_train_params_dict['lag_train'] = -1
+    valid_train_params_dict['lag_train'] = {'var1': -1, 'var2': 5}
     params = TrainModelParams.from_dict(valid_train_params_dict)
 
-    with pytest.raises(ValueError, match='lag_train must be positive, got -1'):
+    with pytest.raises(ValueError, match='lag_train for var1 must be non-negative, got -1'):
         params.validate_business_rules()
 
 
@@ -190,10 +198,10 @@ def test_validate_business_rules_negative_lag_val(valid_train_params_dict):
     """Test validate_business_rules raises error when lag_val is negative."""
     from model_manager.utils.models.train_model_params import TrainModelParams
 
-    valid_train_params_dict['lag_val'] = -2
+    valid_train_params_dict['lag_val'] = {'var1': 3, 'var2': -2}
     params = TrainModelParams.from_dict(valid_train_params_dict)
 
-    with pytest.raises(ValueError, match='lag_val must be positive, got -2'):
+    with pytest.raises(ValueError, match='lag_val for var2 must be non-negative, got -2'):
         params.validate_business_rules()
 
 
@@ -204,7 +212,7 @@ def test_validate_business_rules_negative_window(valid_train_params_dict):
     valid_train_params_dict['window'] = -5
     params = TrainModelParams.from_dict(valid_train_params_dict)
 
-    with pytest.raises(ValueError, match='window must be positive, got -5'):
+    with pytest.raises(ValueError, match='window must be non-negative, got -5'):
         params.validate_business_rules()
 
 
@@ -334,7 +342,7 @@ def test_validate_business_rules_zero_lag_train(valid_train_params_dict):
     """Test validate_business_rules accepts lag_train = 0."""
     from model_manager.utils.models.train_model_params import TrainModelParams
 
-    valid_train_params_dict['lag_train'] = 0
+    valid_train_params_dict['lag_train'] = {'var1': 0, 'var2': 0}
     params = TrainModelParams.from_dict(valid_train_params_dict)
 
     params.validate_business_rules()  # Should not raise
@@ -344,7 +352,7 @@ def test_validate_business_rules_zero_lag_val(valid_train_params_dict):
     """Test validate_business_rules accepts lag_val = 0."""
     from model_manager.utils.models.train_model_params import TrainModelParams
 
-    valid_train_params_dict['lag_val'] = 0
+    valid_train_params_dict['lag_val'] = {'var1': 0, 'var2': 0}
     params = TrainModelParams.from_dict(valid_train_params_dict)
 
     params.validate_business_rules()  # Should not raise

tests/utils/models/test_train_model_result.py

@@ -14,8 +14,8 @@ def sample_params():
     """Create sample TrainModelParams for testing."""
     return TrainModelParams(
         variable_columns=['var1', 'var2'],
-        lag_train=5,
+        lag_train={'var1': 5, 'var2': 5},
-        lag_val=3,
+        lag_val={'var1': 3, 'var2': 3},
         target_variable='target',
         rem_static_win=True,
         low_lim={'var1': 0.0, 'var2': 0.0},
@@ -32,6 +32,14 @@ def sample_params():
         experiment_run_id=123,
         experiment_name='test-experiment',
         removed_intervals=[],
+        model_name='Linear Regression',
+        degree=1,
+        interaction_only=False,
+        nan_treatment='drop',
+        start_date=None,
+        end_date=None,
+        scaler_name='Standard Scaler',
+        support_filters={},
     )
 
 

tests/utils/repository/test_training_repository.py

@@ -33,8 +33,8 @@ def sample_params():
         experiment_name='test_experiment',
         target_variable='target',
         variable_columns=['var1', 'var2', 'var3'],
-        lag_train=0,
+        lag_train={'var1': 0, 'var2': 0, 'var3': 0},
-        lag_val=0,
+        lag_val={'var1': 0, 'var2': 0, 'var3': 0},
         rem_static_win=False,
         low_lim={},
         upp_lim={},
@@ -48,6 +48,14 @@ def sample_params():
         line_separator=',',
         decimal_separator='.',
         removed_intervals=[],
+        model_name='Linear Regression',
+        degree=1,
+        interaction_only=False,
+        nan_treatment='drop',
+        start_date=None,
+        end_date=None,
+        scaler_name='None',
+        support_filters={},
     )
 
 
@@ -106,8 +114,8 @@ class TestExtractModelEquation:
             experiment_name='test',
             target_variable='y',
             variable_columns=['x'],
-            lag_train=0,
+            lag_train={'x': 0},
-            lag_val=0,
+            lag_val={'x': 0},
             rem_static_win=False,
             low_lim={},
             upp_lim={},
@@ -121,6 +129,14 @@ class TestExtractModelEquation:
             line_separator=',',
             decimal_separator='.',
             removed_intervals=[],
+            model_name='Linear Regression',
+            degree=1,
+            interaction_only=False,
+            nan_treatment='drop',
+            start_date=None,
+            end_date=None,
+            scaler_name='None',
+            support_filters={},
         )
 
         # Mock model with single coefficient
@@ -177,6 +193,7 @@ class TestInitDataPreprocessor:
     def test_init_preprocessor_with_scaler(self, training_repo, sample_params):
         """Test preprocessor initialization with scaler enabled."""
         sample_params.use_scaler = True
+        sample_params.scaler_name = 'Standard Scaler'
         preprocessor = training_repo._init_data_preprocessor(sample_params)
 
         assert preprocessor.scaler_name == 'Standard Scaler'
@@ -184,6 +201,7 @@ class TestInitDataPreprocessor:
     def test_init_preprocessor_without_scaler(self, training_repo, sample_params):
         """Test preprocessor initialization without scaler."""
         sample_params.use_scaler = False
+        sample_params.scaler_name = 'None'
         preprocessor = training_repo._init_data_preprocessor(sample_params)
 
         assert preprocessor.scaler_name == 'None'
@@ -211,14 +229,13 @@ class TestInitDataPreprocessor:
 
     def test_init_preprocessor_lag_configuration(self, training_repo, sample_params):
         """Test preprocessor lag configuration."""
-        sample_params.lag_train = 5
+        sample_params.lag_train = {'var1': 5, 'var2': 5, 'var3': 5}
-        sample_params.lag_val = 3
+        sample_params.lag_val = {'var1': 3, 'var2': 3, 'var3': 3}
         preprocessor = training_repo._init_data_preprocessor(sample_params)
 
-        # Check that lag dictionaries are created correctly
+        # Check that lag dictionaries are passed correctly
-        for col in sample_params.variable_columns:
+        assert preprocessor.lag_train == {'var1': 5, 'var2': 5, 'var3': 5}
-            assert preprocessor.lag_train[col] == 5
+        assert preprocessor.lag_transform == {'var1': 3, 'var2': 3, 'var3': 3}
-            assert preprocessor.lag_transform[col] == 3
 
 
 class TestInitScalerDict: