upgini 1.1.266a3254.post2__tar.gz → 1.1.267a3254.post3__tar.gz

{upgini-1.1.266a3254.post2/src/upgini.egg-info → upgini-1.1.267a3254.post3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: upgini
-Version: 1.1.266a3254.post2
+Version: 1.1.267a3254.post3
 Summary: Intelligent data search & enrichment for Machine Learning
 Home-page: https://upgini.com/
 Author: Upgini Developers

{upgini-1.1.266a3254.post2 → upgini-1.1.267a3254.post3}/setup.py

@@ -40,7 +40,7 @@ def send_log(msg: str):
 
 
 here = Path(__file__).parent.resolve()
-version = "1.1.266a3254-2"
+version = "1.1.267a3254-3"
 try:
     send_log(f"Start setup PyLib version {version}")
     setup(

{upgini-1.1.266a3254.post2 → upgini-1.1.267a3254.post3}/src/upgini/autofe/date.py

@@ -54,6 +54,9 @@ class DateDiffType2(PandasOperand, DateDiffMixin):
         return diff
 
 
+_ext_aggregations = {"nunique": (lambda x: len(np.unique(x)), 0), "count": (len, 0)}
+
+
 class DateListDiff(PandasOperand, DateDiffMixin):
     is_binary = True
     has_symmetry_importance = True
@@ -72,18 +75,31 @@ class DateListDiff(PandasOperand, DateDiffMixin):
 
     def calculate_binary(self, left: pd.Series, right: pd.Series) -> pd.Series:
         left = self._convert_to_date(left, self.left_unit)
+        right = right.apply(lambda x: pd.arrays.DatetimeArray(self._convert_to_date(x, self.right_unit)))
+
+        return pd.Series(left - right.values).apply(lambda x: self._agg(self._diff(x)))
 
-        return pd.Series(left.index.map(lambda i: self.reduce(self.map_diff(left.loc[i], right.loc[i]))))
+    def _diff(self, x):
+        x = x / np.timedelta64(1, self.diff_unit)
+        return x[x > 0]
+
+    def _agg(self, x):
+        method = getattr(np, self.aggregation, None)
+        default = np.nan
+        if method is None and self.aggregation in _ext_aggregations:
+            method, default = _ext_aggregations[self.aggregation]
+        elif not callable(method):
+            raise ValueError(f"Unsupported aggregation: {self.aggregation}")
+
+        return method(x) if len(x) > 0 else default
 
 
 class DateListDiffBounded(DateListDiff):
     lower_bound: Optional[int]
     upper_bound: Optional[int]
-    inclusive: Optional[str]
 
     def __init__(self, **data: Any) -> None:
         if "name" not in data:
-            inclusive = data.get("inclusive")
             lower_bound = data.get("lower_bound")
             upper_bound = data.get("upper_bound")
             components = [
@@ -92,18 +108,10 @@ class DateListDiffBounded(DateListDiff):
                 str(lower_bound if lower_bound is not None else "minusinf"),
                 str(upper_bound if upper_bound is not None else "plusinf"),
             ]
-            if inclusive:
-                components.append(inclusive)
             components.append(data.get("aggregation"))
             data["name"] = "_".join(components)
         super().__init__(**data)
 
-    def reduce(self, diff_list: pd.Series) -> float:
-        return diff_list[
-            (diff_list > 0)
-            & (
-                diff_list.between(
-                    self.lower_bound or -np.inf, self.upper_bound or np.inf, inclusive=self.inclusive or "left"
-                )
-            )
-        ].aggregate(self.aggregation)
+    def _agg(self, x):
+        x = x[(x >= (self.lower_bound or -np.inf)) & (x < (self.upper_bound or np.inf))]
+        return super()._agg(x)

{upgini-1.1.266a3254.post2 → upgini-1.1.267a3254.post3}/src/upgini/features_enricher.py

@@ -94,7 +94,7 @@ try:
 except Exception:
     from upgini.utils.fallback_progress_bar import CustomFallbackProgressBar as ProgressBar
 
-from upgini.utils.target_utils import define_task
+from upgini.utils.target_utils import calculate_psi, define_task
 from upgini.utils.warning_counter import WarningCounter
 from upgini.version_validator import validate_version
 
@@ -2226,14 +2226,11 @@ class FeaturesEnricher(TransformerMixin):
             validated_X, self.fit_search_keys, self.logger, self.bundle, self.warning_counter
         )
 
-        has_date = self._get_date_column(self.fit_search_keys) is not None
+        maybe_date_column = self._get_date_column(self.fit_search_keys)
+        has_date = maybe_date_column is not None
         model_task_type = self.model_task_type or define_task(validated_y, has_date, self.logger)
         self._validate_binary_observations(validated_y, model_task_type)
 
-        df = self.__handle_index_search_keys(df, self.fit_search_keys)
-
-        df = self.__correct_target(df)
-
         self.runtime_parameters = get_runtime_params_custom_loss(
             self.loss, model_task_type, self.runtime_parameters, self.logger
         )
@@ -2245,6 +2242,13 @@ class FeaturesEnricher(TransformerMixin):
                 eval_df[EVAL_SET_INDEX] = idx + 1
                 df = pd.concat([df, eval_df])
 
+        df = self.__correct_target(df)
+
+        df = self.__handle_index_search_keys(df, self.fit_search_keys)
+
+        if is_numeric_dtype(df[self.TARGET_NAME]) and has_date:
+            self._validate_PSI(df.sort_values(by=maybe_date_column))
+
         if DEFAULT_INDEX in df.columns:
             msg = self.bundle.get("unsupported_index_column")
             self.logger.info(msg)
@@ -3567,6 +3571,34 @@ class FeaturesEnricher(TransformerMixin):
             self.logger.warning(msg)
             print(msg)
 
+    def _validate_PSI(self, df: pd.DataFrame):
+        if EVAL_SET_INDEX in df.columns:
+            train = df.query(f"{EVAL_SET_INDEX} == 0")
+            eval1 = df.query(f"{EVAL_SET_INDEX} == 1")
+        else:
+            train = df
+            eval1 = None
+
+        # 1. Check train PSI
+        half_train = round(len(train) / 2)
+        part1 = train[:half_train]
+        part2 = train[half_train:]
+        train_psi = calculate_psi(part1[self.TARGET_NAME], part2[self.TARGET_NAME])
+        if train_psi > 0.2:
+            self.warning_counter.increment()
+            msg = self.bundle.get("train_unstable_target").format(train_psi)
+            print(msg)
+            self.logger.warning(msg)
+
+        # 2. Check train-test PSI
+        if eval1 is not None:
+            train_test_psi = calculate_psi(train[self.TARGET_NAME], eval1[self.TARGET_NAME])
+            if train_test_psi > 0.2:
+                self.warning_counter.increment()
+                msg = self.bundle.get("eval_unstable_target").format(train_test_psi)
+                print(msg)
+                self.logger.warning(msg)
+
     def _dump_python_libs(self):
         try:
             from pip._internal.operations.freeze import freeze

{upgini-1.1.266a3254.post2 → upgini-1.1.267a3254.post3}/src/upgini/resource_bundle/strings.properties

@@ -111,7 +111,9 @@ x_is_empty=X is empty
 y_is_empty=y is empty
 x_contains_reserved_column_name=Column name {} is reserved. Please rename column and try again
 missing_generate_feature=\nWARNING: Feature {} specified in `generate_features` is not present in input columns: {}
-x_unstable_by_date=\nWARNING: Your training sample is unstable in number of rows per date. It is recommended to redesign the training sample.
+x_unstable_by_date=\nWARNING: Your training sample is unstable in number of rows per date. It is recommended to redesign the training sample
+train_unstable_target=\nWARNING: Your training sample contains an unstable target event, PSI = {}. This will lead to unstable scoring on deferred samples. It is recommended to redesign the training sample
+eval_unstable_target=\nWARNING: Your training and evaluation samples have a difference in target distribution. PSI = {}. The results will be unstable. It is recommended to redesign the training and evaluation samples
     # eval set validation
 unsupported_type_eval_set=Unsupported type of eval_set: {}. It should be list of tuples with two elements: X and y
 eval_set_invalid_tuple_size=eval_set contains a tuple of size {}. It should contain only pairs of X and y
@@ -198,7 +200,7 @@ email_detected=Emails detected in column `{}`. It will be used as a search key\n
 email_detected_not_registered=Emails detected in column `{}`. It can be used only with api_key from profile.upgini.com\nSee docs to turn off the automatic detection: https://github.com/upgini/upgini/blob/main/README.md#turn-off-autodetection-for-search-key-columns
 phone_detected=Phone numbers detected in column `{}`. It can be used only with api_key from profile.upgini.com\nSee docs to turn off the automatic detection: https://github.com/upgini/upgini/blob/main/README.md#turn-off-autodetection-for-search-key-columns
 phone_detected_not_registered=\nWARNING: Phone numbers detected in column `{}`. It can be used only with api_key from profile.upgini.com\nSee docs to turn off the automatic detection: https://github.com/upgini/upgini/blob/main/README.md#turn-off-autodetection-for-search-key-columns
-target_type_detected=Detected task type: {}\n
+target_type_detected=\nDetected task type: {}\n
 # all_ok_community_invite=Chat with us in Slack community:
 all_ok_community_invite=❓ Support request 
 too_small_for_metrics=Your train dataset contains less than 500 rows. For such dataset Upgini will not calculate accuracy metrics. Please increase the number of rows in the training dataset to calculate accuracy metrics

{upgini-1.1.266a3254.post2 → upgini-1.1.267a3254.post3}/src/upgini/utils/target_utils.py

@@ -177,3 +177,21 @@ def balance_undersample(
 
     logger.info(f"Shape after rebalance resampling: {resampled_data}")
     return resampled_data
+
+
+def calculate_psi(expected: pd.Series, actual: pd.Series) -> float:
+    df = pd.concat([expected, actual])
+
+    # Define the bins for the target variable
+    df_min = df.min()
+    df_max = df.max()
+    bins = [df_min, (df_min + df_max) / 2, df_max]
+
+    # Calculate the base distribution
+    train_distribution = expected.value_counts(bins=bins, normalize=True).sort_index().values
+
+    # Calculate the target distribution
+    test_distribution = actual.value_counts(bins=bins, normalize=True).sort_index().values
+
+    # Calculate the PSI
+    return np.sum((train_distribution - test_distribution) * np.log(train_distribution / test_distribution))

{upgini-1.1.266a3254.post2 → upgini-1.1.267a3254.post3/src/upgini.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: upgini
-Version: 1.1.266a3254.post2
+Version: 1.1.267a3254.post3
 Summary: Intelligent data search & enrichment for Machine Learning
 Home-page: https://upgini.com/
 Author: Upgini Developers

{upgini-1.1.266a3254.post2 → upgini-1.1.267a3254.post3}/tests/test_autofe_operands.py

@@ -1,4 +1,5 @@
 import pandas as pd
+from pydantic import NoneBytes
 from upgini.autofe.date import DateDiff, DateDiffType2, DateListDiff, DateListDiffBounded
 
 from datetime import datetime
@@ -47,7 +48,7 @@ def test_date_diff_list():
     def check(aggregation, expected_name, expected_values):
         operand = DateListDiff(aggregation=aggregation)
         assert operand.name == expected_name
-        assert_series_equal(operand.calculate_binary(df.date1, df.date2), expected_values)
+        assert_series_equal(operand.calculate_binary(df.date1, df.date2).rename(None), expected_values)
 
     check(aggregation="min", expected_name="date_diff_min", expected_values=pd.Series([10530, 10531, None, None]))
     check(aggregation="max", expected_name="date_diff_max", expected_values=pd.Series([10531, 10531, None, None]))
@@ -83,7 +84,7 @@ def test_date_diff_list_bounded():
             diff_unit="Y", aggregation="count", lower_bound=lower_bound, upper_bound=upper_bound
         )
         assert operand.name == expected_name
-        assert_series_equal(operand.calculate_binary(df.date1, df.date2), expected_values)
+        assert_series_equal(operand.calculate_binary(df.date1, df.date2).rename(None), expected_values)
 
     check_num_by_years(0, 18, "date_diff_Y_0_18_count", pd.Series([2, 1, 0, 0, 0]))
     check_num_by_years(18, 23, "date_diff_Y_18_23_count", pd.Series([1, 2, 2, 0, 0]))

{upgini-1.1.266a3254.post2 → upgini-1.1.267a3254.post3}/tests/test_features_enricher.py

@@ -428,13 +428,13 @@ def test_saved_features_enricher(requests_mock: Mocker):
     df.drop(columns="SystemRecordId_473310000", inplace=True)
     train_df = df.head(10000)
     train_features = train_df.drop(columns="target")
-    train_target = train_df["target"]
+    train_target = train_df["target"].copy()
     eval1_df = df[10000:11000].reset_index(drop=True)
     eval1_features = eval1_df.drop(columns="target")
-    eval1_target = eval1_df["target"].reset_index(drop=True)
+    eval1_target = eval1_df["target"].reset_index(drop=True).copy()
     eval2_df = df[11000:12000]
     eval2_features = eval2_df.drop(columns="target")
-    eval2_target = eval2_df["target"]
+    eval2_target = eval2_df["target"].copy()
 
     enricher = FeaturesEnricher(
         search_keys={"phone_num": SearchKey.PHONE, "rep_date": SearchKey.DATE},
@@ -482,6 +482,31 @@ def test_saved_features_enricher(requests_mock: Mocker):
     assert first_feature_info[feature_name_header] == "feature"
     assert first_feature_info[shap_value_header] == 10.1
 
+    # Check imbalanced target metrics
+    random = np.random.RandomState(42)
+    train_random_indices = random.choice(train_target.index, size=9000, replace=False)
+    train_target.loc[train_random_indices] = 0
+
+    metrics = enricher.calculate_metrics(
+        train_features,
+        train_target
+    )
+    expected_metrics = pd.DataFrame(
+        {
+            segment_header: [train_segment],
+            rows_header: [10000],
+            target_mean_header: [0.049],
+            enriched_gini: [0.000985],
+        }
+    )
+    print("Expected metrics: ")
+    print(expected_metrics)
+    print("Actual metrics: ")
+    print(metrics)
+
+    assert metrics is not None
+    assert_frame_equal(expected_metrics, metrics, atol=1e-6)
+
 
 def test_features_enricher_with_demo_key(requests_mock: Mocker):
     url = "http://fake_url2"
@@ -2498,11 +2523,11 @@ def test_diff_target_dups(requests_mock: Mocker):
         assert len(self.data) == 2
         print(self.data)
         assert self.data.loc[0, "date_0e8763"] == 1672531200000
-        assert self.data.loc[0, "feature_2ad562"] == 12
-        assert self.data.loc[0, "target"] == 0
+        assert self.data.loc[0, "feature_2ad562"] == 13
+        assert self.data.loc[0, "target"] == 1
         assert self.data.loc[1, "date_0e8763"] == 1672531200000
-        assert self.data.loc[1, "feature_2ad562"] == 13
-        assert self.data.loc[1, "target"] == 1
+        assert self.data.loc[1, "feature_2ad562"] == 12
+        assert self.data.loc[1, "target"] == 0
         return SearchTask("123", self, rest_client=enricher.rest_client)
 
     Dataset.search = mock_search

{upgini-1.1.266a3254.post2 → upgini-1.1.267a3254.post3}/tests/test_metrics.py

@@ -857,23 +857,23 @@ def test_catboost_metric_binary(requests_mock: Mocker):
     assert metrics_df.loc[0, segment_header] == train_segment
     assert metrics_df.loc[0, rows_header] == 500
     assert metrics_df.loc[0, target_mean_header] == 0.51
-    assert metrics_df.loc[0, baseline_gini] == approx(0.023101)
-    assert metrics_df.loc[0, enriched_gini] == approx(0.090344)
-    assert metrics_df.loc[0, uplift] == approx(0.067243)
+    assert metrics_df.loc[0, baseline_gini] == approx(0.061408)
+    assert metrics_df.loc[0, enriched_gini] == approx(0.071498)
+    assert metrics_df.loc[0, uplift] == approx(0.010090)
 
     assert metrics_df.loc[1, segment_header] == eval_1_segment
     assert metrics_df.loc[1, rows_header] == 250
     assert metrics_df.loc[1, target_mean_header] == 0.452
-    assert metrics_df.loc[1, baseline_gini] == approx(-0.016188)
-    assert metrics_df.loc[1, enriched_gini] == approx(0.014947)
-    assert metrics_df.loc[1, uplift] == approx(0.031135)
+    assert metrics_df.loc[1, baseline_gini] == approx(-0.051702)
+    assert metrics_df.loc[1, enriched_gini] == approx(0.023668)
+    assert metrics_df.loc[1, uplift] == approx(0.075370)
 
     assert metrics_df.loc[2, segment_header] == eval_2_segment
     assert metrics_df.loc[2, rows_header] == 250
     assert metrics_df.loc[2, target_mean_header] == 0.536
-    assert metrics_df.loc[2, baseline_gini] == approx(-0.017138)
-    assert metrics_df.loc[2, enriched_gini] == approx(0.035666)
-    assert metrics_df.loc[2, uplift] == approx(0.052805)
+    assert metrics_df.loc[2, baseline_gini] == approx(0.012674)
+    assert metrics_df.loc[2, enriched_gini] == approx(0.022980)
+    assert metrics_df.loc[2, uplift] == approx(0.010306)
 
 
 def test_catboost_metric_binary_with_cat_features(requests_mock: Mocker):
@@ -984,23 +984,23 @@ def test_catboost_metric_binary_with_cat_features(requests_mock: Mocker):
     assert metrics_df.loc[0, segment_header] == train_segment
     assert metrics_df.loc[0, rows_header] == 500
     assert metrics_df.loc[0, target_mean_header] == 0.51
-    assert metrics_df.loc[0, baseline_gini] == approx(0.102928)
-    assert metrics_df.loc[0, enriched_gini] == approx(0.139437)
-    assert metrics_df.loc[0, uplift] == approx(0.036508)
+    assert metrics_df.loc[0, baseline_gini] == approx(0.027066)
+    assert metrics_df.loc[0, enriched_gini] == approx(0.101601)
+    assert metrics_df.loc[0, uplift] == approx(0.074535)
 
     assert metrics_df.loc[1, segment_header] == eval_1_segment
     assert metrics_df.loc[1, rows_header] == 250
     assert metrics_df.loc[1, target_mean_header] == 0.452
-    assert metrics_df.loc[1, baseline_gini] == approx(-0.074491)
-    assert metrics_df.loc[1, enriched_gini] == approx(-0.052619)
-    assert metrics_df.loc[1, uplift] == approx(0.021872)
+    assert metrics_df.loc[1, baseline_gini] == approx(-0.078548)
+    assert metrics_df.loc[1, enriched_gini] == approx(-0.019663)
+    assert metrics_df.loc[1, uplift] == approx(0.058885)
 
     assert metrics_df.loc[2, segment_header] == eval_2_segment
     assert metrics_df.loc[2, rows_header] == 250
     assert metrics_df.loc[2, target_mean_header] == 0.536
-    assert metrics_df.loc[2, baseline_gini] == approx(0.022002)
-    assert metrics_df.loc[2, enriched_gini] == approx(-0.010950)
-    assert metrics_df.loc[2, uplift] == approx(-0.032952)
+    assert metrics_df.loc[2, baseline_gini] == approx(-0.066572)
+    assert metrics_df.loc[2, enriched_gini] == approx(-0.116598)
+    assert metrics_df.loc[2, uplift] == approx(-0.050026)
 
 
 @pytest.mark.skip()
@@ -1225,23 +1225,23 @@ def test_rf_metric_rmse(requests_mock: Mocker):
     assert metrics_df.loc[0, segment_header] == train_segment
     assert metrics_df.loc[0, rows_header] == 500
     assert metrics_df.loc[0, target_mean_header] == 0.51
-    assert metrics_df.loc[0, baseline_rmse] == approx(0.737054)
-    assert metrics_df.loc[0, enriched_rmse] == approx(0.720624)
-    assert metrics_df.loc[0, uplift] == approx(0.016430)
+    assert metrics_df.loc[0, baseline_rmse] == approx(0.695490)
+    assert metrics_df.loc[0, enriched_rmse] == approx(0.656957)
+    assert metrics_df.loc[0, uplift] == approx(0.038533)
 
     assert metrics_df.loc[1, segment_header] == eval_1_segment
     assert metrics_df.loc[1, rows_header] == 250
     assert metrics_df.loc[1, target_mean_header] == 0.452
-    assert metrics_df.loc[1, baseline_rmse] == approx(0.704719)
-    assert metrics_df.loc[1, enriched_rmse] == approx(0.721444)
-    assert metrics_df.loc[1, uplift] == approx(-0.016725)
+    assert metrics_df.loc[1, baseline_rmse] == approx(0.717178)
+    assert metrics_df.loc[1, enriched_rmse] == approx(0.685107)
+    assert metrics_df.loc[1, uplift] == approx(0.032071)
 
     assert metrics_df.loc[2, segment_header] == eval_2_segment
     assert metrics_df.loc[2, rows_header] == 250
     assert metrics_df.loc[2, target_mean_header] == 0.536
-    assert metrics_df.loc[2, baseline_rmse] == approx(0.690261)
-    assert metrics_df.loc[2, enriched_rmse] == approx(0.694711)
-    assert metrics_df.loc[2, uplift] == approx(-0.004450)
+    assert metrics_df.loc[2, baseline_rmse] == approx(0.678079)
+    assert metrics_df.loc[2, enriched_rmse] == approx(0.718205)
+    assert metrics_df.loc[2, uplift] == approx(-0.040126)
 
 
 def test_default_metric_binary_with_string_feature(requests_mock: Mocker):
@@ -1341,23 +1341,23 @@ def test_default_metric_binary_with_string_feature(requests_mock: Mocker):
     assert metrics_df.loc[0, segment_header] == train_segment
     assert metrics_df.loc[0, rows_header] == 500
     assert metrics_df.loc[0, target_mean_header] == 0.51
-    assert metrics_df.loc[0, baseline_gini] == approx(0.116841)
-    assert metrics_df.loc[0, enriched_gini] == approx(0.076030)
-    assert metrics_df.loc[0, uplift] == approx(-0.040811)
+    assert metrics_df.loc[0, baseline_gini] == approx(-0.034968)
+    assert metrics_df.loc[0, enriched_gini] == approx(-0.090683)
+    assert metrics_df.loc[0, uplift] == approx(-0.055715)
 
     assert metrics_df.loc[1, segment_header] == eval_1_segment
     assert metrics_df.loc[1, rows_header] == 250
     assert metrics_df.loc[1, target_mean_header] == 0.452
-    assert metrics_df.loc[1, baseline_gini] == approx(-0.078160)
-    assert metrics_df.loc[1, enriched_gini] == approx(-0.029288)
-    assert metrics_df.loc[1, uplift] == approx(0.048873)
+    assert metrics_df.loc[1, baseline_gini] == approx(-0.081674)
+    assert metrics_df.loc[1, enriched_gini] == approx(-0.006627)
+    assert metrics_df.loc[1, uplift] == approx(0.075047)
 
     assert metrics_df.loc[2, segment_header] == eval_2_segment
     assert metrics_df.loc[2, rows_header] == 250
     assert metrics_df.loc[2, target_mean_header] == 0.536
-    assert metrics_df.loc[2, baseline_gini] == approx(-0.013484)
-    assert metrics_df.loc[2, enriched_gini] == approx(-0.017486)
-    assert metrics_df.loc[2, uplift] == approx(-0.004002)
+    assert metrics_df.loc[2, baseline_gini] == approx(-0.039166)
+    assert metrics_df.loc[2, enriched_gini] == approx(-0.016457)
+    assert metrics_df.loc[2, uplift] == approx(0.022710)
 
 
 def approx(value: float):

{upgini-1.1.266a3254.post2 → upgini-1.1.267a3254.post3}/tests/test_target_utils.py

@@ -4,7 +4,8 @@ import pytest
 from pandas.testing import assert_frame_equal
 
 from upgini.errors import ValidationError
-from upgini.metadata import SYSTEM_RECORD_ID, TARGET, ModelTaskType
+from upgini.features_enricher import FeaturesEnricher
+from upgini.metadata import SYSTEM_RECORD_ID, TARGET, ModelTaskType, SearchKey
 from upgini.resource_bundle import bundle
 from upgini.utils.target_utils import balance_undersample, define_task
 
@@ -132,3 +133,62 @@ def test_balance_undersaampling_multiclass():
     })
     # Get all of 25% quantile class (b) and minor classes (a) and x2 (or all if less) of major classes
     assert_frame_equal(balanced_df.sort_values(by=SYSTEM_RECORD_ID).reset_index(drop=True), expected_df)
+
+
+def test_binary_psi_calculation():
+    df = pd.DataFrame({
+        "target": [0, 0, 0, 0, 0, 1, 1, 1, 1, 1,   0, 0, 0, 0, 0, 1, 0, 1, 0, 1]
+    })
+    df["date"] = pd.date_range("2020-01-01", "2020-01-20")
+    enricher = FeaturesEnricher(search_keys={"date": SearchKey.DATE})
+    enricher._validate_PSI(df)
+    assert not enricher.warning_counter.has_warnings()
+
+    df = pd.DataFrame({
+        "target": [0, 0, 0, 0, 0, 1, 1, 1, 1, 1,   0, 0, 0, 0, 0, 0, 0, 1, 0, 1]
+    })
+    df["date"] = pd.date_range("2020-01-01", "2020-01-20")
+    enricher = FeaturesEnricher(search_keys={"date": SearchKey.DATE})
+    enricher._validate_PSI(df)
+    assert enricher.warning_counter._count == 1
+
+    df = pd.DataFrame({
+        "target": [0, 0, 0, 0, 0, 1, 1, 1, 1, 1,   0, 0, 0, 0, 0, 1, 0, 1, 0, 1],
+        "eval_set_index": [0] * 10 + [1] * 10,
+    })
+    df["date"] = pd.date_range("2020-01-01", "2020-01-20")
+    enricher = FeaturesEnricher(search_keys={"date": SearchKey.DATE})
+    enricher._validate_PSI(df)
+    assert enricher.warning_counter._count == 1
+
+    df = pd.DataFrame({
+        "target": [0, 0, 0, 0, 0, 1, 1, 1, 1, 1,   0, 0, 0, 0, 0, 0, 0, 1, 0, 1],
+        "eval_set_index": [0] * 10 + [1] * 10,
+    })
+    df["date"] = pd.date_range("2020-01-01", "2020-01-20")
+    enricher = FeaturesEnricher(search_keys={"date": SearchKey.DATE})
+    enricher._validate_PSI(df)
+    assert enricher.warning_counter._count == 2
+
+
+def test_regression_psi_calculation():
+    random = np.random.RandomState(42)
+    df = pd.DataFrame({
+        "target": random.rand(20)
+    })
+    df["date"] = pd.date_range("2020-01-01", "2020-01-20")
+    enricher = FeaturesEnricher(search_keys={"date": SearchKey.DATE})
+    enricher._validate_PSI(df)
+    assert enricher.warning_counter._count == 1
+
+    values1 = random.rand(10)
+    values2 = values1.copy()
+    values2[0] = 0.0
+    values2[9] = 1.0
+    df = pd.DataFrame({
+        "target": list(values1) + list(values2)
+    })
+    df["date"] = pd.date_range("2020-01-01", "2020-01-20")
+    enricher = FeaturesEnricher(search_keys={"date": SearchKey.DATE})
+    enricher._validate_PSI(df)
+    assert not enricher.warning_counter.has_warnings()