diff --git a/tests/sientia/test_models.py b/tests/sientia/test_models.py new file mode 100644 index 0000000..d84fddd --- /dev/null +++ b/tests/sientia/test_models.py @@ -0,0 +1,651 @@ +"""Unit tests for sientia models module.""" + +from unittest.mock import MagicMock, patch + +import numpy as np +import pandas as pd +from pytest import raises + +from model_manager.sientia.models import DataPreprocessor, LinearRegressionModel + +# LinearRegressionModel Tests + + +def test_linear_regression_model_init_default(): + """Test LinearRegressionModel initialization with default parameters.""" + model = LinearRegressionModel() + + assert model.target_variable == '' + assert model.variable_columns is None + assert model.model_params is None + assert model.clipping is None + assert model.weights is None + assert model.q1_target is None + assert model.q3_target is None + + +def test_linear_regression_model_init_with_params(): + """Test LinearRegressionModel initialization with parameters.""" + target = 'target' + variables = ['var1', 'var2'] + params = {'fit_intercept': True} + clipping = {'min': 0, 'max': 100} + weights = {'var1': 0.5, 'var2': 0.3} + + model = LinearRegressionModel( + target_variable=target, + variable_columns=variables, + model_params=params, + clipping=clipping, + weights=weights, + ) + + assert model.target_variable == target + assert model.variable_columns == variables + assert model.model_params == params + assert model.clipping == clipping + assert model.weights == weights + + +def test_linear_regression_model_fit(): + """Test LinearRegressionModel fit method.""" + model = LinearRegressionModel(target_variable='target', variable_columns=['var1', 'var2']) + + data = pd.DataFrame( + {'var1': [1, 2, 3, 4, 5], 'var2': [2, 3, 4, 5, 6], 'target': [3, 5, 7, 9, 11]} + ) + + result = model.fit(data) + + assert result is model + assert model.q1_target is not None + assert model.q3_target is not None + assert model.weights is not None + assert 'Bias' in model.weights + + +def test_linear_regression_model_fit_without_variable_columns(): + """Test LinearRegressionModel fit raises AssertionError without variable_columns.""" + 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.fit(data) + + +def test_linear_regression_model_predict_without_clipping(): + """Test LinearRegressionModel predict without clipping.""" + model = LinearRegressionModel(target_variable='target', variable_columns=['var1', 'var2']) + + train_data = pd.DataFrame( + {'var1': [1, 2, 3, 4, 5], 'var2': [2, 3, 4, 5, 6], 'target': [3, 5, 7, 9, 11]} + ) + model.fit(train_data) + + test_data = pd.DataFrame({'var1': [6, 7], 'var2': [7, 8]}) + predictions = model.predict(test_data) + + assert isinstance(predictions, np.ndarray) + assert len(predictions) == 2 + + +def test_linear_regression_model_predict_with_clipping_max(): + """Test LinearRegressionModel predict with clipping max.""" + model = LinearRegressionModel( + target_variable='target', variable_columns=['var1'], clipping={'min': 0, 'max': 5} + ) + + train_data = pd.DataFrame({'var1': [1, 2, 3, 4, 5], 'target': [1, 2, 3, 4, 5]}) + model.fit(train_data) + + test_data = pd.DataFrame({'var1': [10]}) + predictions = model.predict(test_data) + + assert predictions[0] == model.q3_target + + +def test_linear_regression_model_predict_with_clipping_min(): + """Test LinearRegressionModel predict with clipping min.""" + model = LinearRegressionModel( + target_variable='target', variable_columns=['var1'], clipping={'min': 0, 'max': 10} + ) + + train_data = pd.DataFrame({'var1': [1, 2, 3, 4, 5], 'target': [1, 2, 3, 4, 5]}) + model.fit(train_data) + + test_data = pd.DataFrame({'var1': [-10]}) + predictions = model.predict(test_data) + + assert predictions[0] == model.q1_target + + +def test_linear_regression_model_predict_with_clipping_within_range(): + """Test LinearRegressionModel predict with clipping but value within range.""" + model = LinearRegressionModel( + target_variable='target', variable_columns=['var1'], clipping={'min': 0, 'max': 10} + ) + + train_data = pd.DataFrame({'var1': [1, 2, 3, 4, 5], 'target': [1, 2, 3, 4, 5]}) + model.fit(train_data) + + test_data = pd.DataFrame({'var1': [3]}) + predictions = model.predict(test_data) + + # Prediction should be within range and not clipped + assert 0 <= predictions[0] <= 10 + + +# DataPreprocessor Tests + + +def test_data_preprocessor_init_default(): + """Test DataPreprocessor initialization with default parameters.""" + preprocessor = DataPreprocessor() + + assert preprocessor.date_column == '' + assert preprocessor.target_variable == '' + assert preprocessor.input_columns is None + assert preprocessor.nan_treatment is None + assert preprocessor.lag_train == {} + assert preprocessor.lag_transform == {} + assert preprocessor.scaler is None + + +def test_data_preprocessor_init_with_standard_scaler(): + """Test DataPreprocessor initialization with Standard Scaler.""" + preprocessor = DataPreprocessor(scaler_name='Standard Scaler') + + assert preprocessor.scaler is not None + + +def test_data_preprocessor_init_with_none_scaler(): + """Test DataPreprocessor initialization with None scaler.""" + preprocessor = DataPreprocessor(scaler_name='None') + + assert preprocessor.scaler is None + + +def test_data_preprocessor_init_with_unknown_scaler(): + """Test DataPreprocessor initialization with unknown scaler.""" + preprocessor = DataPreprocessor(scaler_name='Unknown') + + assert preprocessor.scaler is None + + +def test_data_preprocessor_init_with_custom_steps_order(): + """Test DataPreprocessor initialization with custom steps order.""" + custom_steps = ['Normalization', 'Feature Creation'] + preprocessor = DataPreprocessor(steps_order=custom_steps) + + assert 'Normalization' in preprocessor.steps_order + assert 'Feature Creation' in preprocessor.steps_order + assert len(preprocessor.steps_order) == 7 + + +def test_data_preprocessor_get_scaler(): + """Test DataPreprocessor get_scaler method.""" + preprocessor = DataPreprocessor(scaler_name='Standard Scaler') + + scaler = preprocessor.get_scaler() + + assert scaler is not None + + +@patch('model_manager.sientia.models.treat_nan') +def test_data_preprocessor_treat_discontinuities_with_treatment(mock_treat_nan): + """Test treat_discontinuities with nan_treatment.""" + preprocessor = DataPreprocessor(nan_treatment='drop') + data = pd.DataFrame({'col1': [1, 2, np.nan]}) + expected_data = pd.DataFrame({'col1': [1, 2]}) + mock_treat_nan.return_value = expected_data + + result = preprocessor.treat_discontinuities(data) + + mock_treat_nan.assert_called_once_with(data, 'drop') + pd.testing.assert_frame_equal(result, expected_data) + + +def test_data_preprocessor_treat_discontinuities_without_treatment(): + """Test treat_discontinuities without nan_treatment.""" + preprocessor = DataPreprocessor() + data = pd.DataFrame({'col1': [1, 2, 3]}) + + result = preprocessor.treat_discontinuities(data) + + pd.testing.assert_frame_equal(result, data) + + +def test_data_preprocessor_lag_selection_with_lag(): + """Test lag_selection with lag.""" + preprocessor = DataPreprocessor() + data = pd.DataFrame({'var1': [1, 2, 3, 4, 5]}) + lag_dict = {'var1': 1} + + result = preprocessor.lag_selection(data, lag_dict) + + assert len(result) == 4 + assert result['var1'].iloc[0] == 1 + + +def test_data_preprocessor_lag_selection_without_lag(): + """Test lag_selection without lag.""" + preprocessor = DataPreprocessor() + data = pd.DataFrame({'var1': [1, 2, 3, 4, 5]}) + lag_dict = {} + + result = preprocessor.lag_selection(data, lag_dict) + + assert len(result) == 5 + + +def test_data_preprocessor_lag_selection_with_zero_lag(): + """Test lag_selection with zero lag.""" + preprocessor = DataPreprocessor() + data = pd.DataFrame({'var1': [1, 2, 3, 4, 5]}) + lag_dict = {'var1': 0} + + result = preprocessor.lag_selection(data, lag_dict) + + assert len(result) == 5 + + +@patch('model_manager.sientia.models.TimeSeriesDiscontinuityAnalyzer') +def test_data_preprocessor_treat_static_windows(mock_analyzer_class): + """Test treat_static_windows method.""" + preprocessor = DataPreprocessor(static_threshold=3) + data = pd.DataFrame({'col1': [1, 1, 1, 2, 3]}) + + mock_analyzer = MagicMock() + mock_analyzer_class.return_value = mock_analyzer + mock_analyzer.get_treated_data.return_value = data + + preprocessor.treat_static_windows(data) + + mock_analyzer.infer_frequency.assert_called_once() + assert mock_analyzer.identify_static_windows.called + assert mock_analyzer.treat_static_windows.called + + +def test_data_preprocessor_treat_static_windows_without_threshold(): + """Test treat_static_windows without threshold.""" + preprocessor = DataPreprocessor() + data = pd.DataFrame({'col1': [1, 2, 3]}) + + result = preprocessor.treat_static_windows(data) + + pd.testing.assert_frame_equal(result, data) + + +@patch('model_manager.sientia.models.limit_dataset') +def test_data_preprocessor_adjust_limits(mock_limit_dataset): + """Test adjust_limits method.""" + preprocessor = DataPreprocessor(low_lim={'col1': 0}, upp_lim={'col1': 10}) + data = pd.DataFrame({'col1': [1, 2, 3]}) + expected_data = pd.DataFrame({'col1': [1, 2, 3]}) + mock_limit_dataset.return_value = (expected_data, {'col1': 0}, {'col1': 10}) + + result = preprocessor.adjust_limits(data) + + mock_limit_dataset.assert_called_once() + pd.testing.assert_frame_equal(result, expected_data) + + +@patch('model_manager.sientia.models.create_features') +def test_data_preprocessor_create_features(mock_create_features): + """Test create_features method.""" + preprocessor = DataPreprocessor( + self_operations=['{var1}_{pow}_{2}'], cross_operations=['{var1}_{*}_{var2}'] + ) + data = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4]}) + expected_data = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var1_pow_2': [1, 4, 9]}) + mock_create_features.return_value = expected_data + + result = preprocessor.create_features(data) + + mock_create_features.assert_called_once() + pd.testing.assert_frame_equal(result, expected_data) + + +def test_data_preprocessor_create_ar(): + """Test create_ar method.""" + preprocessor = DataPreprocessor(target_variable='target', ar_var='ar_target') + data = pd.DataFrame({'target': [1, 2, 3, 4, 5]}) + + result = preprocessor.create_ar(data) + + assert 'ar_target' in result.columns + assert len(result) == 4 + + +def test_data_preprocessor_create_ar_without_ar_var(): + """Test create_ar without ar_var.""" + preprocessor = DataPreprocessor(target_variable='target') + data = pd.DataFrame({'target': [1, 2, 3, 4, 5]}) + + result = preprocessor.create_ar(data) + + assert len(result) == 5 + + +def test_data_preprocessor_create_lags(): + """Test create_lags method.""" + preprocessor = DataPreprocessor(created_lags={'var1': 1}) + data = pd.DataFrame({'var1': [1, 2, 3, 4, 5]}) + + result = preprocessor.create_lags(data) + + assert 'var1_lag1' in result.columns + assert len(result) == 4 + + +def test_data_preprocessor_create_lags_with_zero_lag(): + """Test create_lags with zero lag.""" + preprocessor = DataPreprocessor(created_lags={'var1': 0}) + data = pd.DataFrame({'var1': [1, 2, 3, 4, 5]}) + + result = preprocessor.create_lags(data) + + assert 'var1_lag0' not in result.columns + assert len(result) == 5 + + +def test_data_preprocessor_create_lags_without_created_lags(): + """Test create_lags without created_lags.""" + preprocessor = DataPreprocessor() + data = pd.DataFrame({'var1': [1, 2, 3, 4, 5]}) + + result = preprocessor.create_lags(data) + + assert len(result) == 5 + + +def test_data_preprocessor_create_lags_with_missing_column(): + """Test create_lags with missing column.""" + preprocessor = DataPreprocessor(created_lags={'var2': 1}) + data = pd.DataFrame({'var1': [1, 2, 3, 4, 5]}) + + result = preprocessor.create_lags(data) + + assert 'var2_lag1' not in result.columns + + +def test_data_preprocessor_fit_with_x_and_y(): + """Test fit method with x and y.""" + preprocessor = DataPreprocessor( + target_variable='target', + input_columns=['var1', 'var2'], + steps_order=['Discontinuity Treatment'], + ) + x = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4]}) + y = pd.Series([3, 5, 7], name='target') + + result = preprocessor.fit(x, y) + + assert result is preprocessor + + +def test_data_preprocessor_fit_with_only_x(): + """Test fit method with only x.""" + preprocessor = DataPreprocessor( + target_variable='target', + input_columns=['var1', 'var2'], + steps_order=['Discontinuity Treatment'], + ) + x = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'target': [3, 5, 7]}) + + result = preprocessor.fit(x) + + assert result is preprocessor + + +def test_data_preprocessor_fit_without_data(): + """Test fit method without data.""" + preprocessor = DataPreprocessor(target_variable='target', input_columns=['var1']) + + with raises(ValueError, match='No data was provided'): + preprocessor.fit(None, None) + + +def test_data_preprocessor_fit_without_input_columns(): + """Test fit method without input_columns.""" + preprocessor = DataPreprocessor(target_variable='target') + x = pd.DataFrame({'var1': [1, 2, 3], 'target': [3, 5, 7]}) + + with raises(AssertionError, match='input_columns must be set'): + preprocessor.fit(x) + + +def test_data_preprocessor_fit_with_normalization(): + """Test fit method with normalization.""" + preprocessor = DataPreprocessor( + target_variable='target', + input_columns=['var1', 'var2'], + scaler_name='Standard Scaler', + scaler_params={}, + steps_order=['Normalization'], + ) + x = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'target': [3, 5, 7]}) + + result = preprocessor.fit(x) + + assert result is preprocessor + assert preprocessor.scaler_params is not None + + +def test_data_preprocessor_transform_with_timestamp(): + """Test transform method with timestamp column.""" + preprocessor = DataPreprocessor( + target_variable='target', input_columns=['var1'], steps_order=['Discontinuity Treatment'] + ) + x = pd.DataFrame({'timestamp': [1, 2, 3], 'var1': [1, 2, 3], 'target': [3, 5, 7]}) + + result = preprocessor.transform(x) + + assert 'timestamp' not in result.columns + + +def test_data_preprocessor_transform_without_timestamp(): + """Test transform method without timestamp column.""" + preprocessor = DataPreprocessor( + target_variable='target', input_columns=['var1'], steps_order=['Discontinuity Treatment'] + ) + x = pd.DataFrame({'var1': [1, 2, 3], 'target': [3, 5, 7]}) + + result = preprocessor.transform(x) + + assert 'var1' in result.columns + + +def test_data_preprocessor_transform_without_input_columns(): + """Test transform method without input_columns.""" + preprocessor = DataPreprocessor(target_variable='target') + x = pd.DataFrame({'var1': [1, 2, 3], 'target': [3, 5, 7]}) + + with raises(AssertionError, match='input_columns must be set'): + preprocessor.transform(x) + + +def test_data_preprocessor_transform_with_feature_creation(): + """Test transform method with feature creation.""" + preprocessor = DataPreprocessor( + target_variable='target', + input_columns=['var1'], + self_operations=['{var1}_{pow}_{2}'], + steps_order=['Feature Creation'], + ) + x = pd.DataFrame({'var1': [1, 2, 3], 'target': [3, 5, 7]}) + + with patch('model_manager.sientia.models.create_features') as mock_create: + mock_create.return_value = x + preprocessor.transform(x) + mock_create.assert_called_once() + + +def test_data_preprocessor_transform_with_lag_creation(): + """Test transform method with lag creation.""" + preprocessor = DataPreprocessor( + target_variable='target', + input_columns=['var1', 'ar_target'], + ar_var='ar_target', + created_lags={'var1': 1}, + steps_order=['Lag Creation'], + ) + x = pd.DataFrame({'var1': [1, 2, 3, 4, 5], 'target': [3, 5, 7, 9, 11]}) + + result = preprocessor.transform(x) + + assert 'ar_target' in result.columns + assert 'var1_lag1' in result.columns + + +def test_data_preprocessor_transform_with_normalization(): + """Test transform method with normalization.""" + preprocessor = DataPreprocessor( + target_variable='target', + input_columns=['var1'], + scaler_name='Standard Scaler', + scaler_params={}, + steps_order=['Normalization'], + ) + train_x = pd.DataFrame({'var1': [1, 2, 3], 'target': [3, 5, 7]}) + preprocessor.fit(train_x) + + test_x = pd.DataFrame({'var1': [4, 5, 6], 'target': [9, 11, 13]}) + result = preprocessor.transform(test_x) + + assert 'var1' in result.columns + + +def test_data_preprocessor_get_required_columns_with_self_operations(): + """Test get_required_columns with self_operations.""" + preprocessor = DataPreprocessor(self_operations=['{var1}_{pow}_{2}']) + existing_columns = ['var1', 'var2'] + + result = preprocessor.get_required_columns(existing_columns) + + assert 'var1' in result + + +def test_data_preprocessor_get_required_columns_with_cross_operations(): + """Test get_required_columns with cross_operations.""" + preprocessor = DataPreprocessor(cross_operations=['{var1}_{*}_{var2}']) + existing_columns = ['var1', 'var2'] + + result = preprocessor.get_required_columns(existing_columns) + + assert 'var1' in result + assert 'var2' in result + + +def test_data_preprocessor_get_required_columns_with_created_lags(): + """Test get_required_columns with created_lags.""" + preprocessor = DataPreprocessor(created_lags={'var1': 1}) + existing_columns = ['var1', 'var2'] + + result = preprocessor.get_required_columns(existing_columns) + + assert 'var1' in result + + +def test_data_preprocessor_get_required_columns_with_missing_columns(): + """Test get_required_columns with missing columns in existing_columns.""" + preprocessor = DataPreprocessor( + self_operations=['{var3}_{pow}_{2}'], cross_operations=['{var4}_{*}_{var5}'] + ) + existing_columns = ['var1', 'var2'] + + result = preprocessor.get_required_columns(existing_columns) + + assert 'var3' in result + assert 'var4' in result + assert 'var5' in result + + +def test_data_preprocessor_get_required_columns_removes_duplicates(): + """Test get_required_columns removes duplicates from self_operations.""" + preprocessor = DataPreprocessor(self_operations=['{var1}_{pow}_{2}'], created_lags={'var1': 1}) + existing_columns = ['var1', 'var2'] + + result = preprocessor.get_required_columns(existing_columns) + + # var1 is in existing_columns, so it should not be in required_columns + assert 'var1' not in result or result.count('var1') <= 1 + + +def test_data_preprocessor_get_required_columns_removes_created_lags(): + """Test get_required_columns removes columns from created_lags when column is in created_lags dict - covers line 292.""" + preprocessor = DataPreprocessor(created_lags={'var1': 1, 'var2': 1}) + existing_columns = ['var3'] + + result = preprocessor.get_required_columns(existing_columns) + + # var1 and var2 should be removed because they're in created_lags dict and not in existing_columns + assert 'var1' not in result + assert 'var2' not in result + + +def test_data_preprocessor_fit_all_steps(): + """Test fit method with all steps.""" + preprocessor = DataPreprocessor( + target_variable='target', + input_columns=['var1'], + nan_treatment='drop', + lag_train={'var1': 1}, + static_threshold=3, + low_lim={'var1': 0}, + upp_lim={'var1': 10}, + scaler_name='Standard Scaler', + scaler_params={}, + ) + x = pd.DataFrame({'var1': [1, 2, 3, 4, 5], 'target': [3, 5, 7, 9, 11]}) + + with patch('model_manager.sientia.models.treat_nan') as mock_treat: + with patch('model_manager.sientia.models.TimeSeriesDiscontinuityAnalyzer'): + with patch('model_manager.sientia.models.limit_dataset') as mock_limit: + mock_treat.return_value = x + mock_limit.return_value = (x, {'var1': 0}, {'var1': 10}) + result = preprocessor.fit(x) + + assert result is preprocessor + + +def test_data_preprocessor_transform_all_steps(): + """Test transform method with all steps.""" + preprocessor = DataPreprocessor( + target_variable='target', + input_columns=['var1'], + nan_treatment='drop', + lag_transform={'var1': 1}, + static_threshold=3, + low_lim={'var1': 0}, + upp_lim={'var1': 10}, + scaler_name='Standard Scaler', + scaler_params={}, + self_operations=['{var1}_{pow}_{2}'], + ar_var='ar_target', + created_lags={'var1': 1}, + ) + train_x = pd.DataFrame({'var1': [1, 2, 3, 4, 5], 'target': [3, 5, 7, 9, 11]}) + + with patch('model_manager.sientia.models.treat_nan') as mock_treat: + with patch('model_manager.sientia.models.TimeSeriesDiscontinuityAnalyzer'): + with patch('model_manager.sientia.models.limit_dataset') as mock_limit: + with patch('model_manager.sientia.models.create_features') as mock_create: + mock_treat.return_value = train_x + mock_limit.return_value = (train_x, {'var1': 0}, {'var1': 10}) + mock_create.return_value = train_x + preprocessor.fit(train_x) + + test_x = pd.DataFrame({'var1': [6, 7, 8, 9, 10], 'target': [13, 15, 17, 19, 21]}) + + with patch('model_manager.sientia.models.treat_nan') as mock_treat: + with patch('model_manager.sientia.models.TimeSeriesDiscontinuityAnalyzer'): + with patch('model_manager.sientia.models.limit_dataset') as mock_limit: + with patch('model_manager.sientia.models.create_features') as mock_create: + mock_treat.return_value = test_x + mock_limit.return_value = (test_x, {'var1': 0}, {'var1': 10}) + mock_create.return_value = test_x + result = preprocessor.transform(test_x) + + assert isinstance(result, pd.DataFrame)