Repository: feature-engine/feature_engine

Issue:
PSI feature selection for categorical features
Currently, PSI feature selection as implemented in selection/drop_psi_features.py can only be applied to numerical features: it first converts these numerical features to categoricals via bucketing, then calculates PSI for the discretised numerical features. 

It would be useful to also be able to use DropHighPSIFeatures on categorical features. I believe the implementation is fairly straightforward.

I'd be happy to pick up this issue if it is of interest to others. 


Discussion:
solegalli: Hey @dlaprins and @ClaudioSalvatoreArcidiacono

I've got the impression that rebasing didn't work in #657 because I can see the files that we merged in #660 in there.

I tried to rebase again, but got a lot of errors, so decided to start afresh.

The changes here in summary:

- I mostly tidied the docstrings
- fixed doc error
- did minor cosmetics in the class code
- added some tests for p-value

Would you be able to have a quick look and if OK we merge?

Thank you!

Patch:
diff --git a/docs/index.rst b/docs/index.rst
index e121856..c2e3088 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -267,4 +267,5 @@ Table of Contents
    contribute/index
    about/index
    whats_new/index
+   versions/index
    donate
diff --git a/docs/versions/index.rst b/docs/versions/index.rst
new file mode 100644
index 0000000..be193f0
--- /dev/null
+++ b/docs/versions/index.rst
@@ -0,0 +1,12 @@
+Other versions
+==============
+
+Web-based documentation is available for versions listed below:
+
+- `Feature-engine 1.6 <https://feature-engine.trainindata.com/en/1.6.x/index.html>`_
+- `Feature-engine 1.5 <https://feature-engine.trainindata.com/en/1.5.x/index.html>`_
+- `Feature-engine 1.4 <https://feature-engine.trainindata.com/en/1.4.x/index.html>`_
+- `Feature-engine 1.3 <https://feature-engine.trainindata.com/en/1.3.x/index.html>`_
+- `Feature-engine 1.2 <https://feature-engine.trainindata.com/en/1.2.x/index.html>`_
+- `Feature-engine 1.1 <https://feature-engine.trainindata.com/en/1.1.x/index.html>`_
+- `Feature-engine 0.6 <https://feature-engine.trainindata.com/en/0.6.x_a/>`_
diff --git a/feature_engine/creation/relative_features.py b/feature_engine/creation/relative_features.py
index 55aa277..b8ee009 100644
--- a/feature_engine/creation/relative_features.py
+++ b/feature_engine/creation/relative_features.py
@@ -77,6 +77,10 @@ class RelativeFeatures(BaseCreation):
         one or more of the following strings: 'add', 'mul','sub', 'div', truediv,
         'floordiv', 'mod', 'pow'.
 
+    fill_value: int, float, default=None
+        When dividing by zero, this value is used in place of infinity. If None,
+        then an error will be raised when dividing by zero.
+
     {missing_values}
 
     {drop_original}
@@ -128,6 +132,7 @@ class RelativeFeatures(BaseCreation):
         variables: List[Union[str, int]],
         reference: List[Union[str, int]],
         func: List[str],
+        fill_value: Union[int, float, None] = None,
         missing_values: str = "ignore",
         drop_original: bool = False,
     ) -> None:
@@ -163,10 +168,16 @@ class RelativeFeatures(BaseCreation):
                 "Supported functions are {}. ".format(", ".join(_PERMITTED_FUNCTIONS))
             )
 
+        if fill_value is not None and not isinstance(fill_value, (float, int)):
+            raise ValueError(
+                "fill_value must be a float, integer or None. "
+                f"Got {fill_value} instead."
+            )
         super().__init__(missing_values, drop_original)
         self.variables = variables
         self.reference = reference
         self.func = func
+        self.fill_value = fill_value
 
     def transform(self, X: pd.DataFrame) -> pd.DataFrame:
         """
@@ -213,17 +224,6 @@ class RelativeFeatures(BaseCreation):
             X[varname] = X[self.variables].sub(X[reference], axis=0)
         return X
 
-    def _div(self, X):
-        for reference in self.reference:
-            if (X[reference] == 0).any():
-                raise ValueError(
-                    "Some of the reference variables contain 0 as values. Check and "
-                    "remove those before using this transformer."
-                )
-            varname = [f"{var}_div_{reference}" for var in self.variables]
-            X[varname] = X[self.variables].div(X[reference], axis=0)
-        return X
-
     def _add(self, X):
         for reference in self.reference:
             varname = [f"{var}_add_{reference}" for var in self.variables]
@@ -236,38 +236,60 @@ class RelativeFeatures(BaseCreation):
             X[varname] = X[self.variables].mul(X[reference], axis=0)
         return X
 
-    def _truediv(self, X):
+    def _div(self, X):
+        for reference in self.reference:
+            zeros_ix, contains_zero = self._find_zeroes_in_reference(X, reference)
 
+            if self.fill_value is None and contains_zero:
+                self._raise_error_when_zero_in_denominator()
+
+            varname = [f"{var}_div_{reference}" for var in self.variables]
+            X[varname] = X[self.variables].div(X[reference], axis=0)
+
+            if contains_zero:
+                X.loc[zeros_ix, varname] = self.fill_value
+        return X
+
+    def _truediv(self, X):
         for reference in self.reference:
-            if (X[reference] == 0).any():
-                raise ValueError(
-                    "Some of the reference variables contain 0 as values. Check and "
-                    "remove those before using this transformer."
-                )
+            zeros_ix, contains_zero = self._find_zeroes_in_reference(X, reference)
+
+            if self.fill_value is None and contains_zero:
+                self._raise_error_when_zero_in_denominator()
+
             varname = [f"{var}_truediv_{reference}" for var in self.variables]
             X[varname] = X[self.variables].truediv(X[reference], axis=0)
+
+            if contains_zero:
+                X.loc[zeros_ix, varname] = self.fill_value
         return X
 
     def _floordiv(self, X):
         for reference in self.reference:
-            if (X[reference] == 0).any():
-                raise ValueError(
-                    "Some of the reference variables contain 0 as values. Check and "
-                    "remove those before using this transformer."
-                )
+            zeros_ix, contains_zero = self._find_zeroes_in_reference(X, reference)
+
+            if self.fill_value is None and contains_zero:
+                self._raise_error_when_zero_in_denominator()
+
             varname = [f"{var}_floordiv_{reference}" for var in self.variables]
             X[varname] = X[self.variables].floordiv(X[reference], axis=0)
+
+            if contains_zero:
+                X.loc[zeros_ix, varname] = self.fill_value
         return X
 
     def _mod(self, X):
         for reference in self.reference:
-            if (X[reference] == 0).any():
-                raise ValueError(
-                    "Some of the reference variables contain 0 as values. Check and "
-                    "remove those before using this transformer."
-                )
+            zeros_ix, contains_zero = self._find_zeroes_in_reference(X, reference)
+
+            if self.fill_value is None and contains_zero:
+                self._raise_error_when_zero_in_denominator()
+
             varname = [f"{var}_mod_{reference}" for var in self.variables]
             X[varname] = X[self.variables].mod(X[reference], axis=0)
+
+            if contains_zero:
+                X.loc[zeros_ix, varname] = self.fill_value
         return X
 
     def _pow(self, X):
@@ -276,6 +298,18 @@ class RelativeFeatures(BaseCreation):
             X[varname] = X[self.variables].pow(X[reference], axis=0)
         return X
 
+    def _raise_error_when_zero_in_denominator(self):
+        raise ValueError(
+            "Some of the reference variables contain zeroes. Division by zero "
+            "does not exist. Replace zeros before using this transformer for division "
+            "or set `fill_value` to a number."
+        )
+
+    def _find_zeroes_in_reference(self, X, var):
+        zero_ix = X[var] == 0
+        zero_bool = (zero_ix).any()
+        return zero_ix, zero_bool
+
     def _get_new_features_name(self) -> List:
         """Return names of the created features."""
 
diff --git a/feature_engine/encoding/rare_label.py b/feature_engine/encoding/rare_label.py
index 4859b2c..c7afb18 100644
--- a/feature_engine/encoding/rare_label.py
+++ b/feature_engine/encoding/rare_label.py
@@ -218,7 +218,7 @@ class RareLabelEncoder(CategoricalInitMixinNA, CategoricalMethodsMixin):
                     "indicated in n_categories. Thus, all categories will be "
                     "considered frequent".format(var)
                 )
-                self.encoder_dict_[var] = X[var].unique()
+                self.encoder_dict_[var] = list(X[var].unique())
 
         self.variables_ = variables_
         self._get_feature_names_in(X)
@@ -247,19 +247,19 @@ class RareLabelEncoder(CategoricalInitMixinNA, CategoricalMethodsMixin):
             _check_optional_contains_na(X, self.variables_)
 
             for feature in self.variables_:
-                X[feature] = np.where(
-                    X[feature].isin(self.encoder_dict_[feature]),
-                    X[feature],
-                    self.replace_with,
-                )
+                if X[feature].dtype == "category":
+                    X[feature] = X[feature].cat.add_categories(self.replace_with)
+                X.loc[
+                    ~X[feature].isin(self.encoder_dict_[feature]), feature
+                ] = self.replace_with
 
         else:
             for feature in self.variables_:
-                X[feature] = np.where(
-                    X[feature].isin(self.encoder_dict_[feature] + [np.nan]),
-                    X[feature],
-                    self.replace_with,
-                )
+                if X[feature].dtype == "category":
+                    X[feature] = X[feature].cat.add_categories(self.replace_with)
+                X.loc[
+                    ~X[feature].isin(self.encoder_dict_[feature] + [np.nan]), feature
+                ] = self.replace_with
 
         return X
 
diff --git a/feature_engine/selection/drop_psi_features.py b/feature_engine/selection/drop_psi_features.py
index 2de3803..1d83784 100644
--- a/feature_engine/selection/drop_psi_features.py
+++ b/feature_engine/selection/drop_psi_features.py
@@ -27,36 +27,63 @@ from feature_engine.discretisation import (
 from feature_engine.selection._docstring import (
     _get_support_docstring,
     _variables_attribute_docstring,
-    _variables_numerical_docstring,
 )
 from feature_engine.selection.base_selector import BaseSelector
+from feature_engine.tags import _return_tags
 from feature_engine.variable_handling._init_parameter_checks import (
     _check_init_parameter_variables,
 )
 from feature_engine.variable_handling.variable_type_selection import (
-    find_or_check_numerical_variables,
+    find_categorical_and_numerical_variables,
 )
 
 Variables = Union[None, int, str, List[Union[str, int]]]
 
+PSI = """PSI = sum ( (test_i - basis_i) x ln(test_i/basis_i) )""".rstrip()
+
 
 @Substitution(
     confirm_variables=_confirm_variables_docstring,
-    variables=_variables_numerical_docstring,
     variables_=_variables_attribute_docstring,
     feature_names_in_=_feature_names_in_docstring,
     n_features_in_=_n_features_in_docstring,
     fit_transform=_fit_transform_docstring,
     get_support=_get_support_docstring,
+    psi=PSI,
 )
 class DropHighPSIFeatures(BaseSelector):
     r"""
-    DropHighPSIFeatures drops features which Population Stability Index (PSI) value is
-    above a given threshold. The PSI of a numerical feature is an indication of the
-    shift in its distribution; a feature with high PSI could therefore be considered
+    DropHighPSIFeatures() drops features which Population Stability Index (PSI) is
+    above a given threshold.
+
+    The PSI is used to compare distributions. Higher PSI values mean greater changes in
+    a feature's distribution. Therefore, a feature with high PSI can be considered
     unstable.
 
-    A bigger PSI value indicates a bigger shift in the feature distribution.
+    To compute the PSI, DropHighPSIFeatures() splits the dataset in two: a basis and
+    a test set. Then, it compares the distribution of each feature between those sets.
+
+    To determine the PSI, continuous features are sorted into discrete intervals, and
+    then, the number of observations per interval are compared between the 2
+    distributions.
+
+    The PSI is calculated as:
+
+    {psi}
+
+    where `basis` and `test` are the 2 datasets, `i` refers to each interval, and then,
+    `test_i` and `basis_i` are the number of observations in interval i in each data
+    set.
+
+    The PSI has traditionally been used to assess changes in distributions of
+    continuous variables.
+
+    In version 1.7, we extended the functionality of DropHighPSIFeatures() to
+    calculate the PSI for categorical features as well. In this case, `i` is each
+    unique category, and `test_i` and `basis_i` are the number of observations in
+    category i.
+
+    **Threshold**
 
     Different thresholds can be used to assess the magnitude of the distribution shift
     according to the PSI value. The most commonly used thresholds are:
@@ -65,16 +92,22 @@ class DropHighPSIFeatures(BaseSelector):
     - Above 25%, the variable has experienced a major shift.
     - Between those two values, the shift is intermediate.
 
-    To compute the PSI the DropHighPSIFeatures splits the dataset in two:
+    **Data split**
 
-    First and foremost, the user should enter one variable which will be used to guide
-    the data split. This variable can be of any data type. If the user does not enter a
-    variable name, DropHighPSIFeatures will use the dataframe index.
+    To compute the PSI, DropHighPSIFeatures() splits the dataset in two: a basis and
+    a test set. Then, it compares the distribution of each feature between those sets.
 
-    Second, the user has the option to specify a proportion of observations to put in
-    each data set, or alternatively, provide a cut-off value.
+    There are various options to split a dataset:
 
-    If the user specifies a proportion through the `split_frac` parameter, the data will
+    First, you can indicate which variable should be used to guide the data split. This
+    variable can be of any data type. If you do not enter a variable name,
+    DropHighPSIFeatures() will use the dataframe index.
+
+    Next, you need to specify how that variable (or the index) should be used to split
+    the data. You can specify a proportion of observations to be put in each data set,
+    or alternatively, provide a cut-off value.
+
+    If you specify a proportion through the `split_frac` parameter, the data will
     be sorted to accommodate that proportion. If `split_frac` is 0.5, 50% of the
     observations will go to either basis or test sets. If `split_frac` is 0.6, 60% of
     the samples will go to the basis data set and the remaining 40% to the test set.
@@ -82,14 +115,13 @@ class DropHighPSIFeatures(BaseSelector):
     If `split_distinct` is True, the data will be sorted considering unique values in
     the selected variables. Check the parameter below for more details.
 
-    If the user defines a numeric cut-off value or a specific date using the `cut_off`
+    If you define a numeric cut-off value or a specific date using the `cut_off`
     parameter, the observations with value <= cut-off will go to the basis data set and
-    the remaining ones to the test set. For categorical values this means they are
-    sorted alphabetically and cut accordingly.
+    the remaining ones to the test set. If the variable used to guide the split is
+    categorical, its values are sorted alphabetically and cut accordingly.
 
-    If the user passes a list of values in the `cut-off`, the observations with the
-    values in the list, will go to the basis set, and the remaining ones to the test
-    set.
+    If you pass a list of values in the `cut-off`, the observations with the values in
+    the list, will go to the basis set, and the remaining ones to the test set.
 
     More details in the :ref:`User Guide <psi_selection>`.
 
@@ -117,13 +149,14 @@ class DropHighPSIFeatures(BaseSelector):
         instead of being applied to the whole vector of values. For example, if the
         values in `split_col` are [1, 1, 1, 1, 2, 2, 3, 4] and `split_frac` is
         0.5, we have the following:
-        - `split_distinct=False` splits the vector in two equally sized parts:
-        [1, 1, 1, 1] and [2, 2, 3, 4]. This involves that 2 dataframes with 4
-        observations each are used for the PSI calculations.
-        - `split_distinct=True` computes the vector of unique values in `split_col`
-        ([1, 2, 3, 4]) and splits that vector in two equal parts: [1, 2] and [3, 4].
-        The number of observations in the two dataframes used for the PSI calculations
-        is respectively 6 ([1, 1, 1, 1, 2, 2]) and 2 ([3, 4]).
+
+            - `split_distinct=False` splits the vector in two equally sized parts:
+                [1, 1, 1, 1] and [2, 2, 3, 4]. This involves that 2 dataframes with 4
+                observations each are used for the PSI calculations.
+            - `split_distinct=True` computes the vector of unique values in `split_col`
+                ([1, 2, 3, 4]) and splits that vector in two equal parts: [1, 2] and
+                [3, 4]. The number of observations in the two dataframes used for the
+                PSI calculations is respectively 6 ([1, 1, 1, 1, 2, 2]) and 2 ([3, 4]).
 
     cut_off: int, float, date or list, default=None
         Threshold to split the dataset based on the `split_col` variable. If int, float
@@ -143,8 +176,20 @@ class DropHighPSIFeatures(BaseSelector):
         The threshold to drop a feature. If the PSI for a feature is >= threshold, the
         feature will be dropped. The most common threshold values are 0.25 (large shift)
         and 0.10 (medium shift).
-        If 'auto', the threshold will be calculated based on the size of the base and
-        target dataset and the number of bins.
+        If 'auto', the threshold will be calculated based on the size of the basis and
+        test dataset and the number of bins as:
+
+                threshold = χ2(q, B−1) × (1/N + 1/M)
+
+        where:
+
+            - q = quantile of the distribution (or 1 - p-value),
+            - B = number of bins/categories,
+            - N = size of basis dataset,
+            - M = size of test dataset.
+
+        See formula (5.2) from reference [1].
+
 
     bins: int, default = 10
         Number of bins or intervals. For continuous features with good value spread, 10
@@ -169,7 +214,17 @@ class DropHighPSIFeatures(BaseSelector):
         when determining the PSI for that particular feature. If 'raise' the transformer
         will raise an error and features will not be selected.
 
-    {variables}
+    p_value: float, default = 0.001
+        The p-value to test the null hypothesis that there is no feature drift. In that
+        case, the PSI-value approximates a random variable that follows a chi-square
+        distribution. See [1] for details. This parameter is used only if `threshold`
+        is set to 'auto'.
+
+    variables: int, str, list, default = None
+        The list of variables to evaluate. If `None`, the transformer will evaluate all
+        numerical variables in the dataset. If `"all"` the transformer will evaluate all
+        categorical and numerical variables in the dataset. Alternatively, the
+        transformer will evaluate the variables indicated in the list or string.
 
     {confirm_variables}
 
@@ -219,8 +274,10 @@ class DropHighPSIFeatures(BaseSelector):
 
     >>> import pandas as pd
     >>> from feature_engine.selection import DropHighPSIFeatures
-    >>> X = pd.DataFrame(dict(x1 = [1,1,0,0,0,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0],
-    >>>                 x2 = [32,87,6,32,11,44,8,7,9,0,32,87,6,32,11,44,8,7,9,0]))
+    >>> X = pd.DataFrame(dict(
+    >>>         x1 = [1,1,0,0,0,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0],
+    >>>         x2 = [32,87,6,32,11,44,8,7,9,0,32,87,6,32,11,44,8,7,9,0],
+    >>>         ))
     >>> psi = DropHighPSIFeatures()
     >>> psi.fit_transform(X)
         x2
@@ -251,6 +308,7 @@ class DropHighPSIFeatures(BaseSelector):
         missing_values: str = "raise",
         variables: Variables = None,
         confirm_variables: bool = False,
+        p_value: float = 0.001,
     ):
 
         if not isinstance(split_col, (str, int, type(None))):
@@ -319,6 +377,11 @@ class DropHighPSIFeatures(BaseSelector):
                     f"or choose another splitting criteria."
                 )
 
+        if not isinstance(p_value, float) or p_value < 0 or p_value > 1:
+            raise ValueError(
+                f"p_value must be a float between 0 and 1. Got {p_value} instead."
+            )
+
         super().__init__(confirm_variables)
 
         # Check the variables before assignment.
@@ -335,6 +398,7 @@ class DropHighPSIFeatures(BaseSelector):
         self.strategy = strategy
         self.min_pct_empty_bins = min_pct_empty_bins
         self.missing_values = missing_values
+        self.p_value = p_value
 
     def fit(self, X: pd.DataFrame, y: pd.Series = None):
         """
@@ -351,21 +415,18 @@ class DropHighPSIFeatures(BaseSelector):
         # check input dataframe
         X = check_X(X)
 
-        # If required exclude variables that are not in the input dataframe
-        self._confirm_variables(X)
-
-        # find numerical variables or check those entered are present in the dataframe
-        self.variables_ = find_or_check_numerical_variables(X, self.variables_)
+        # select variables to evaluate
+        cat_variables_, num_variables_ = self._select_variables(X)
 
-        # Remove the split_col from the variables list. It might be added if the
-        # variables are not defined at initialization.
-        if self.split_col in self.variables_:
-            self.variables_.remove(self.split_col)
+        # check that split column is in the dataframe and remove from variable lists
+        cat_variables_, num_variables_ = self._check_split_column(
+            X, cat_variables_, num_variables_
+        )
 
         if self.missing_values == "raise":
             # check if dataset contains na or inf
-            _check_contains_na(X, self.variables_)
-            _check_contains_inf(X, self.variables_)
+            _check_contains_na(X, num_variables_ + cat_variables_)
+            _check_contains_inf(X, num_variables_)
 
         # Split the dataframe into basis and test.
         basis_df, test_df = self._split_dataframe(X)
@@ -387,26 +448,37 @@ class DropHighPSIFeatures(BaseSelector):
         if self.switch:
             test_df, basis_df = basis_df, test_df
 
-        if self.threshold == "auto":
-            threshold = self._calculate_auto_threshold(
-                basis_df.shape[0], test_df.shape[0]
-            )
-        else:
-            threshold = self.threshold
+        # Set up parameters for numerical features
+        if len(num_variables_) > 0:
+
+            # Set up the discretizer for numerical features
+            if self.strategy == "equal_width":
+                bucketer = EqualWidthDiscretiser(bins=self.bins)
+            else:
+                bucketer = EqualFrequencyDiscretiser(q=self.bins)
+
+            # Set up the threshold for numerical features
+            if self.threshold == "auto":
+                threshold_num = self._calculate_auto_threshold(
+                    basis_df.shape[0],
+                    test_df.shape[0],
+                    self.bins,
+                )
+            else:
+                threshold_num = self.threshold
 
-        # set up the discretizer
-        if self.strategy == "equal_width":
-            bucketer = EqualWidthDiscretiser(bins=self.bins)
-        else:
-            bucketer = EqualFrequencyDiscretiser(q=self.bins)
+        # Set up the generic threshold for categorical features if used
+        if len(cat_variables_) > 0:
+            if self.threshold != "auto":
+                threshold_cat = self.threshold
 
         # Compute the PSI by looping over the features
         self.psi_values_ = {}
         self.features_to_drop_ = []
 
-        for feature in self.variables_:
-            # Discretize the features.
-
+        # Compute PSI for numerical features
+        for feature in num_variables_:
+            # Discretize feature
             basis_discrete = bucketer.fit_transform(basis_df[[feature]].dropna())
             test_discrete = bucketer.transform(test_df[[feature]].dropna())
 
@@ -419,15 +491,118 @@ class DropHighPSIFeatures(BaseSelector):
             self.psi_values_[feature] = np.sum(
                 (test_distrib - basis_distrib) * np.log(test_distrib / basis_distrib)
             )
+
             # Assess if feature should be dropped
-            if self.psi_values_[feature] > threshold:
+            if self.psi_values_[feature] > threshold_num:
                 self.features_to_drop_.append(feature)
 
+        # Compute the PSI for categorical features
+        for feature in cat_variables_:
+            basis_discrete = basis_df[[feature]]
+            test_discrete = test_df[[feature]]
+
+            # Determine percentage of observations per bin
+            basis_distrib, test_distrib = self._observation_frequency_per_bin(
+                basis_discrete, test_discrete
+            )
+
+            # Calculate the PSI value
+            self.psi_values_[feature] = np.sum(
+                (test_distrib - basis_distrib) * np.log(test_distrib / basis_distrib)
+            )
+
+            # Determine the appropriate threshold for the categorical feature
+            if self.threshold == "auto":
+                n_bins_cat = X[feature].nunique()
+                threshold_cat = self._calculate_auto_threshold(
+                    basis_df.shape[0],
+                    test_df.shape[0],
+                    n_bins_cat,
+                )
+
+            # Assess if feature should be dropped
+            if self.psi_values_[feature] > threshold_cat:
+                self.features_to_drop_.append(feature)
+
+        # store analyzed variables
+        self.variables_ = num_variables_ + cat_variables_
         # save input features
         self._get_feature_names_in(X)
 
         return self
 
+    def _select_variables(self, X: pd.DataFrame):
+        """Based on the user input to the `variables` attribute in init, find or check
+        the variables for which the PSI should be calculated.
+
+        If `None`, select all numerical variables.
+        If `"all", select all numerical and categorical variables.
+        If string or list, check that the variables are numerical or categorical.
+        """
+
+        if self.variables is None:
+            num_variables = list(X.select_dtypes(include="number").columns)
+            if len(num_variables) == 0:
+                raise ValueError(
+                    "No numerical variables found in this dataframe. Please check "
+                    "variable format with pandas dtypes."
+                )
+            cat_variables: List[Union[str, int]] = []
+
+        elif self.variables == "all":
+            (
+                cat_variables,
+                num_variables,
+            ) = find_categorical_and_numerical_variables(X, None)
+
+        else:
+            if not isinstance(self.variables, list):
+                variables = [self.variables]
+            else:
+                variables = self.variables
+
+            if self.confirm_variables is True:
+                variables = [var for var in variables if var in X.columns]
+                # Raise an error if no column is left to work with.
+                if len(variables) == 0:
+                    raise ValueError(
+                        "After confirming variables, no variable remains. At least 1 "
+                        "variable is required for the selection."
+                    )
+
+            (
+                cat_variables,
+                num_variables,
+            ) = find_categorical_and_numerical_variables(X, variables)
+
+        return cat_variables, num_variables
+
+    def _check_split_column(
+        self,
+        X: pd.DataFrame,
+        cat_variables: List[Union[str, int]],
+        num_variables: List[Union[str, int]],
+    ):
+        """Check that split_col is in the dataframe and remove from numerical and
+        categorical variable lists if necessary.
+
+        It will get added if the variables are selected automatically.
+        """
+        if self.split_col is not None:
+            # check that split_col is in the dataframe.
+            if self.split_col not in X.columns:
+                raise ValueError(f"{self.split_col} is not in the dataframe.")
+
+            # Remove the split_col from variables lists. Happens when variables are
+            # selected by transformer.
+            if self.variables is None or self.variables == "all":
+                if self.split_col in num_variables:
+                    num_variables.remove(self.split_col)
+                elif self.split_col in cat_variables:
+                    cat_variables.remove(self.split_col)
+
+        return cat_variables, num_variables
+
     def _observation_frequency_per_bin(self, basis, test):
         """
         Obtain the fraction of observations per interval.
@@ -464,12 +639,13 @@ class DropHighPSIFeatures(BaseSelector):
                 how="outer",
             )
             .fillna(self.min_pct_empty_bins)
+            .replace(to_replace=0, value=self.min_pct_empty_bins)
         )
         distributions.columns = ["basis", "test"]
 
         return distributions.basis, distributions.test
 
-    def _split_dataframe(self, X):
+    def _split_dataframe(self, X: pd.DataFrame):
         """
         Split dataframe according to a cut-off value and return two dataframes: the
         basis dataframe contains all observations <= cut_off and the test dataframe the
@@ -495,13 +671,13 @@ class DropHighPSIFeatures(BaseSelector):
         """
 
         # Identify the values according to which the split must be done.
-        if not self.split_col:
+        if self.split_col is None:
             reference = pd.Series(X.index)
         else:
             reference = X[self.split_col]
 
         # Raise an error if there are missing values in the reference column.
-        if reference.isna().sum() != 0:
+        if reference.isna().any():
             raise ValueError(
                 f"There are {reference.isna().sum()} missing values in the reference"
                 "variable. Missing data are not allowed in the variable used to "
@@ -547,7 +723,7 @@ class DropHighPSIFeatures(BaseSelector):
 
         Parameters
         ----------
-        split_column : pd.Series.
+        split_column: pd.Series.
             Series for which the nth quantile will be computed.
 
         Returns
@@ -577,10 +753,37 @@ class DropHighPSIFeatures(BaseSelector):
 
         return cut_off
 
-    def _calculate_auto_threshold(self, N, M, q=0.999):
-        # threshold = χ2(q,B−1) × (1/N + 1/M)
-        # where q - quantile (or 1 - p-value) B - number of bins,
-        # N - size of basis dataset, M - size of test dataset
-        # see formula (5.2) from reference
-        # taking q = 0.999 to get higher threshold
-        return stats.chi2.ppf(q, self.bins - 1) * (1.0 / N + 1.0 / M)
+    def _calculate_auto_threshold(self, N, M, bins):
+        """Threshold computation for chi-square test.
+
+        The threshold is given by:
+
+            threshold = χ2(q,B−1) × (1/N + 1/M)
+
+        where:
+
+        q = quantile of the distribution (or 1 - p-value),
+        B = number of bins/categories,
+        N = size of basis dataset,
+        M = size of test dataset.
+        See formula (5.2) from reference [1] in the class docstring.
+
+        Parameters
+        ----------
+        N: float or int
+        M: float or int
+        bins: int
+
+        Returns
+        -------
+        float
+        """
+        return stats.chi2.ppf(1 - self.p_value, bins - 1) * (1.0 / N + 1.0 / M)
+
+    def _more_tags(self):
+        tags_dict = _return_tags()
+        tags_dict["variables"] = "pass"
+        # add additional test that fails
+        tags_dict["_xfail_checks"]["check_estimators_nan_inf"] = "transformer allows NA"
+
+        return tags_dict


Test patch:
diff --git a/tests/estimator_checks/estimator_checks.py b/tests/estimator_checks/estimator_checks.py
index 17388de..06aec8a 100644
--- a/tests/estimator_checks/estimator_checks.py
+++ b/tests/estimator_checks/estimator_checks.py
@@ -84,6 +84,8 @@ def check_feature_engine_estimator(estimator, needs_group: bool = False):
             check_all_types_variables_assignment(estimator, needs_group=needs_group)
         elif tags["variables"] == "datetime":
             check_datetime_variables_assignment(estimator)
+        else:
+            pass
 
     # Tests based on transformer's init parameters
     if hasattr(estimator, "cv"):
diff --git a/tests/test_creation/test_relative_features.py b/tests/test_creation/test_relative_features.py
index d0db640..6ff77b2 100644
--- a/tests/test_creation/test_relative_features.py
+++ b/tests/test_creation/test_relative_features.py
@@ -53,6 +53,17 @@ def test_error_if_func_not_supported(_func):
         )
 
 
+@pytest.mark.parametrize("_fill_value", [(2, 3.3), ["test"], "python"])
+def test_error_if_fill_value_not_permitted(_fill_value):
+    with pytest.raises(ValueError):
+        RelativeFeatures(
+            variables=["Age"],
+            reference=["Marks"],
+            func=["sub", "div", "add", "mul"],
+            fill_value=_fill_value,
+        )
+
+
 def test_error_when_drop_original_not_bool():
     for drop_original in ["True", [True]]:
         with pytest.raises(ValueError):
@@ -318,7 +329,7 @@ def test_when_df_cols_are_integers(df_vartypes):
 
 
 @pytest.mark.parametrize("_func", [["div"], ["truediv"], ["floordiv"], ["mod"]])
-def test_error_when_division_by_zero(_func, df_vartypes):
+def test_error_when_division_by_zero_and_fill_value_is_none(_func, df_vartypes):
 
     df_zero = df_vartypes.copy()
     df_zero.loc[1, "Marks"] = 0
@@ -329,9 +340,53 @@ def test_error_when_division_by_zero(_func, df_vartypes):
         func=_func,
     )
     transformer.fit(df_vartypes)
-    with pytest.raises(ValueError):
+
+    with pytest.raises(ValueError) as record:
         transformer.transform(df_zero)
 
+    msg = (
+            "Some of the reference variables contain zeroes. Division by zero "
+            "does not exist. Replace zeros before using this transformer for division "
+            "or set `fill_value` to a number."
+        )
+    # check that the error message matches
+    assert str(record.value) == msg
+
+
+@pytest.mark.parametrize("_fill_value, _func", [
+    (111.111, ["div"]),
+    (999, ["div"]),
+    (111.111, ["truediv"]),
+    (999, ["truediv"]),
+    (111.111, ["floordiv"]),
+    (999, ["floordiv"]),
+    (111.111, ["mod"]),
+    (999, ["mod"]),
+])
+def test_fill_values_when_division_by_zero(
+        _fill_value, _func, df_vartypes
+):
+    df_zero = df_vartypes.copy()
+    df_zero.loc[2, "Marks"] = 0
+    df_zero.loc[1, "Age"] = np.nan
+    df_zero.loc[3, "Age"] = np.inf
+
+    transformer = RelativeFeatures(
+        variables=["Age"],
+        reference=["Marks"],
+        fill_value=_fill_value,
+        func=_func,
+        missing_values="ignore",
+    )
+
+    X = transformer.fit_transform(df_zero)
+
+    new_var = f"Age_{_func[0]}_Marks"
+
+    assert X.loc[2, new_var] == _fill_value
+    np.testing.assert_equal(X.loc[1, "Age"], np.nan)
+    np.testing.assert_equal(X.loc[3, "Age"], np.inf)
+
 
 @pytest.mark.parametrize("_drop", [True, False])
 def test_get_feature_names_out(_drop, df_vartypes):
diff --git a/tests/test_encoding/test_rare_label_encoder.py b/tests/test_encoding/test_rare_label_encoder.py
index 1a17553..ba85299 100644
--- a/tests/test_encoding/test_rare_label_encoder.py
+++ b/tests/test_encoding/test_rare_label_encoder.py
@@ -54,6 +54,34 @@ def test_defo_params_plus_automatically_find_variables(df_enc_big):
     pd.testing.assert_frame_equal(X, df)
 
 
+def test_when_varnames_are_numbers(df_enc_big):
+    input_df = df_enc_big.copy()
+    input_df.columns = [1, 2, 3]
+
+    encoder = RareLabelEncoder(
+        tol=0.06, n_categories=5, variables=None, replace_with="Rare"
+    )
+    X = encoder.fit_transform(input_df)
+
+    # expected output
+    df = {
+        1: ["A"] * 6 + ["B"] * 10 + ["C"] * 4 + ["D"] * 10 + ["Rare"] * 4 + ["G"] * 6,
+        2: ["A"] * 10 + ["B"] * 6 + ["C"] * 4 + ["D"] * 10 + ["Rare"] * 4 + ["G"] * 6,
+        3: ["A"] * 4 + ["B"] * 6 + ["C"] * 10 + ["D"] * 10 + ["Rare"] * 4 + ["G"] * 6,
+    }
+    df = pd.DataFrame(df)
+
+    frequenc_cat = {
+        1: ["B", "D", "A", "G", "C"],
+        2: ["A", "D", "B", "G", "C"],
+        3: ["C", "D", "B", "G", "A"],
+    }
+
+    assert encoder.variables_ == [1, 2, 3]
+    assert encoder.encoder_dict_ == frequenc_cat
+    pd.testing.assert_frame_equal(X, df)
+
+
 def test_correctly_ignores_nan_in_transform(df_enc_big):
     encoder = RareLabelEncoder(
         tol=0.06,
@@ -102,7 +130,7 @@ def test_correctly_ignores_nan_in_fit(df_enc_big):
         n_categories=3,
         missing_values="ignore",
     )
-    X = encoder.fit_transform(df)
+    encoder.fit(df)
 
     # expected:
     frequenc_cat = {
@@ -134,6 +162,90 @@ def test_correctly_ignores_nan_in_fit(df_enc_big):
     pd.testing.assert_frame_equal(X, tt)
 
 
+def test_correctly_ignores_nan_in_fit_when_var_is_numerical(df_enc_big):
+
+    df = df_enc_big.copy()
+    df["var_C"] = [
+        1,
+        1,
+        1,
+        1,
+        2,
+        2,
+        2,
+        2,
+        2,
+        2,
+        3,
+        3,
+        3,
+        3,
+        3,
+        3,
+        3,
+        3,
+        3,
+        3,
+        4,
+        4,
+        4,
+        4,
+        4,
+        4,
+        4,
+        4,
+        4,
+        4,
+        5,
+        5,
+        6,
+        6,
+        np.nan,
+        np.nan,
+        np.nan,
+        np.nan,
+        np.nan,
+        np.nan,
+    ]
+
+    encoder = RareLabelEncoder(
+        tol=0.06,
+        n_categories=3,
+        missing_values="ignore",
+        ignore_format=True,
+    )
+    encoder.fit(df)
+
+    # expected:
+    frequenc_cat = {
+        "var_A": ["B", "D", "A", "G", "C"],
+        "var_B": ["A", "D", "B", "G", "C"],
+        "var_C": [3, 4, 2, 1],
+    }
+    assert encoder.encoder_dict_ == frequenc_cat
+
+    # input
+    t = pd.DataFrame(
+        {
+            "var_A": ["A", np.nan, "J", "G"],
+            "var_B": ["A", np.nan, "J", "G"],
+            "var_C": [3, np.nan, 9, 10],
+        }
+    )
+
+    # expected
+    tt = pd.DataFrame(
+        {
+            "var_A": ["A", np.nan, "Rare", "G"],
+            "var_B": ["A", np.nan, "Rare", "G"],
+            "var_C": [3.0, np.nan, "Rare", "Rare"],
+        }
+    )
+
+    X = encoder.transform(t)
+    pd.testing.assert_frame_equal(X, tt, check_dtype=False)
+
+
 def test_user_provides_grouping_label_name_and_variable_list(df_enc_big):
     # test case 2: user provides alternative grouping value and variable list
     encoder = RareLabelEncoder(
@@ -316,12 +428,84 @@ def test_variables_cast_as_category(df_enc_big):
         + ["G"] * 6,
     }
     df = pd.DataFrame(df)
+    df["var_B"] = pd.Categorical(df["var_B"])
 
     # test fit attr
     assert encoder.variables_ == ["var_A", "var_B", "var_C"]
     assert encoder.n_features_in_ == 3
     # test transform output
-    pd.testing.assert_frame_equal(X, df)
+    pd.testing.assert_frame_equal(X, df, check_categorical=False)
+
+
+def test_variables_cast_as_category_with_na_in_transform(df_enc_big):
+    encoder = RareLabelEncoder(
+        tol=0.06,
+        n_categories=5,
+        variables=None,
+        replace_with="Rare",
+        missing_values="ignore",
+    )
+
+    df_enc_big = df_enc_big.copy()
+    df_enc_big["var_B"] = df_enc_big["var_B"].astype("category")
+    encoder.fit(df_enc_big)
+
+    # input
+    t = pd.DataFrame(
+        {
+            "var_A": ["A", np.nan, "J", "G"],
+            "var_B": ["A", np.nan, "J", "G"],
+            "var_C": ["A", np.nan, "J", "G"],
+        }
+    )
+    t["var_B"] = pd.Categorical(t["var_B"])
+
+    # expected
+    tt = pd.DataFrame(
+        {
+            "var_A": ["A", np.nan, "Rare", "G"],
+            "var_B": ["A", np.nan, "Rare", "G"],
+            "var_C": ["A", np.nan, "Rare", "G"],
+        }
+    )
+    tt["var_B"] = pd.Categorical(tt["var_B"])
+    pd.testing.assert_frame_equal(encoder.transform(t), tt, check_categorical=False)
+
+
+def test_variables_cast_as_category_with_na_in_fit(df_enc_big):
+
+    df = df_enc_big.copy()
+    df.loc[df["var_C"] == "G", "var_C"] = np.nan
+    df["var_C"] = df["var_C"].astype("category")
+
+    encoder = RareLabelEncoder(
+        tol=0.06,
+        n_categories=3,
+        missing_values="ignore",
+    )
+    encoder.fit(df)
+
+    # input
+    t = pd.DataFrame(
+        {
+            "var_A": ["A", np.nan, "J", "G"],
+            "var_B": ["A", np.nan, "J", "G"],
+            "var_C": ["C", np.nan, "J", "G"],
+        }
+    )
+    t["var_C"] = pd.Categorical(t["var_C"])
+
+    # expected
+    tt = pd.DataFrame(
+        {
+            "var_A": ["A", np.nan, "Rare", "G"],
+            "var_B": ["A", np.nan, "Rare", "G"],
+            "var_C": ["C", np.nan, "Rare", "Rare"],
+        }
+    )
+    tt["var_C"] = pd.Categorical(tt["var_C"])
+
+    pd.testing.assert_frame_equal(encoder.transform(t), tt, check_categorical=False)
 
 
 def test_inverse_transform_raises_not_implemented_error(df_enc_big):
diff --git a/tests/test_selection/test_drop_high_psi_features.py b/tests/test_selection/test_drop_high_psi_features.py
index 004a93b..02ea93b 100644
--- a/tests/test_selection/test_drop_high_psi_features.py
+++ b/tests/test_selection/test_drop_high_psi_features.py
@@ -1,6 +1,6 @@
-from datetime import date
-
 import math
+from datetime import date, datetime
+
 import numpy as np
 import pandas as pd
 import pytest
@@ -26,9 +26,16 @@ def df():
     colnames = ["var_" + str(i) for i in range(6)]
     X = pd.DataFrame(X, columns=colnames)
 
+    # Add a categorical column that does not drift
+    X["cat_1"] = ["A", "B"] * int(X.shape[0] / 2)
+
     # Add drifted features that will be dropped during transformation.
     X["drift_1"] = [number for number in range(X.shape[0])]
     X["drift_2"] = [number / 2 for number in range(X.shape[0])]
+    X["drift_cat_1"] = ["A" for _ in range(int(X.shape[0] / 2))] + [
+        "B" for _ in range(int(X.shape[0] / 2))
+    ]
+    X["drift_cat_1"] = X["drift_cat_1"].astype("category")
 
     return X
 
@@ -40,17 +47,71 @@ def df_mixed_types():
             "A": [it for it in range(0, 20)],
             "B": [1, 2, 2, 1] * 5,
             "C": ["A", "B", "D", "D"] * 5,
-            "time": [date(2019, 1, it + 1) for it in range(20)],
+            "time": [datetime(2019, 1, it + 1) for it in range(20)],
         }
     )
 
     return df
 
 
+EXPECTED_PSI_NUM = {
+    "var_0": 0.043828484052281,
+    "var_1": 0.040929870747665395,
+    "var_2": 0.04330418495156895,
+    "var_3": 0.03773286532548153,
+    "var_4": 0.05047388515663041,
+    "var_5": 0.014717735595712466,
+    "drift_1": 8.283089355027482,
+    "drift_2": 8.283089355027482,
+}
+
+EXPECTED_PSI_ALL = {
+    "var_0": 0.043828484052281,
+    "var_1": 0.040929870747665395,
+    "var_2": 0.04330418495156895,
+    "var_3": 0.03773286532548153,
+    "var_4": 0.05047388515663041,
+    "var_5": 0.014717735595712466,
+    "drift_1": 8.283089355027482,
+    "drift_2": 8.283089355027482,
+    "cat_1": 0.0,
+    "drift_cat_1": 18.41883867587797,
+}
+
+EXPECTED_PSI_NUM_FEW = {
+    "var_2": 0.04330418495156895,
+    "var_3": 0.03773286532548153,
+    "drift_1": 8.283089355027482,
+    "drift_2": 8.283089355027482,
+}
+
+EXPECTED_PSI_MIXED = {
+    "var_0": 0.043828484052281,
+    "drift_1": 8.283089355027482,
+    "drift_cat_1": 18.41883867587797,
+}
+
+EXPECTED_PSI_CAT_FEW = {
+    "cat_1": 0.0,
+    "drift_cat_1": 18.41883867587797,
+}
+
+EXPECTED_PSI_STRING = {"var_0": 0.043828484052281}
+
+_input_output = [
+    (None, EXPECTED_PSI_NUM),
+    ("all", EXPECTED_PSI_ALL),
+    (["var_2", "var_3", "drift_1", "drift_2"], EXPECTED_PSI_NUM_FEW),
+    (["cat_1", "drift_cat_1"], EXPECTED_PSI_CAT_FEW),
+    (["var_0", "drift_1", "drift_cat_1"], EXPECTED_PSI_MIXED),
+    ("var_0", EXPECTED_PSI_STRING),
+]
+
+
 # ====  test  main functionality of the class ====
-def test_fit_attributes(df):
+@pytest.mark.parametrize("variables, expected_psi", _input_output)
+def test_fit_attributes(variables, expected_psi, df):
     """Check the value of the fit attributes.
-
     The expected PSI values used in the assertion were determined using
     the Probatus package.
     ```
@@ -63,67 +124,38 @@ def test_fit_attributes(df):
     psi = psi_calculator.compute(train_df, test_df)
     ```
     """
-    transformer = DropHighPSIFeatures()
-    transformer.fit_transform(df)
-
-    expected_psi = {
-        "var_0": 0.043828484052281,
-        "var_1": 0.040929870747665395,
-        "var_2": 0.04330418495156895,
-        "var_3": 0.03773286532548153,
-        "var_4": 0.05047388515663041,
-        "var_5": 0.014717735595712466,
-        "drift_1": 8.283089355027482,
-        "drift_2": 8.283089355027482,
-    }
+    transformer = DropHighPSIFeatures(variables=variables)
+    dft = transformer.fit_transform(df)
 
-    assert transformer.variables_ == [
-        "var_0",
-        "var_1",
-        "var_2",
-        "var_3",
-        "var_4",
-        "var_5",
-        "drift_1",
-        "drift_2",
-    ]
+    assert transformer.variables_ == list(expected_psi.keys())
     assert transformer.psi_values_ == pytest.approx(expected_psi, 12)
-    assert transformer.features_to_drop_ == ["drift_1", "drift_2"]
+    assert transformer.features_to_drop_ == [
+        var for var in expected_psi.keys() if "drift" in var
+    ]
+    pd.testing.assert_frame_equal(dft, df.drop(transformer.features_to_drop_, axis=1))
 
 
 def test_auto_threshold_calculation():
-    """Check the results of 'auto' threshold calculation
-    """
-    transformer = DropHighPSIFeatures(threshold='auto', bins=10)
+    """Check the results of 'auto' threshold calculation"""
+    transformer = DropHighPSIFeatures(threshold="auto", p_value=0.001, bins=10)
     assert math.isclose(
-        transformer._calculate_auto_threshold(
-            N=500,
-            M=500,
-            q=0.999
-        ),
-        0.11150865948502628
+        transformer._calculate_auto_threshold(N=500, M=500, bins=10),
+        0.11150865948502628,
     )
-    transformer = DropHighPSIFeatures(threshold='auto', bins=32)
+    transformer = DropHighPSIFeatures(threshold="auto", p_value=0.05, bins=32)
     assert math.isclose(
-        transformer._calculate_auto_threshold(
-            N=1000,
-            M=1500,
-            q=0.95
-        ),
-        0.07497557213394188
+        transformer._calculate_auto_threshold(N=1000, M=1500, bins=32),
+        0.07497557213394188,
     )
-    transformer = DropHighPSIFeatures(threshold='auto', bins=42)
+    transformer = DropHighPSIFeatures(threshold="auto", p_value=0.01, bins=42)
     assert math.isclose(
-        transformer._calculate_auto_threshold(
-            N=777,
-            M=666,
-            q=0.99
-        ),
-        0.18111345503169146
+        transformer._calculate_auto_threshold(N=777, M=666, bins=42),
+        0.18111345503169146,
     )
 
 
-def test_fit_attributes_auto(df):
+@pytest.mark.parametrize("variables, expected_psi", _input_output)
+def test_fit_attributes_with_autothreshold(variables, expected_psi, df):
     """Check the value of the fit attributes.
     The expected PSI values used in the assertion were determined using
     the Probatus package.
@@ -137,31 +169,34 @@ def test_fit_attributes_auto(df):
     psi = psi_calculator.compute(train_df, test_df)
     ```
     """
-    transformer = DropHighPSIFeatures(threshold='auto', bins=10)
-    transformer.fit_transform(df)
-
-    expected_psi = {
-        "var_0": 0.043828484052281,
-        "var_1": 0.040929870747665395,
-        "var_2": 0.04330418495156895,
-        "var_3": 0.03773286532548153,
-        "var_4": 0.05047388515663041,
-        "var_5": 0.014717735595712466,
+    transformer = DropHighPSIFeatures(threshold="auto", variables=variables, bins=10)
+    transformer.fit(df)
+
+    assert transformer.psi_values_ == pytest.approx(expected_psi, 12)
+    assert transformer.features_to_drop_ == [
+        var for var in expected_psi.keys() if "drift" in var
+    ]
+
+
+def test_calculation_when_strategy_equal_width(df):
+    transformer = DropHighPSIFeatures(strategy="equal_width")
+    transformer.fit(df)
+
+    expected = {
+        "var_0": 0.014858665472468786,
+        "var_1": 0.04514737836588022,
+        "var_2": 0.03431479397506742,
+        "var_3": 0.04298209189840294,
+        "var_4": 0.02385796430263416,
+        "var_5": 0.046809664317794444,
         "drift_1": 8.283089355027482,
         "drift_2": 8.283089355027482,
     }
-    assert transformer.variables_ == [
-        "var_0",
-        "var_1",
-        "var_2",
-        "var_3",
-        "var_4",
-        "var_5",
-        "drift_1",
-        "drift_2",
+
+    assert transformer.psi_values_ == pytest.approx(expected, 12)
+    assert transformer.features_to_drop_ == [
+        var for var in expected.keys() if "drift" in var
     ]
-    assert transformer.psi_values_ == pytest.approx(expected_psi, 12)
-    assert transformer.features_to_drop_ == ["drift_1", "drift_2"]
 
 
 # ================ test init parameters =================
@@ -180,6 +215,7 @@ default_dict = {
     "min_pct_empty_bins": 0.0001,
     "missing_values": "raise",
     "variables": None,
+    "p_value": 0.001,
 }
 
 args_dict = {
@@ -194,6 +230,7 @@ args_dict = {
     "min_pct_empty_bins": 0.1,
     "missing_values": "ignore",
     "variables": ["chau", "adios"],
+    "p_value": 0.2,
 }
 
 init_dict = [(None, default_dict), (args_dict, args_dict)]
@@ -249,35 +286,109 @@ def test_init_value_error_is_raised():
         DropHighPSIFeatures(min_pct_empty_bins=-1)
 
 
+@pytest.mark.parametrize("p_value", ["hola", -1, 10])
+def test_p_value_not_allowed(p_value):
+    with pytest.raises(ValueError):
+        DropHighPSIFeatures(p_value=p_value)
+
+
 # ================= test fit() functionality ==================
 
 
 def test_split_col_not_included_in_variables(df):
-    """Check that the split columns is not included among the features
+    """Check that the split column is not included among the features
     to evaluate when these are selected automatically."""
     transformer = DropHighPSIFeatures(split_col="var_3", variables=None)
     transformer.fit(df)
+    assert "var_3" not in transformer.variables_
+    assert "var_3" not in transformer.psi_values_.keys()
 
-    assert transformer.variables is None
+    transformer = DropHighPSIFeatures(split_col="var_3", variables="all")
+    transformer.fit(df)
     assert "var_3" not in transformer.variables_
     assert "var_3" not in transformer.psi_values_.keys()
 
+    transformer = DropHighPSIFeatures(split_col="cat_1", variables="all")
+    transformer.fit(df)
+    assert "cat_1" not in transformer.variables_
+    assert "cat_1" not in transformer.psi_values_.keys()
+
+
+def test_error_split_col_not_in_df(df):
+    transformer = DropHighPSIFeatures(variables=None, split_col="var_0")
+    data = df.copy()
+    data = data.drop(["var_0"], axis=1)
+    msg = "var_0 is not in the dataframe."
+    with pytest.raises(ValueError) as record:
+        transformer.fit(data)
+    assert str(record.value) == msg
+
+
+_input_output = [
+    (["var_2", "var_3", "drift_1", "drift_2"], ["drift_1"], EXPECTED_PSI_NUM_FEW),
+    (["var_0", "drift_1", "drift_cat_1"], ["drift_cat_1"], EXPECTED_PSI_MIXED),
+]
+
+
+@pytest.mark.parametrize("variables, variable, expected_psi", _input_output)
+def test_confirm_variables(variables, variable, expected_psi, df):
+    data = df.copy()
+    data = data.drop(variable, axis=1)
+    del expected_psi[variable[0]]
+
+    transformer = DropHighPSIFeatures(variables=variables, confirm_variables=True)
+    transformer.fit(data)
+
+    assert transformer.variables_ == list(expected_psi.keys())
+    assert transformer.psi_values_ == pytest.approx(expected_psi, 12)
+    assert transformer.features_to_drop_ == [
+        var for var in expected_psi.keys() if "drift" in var
+    ]
+
+
+def test_error_if_variables_is_none_and_no_numerical_in_df(df):
+    transformer = DropHighPSIFeatures(variables=None)
+    msg = (
+        "No numerical variables found in this dataframe. Please check "
+        "variable format with pandas dtypes."
+    )
+    with pytest.raises(ValueError) as record:
+        transformer.fit(df[["cat_1", "drift_cat_1"]])
+
+    assert str(record.value) == msg
+
+
+def test_error_if_confirm_variables_returns_empty_list(df):
+    transformer = DropHighPSIFeatures(
+        variables=["cat_1", "drift_cat_1"], confirm_variables=True
+    )
+    data = df.copy()
+    data = data.drop(["cat_1", "drift_cat_1"], axis=1)
+    msg = (
+        "After confirming variables, no variable remains. At least 1 "
+        "variable is required for the selection."
+    )
+    with pytest.raises(ValueError) as record:
+        transformer.fit(data)
+
+    assert str(record.value) == msg
+
 
 def test_error_if_na_in_split_col(df):
     """Test an error is raised if the split column contains missing values."""
     data = df.copy()
-    data.iloc[15, data.columns.get_loc("var_3")] = np.nan
+    data.loc[15, "var_3"] = np.nan
 
     transformer = DropHighPSIFeatures(split_col="var_3")
 
     with pytest.raises(ValueError):
-        transformer.fit_transform(data)
+        transformer.fit(data)
 
 
 def test_raise_error_if_na_in_df(df):
     """Test an error is raised when missing values is set to raise."""
     data = df.copy()
-    data.iloc[15, data.columns.get_loc("var_3")] = np.nan
+    data.loc[15, "var_3"] = np.nan
 
     transformer = DropHighPSIFeatures(missing_values="raise")
 
@@ -288,20 +399,24 @@ def test_raise_error_if_na_in_df(df):
 def test_missing_value_ignored(df):
     """Test if PSI are computed when missing values are present in the dataframe."""
     data = df.copy()
-    data.iloc[15, data.columns.get_loc("var_3")] = np.nan
-
-    var_col = [col for col in data if "var" in col]
+    data.loc[15, "var_3"] = np.nan
 
     transformer = DropHighPSIFeatures(missing_values="ignore")
     transformed = transformer.fit_transform(data)
 
-    pd.testing.assert_frame_equal(transformed, data[var_col])
+    assert transformer.psi_values_ == pytest.approx(EXPECTED_PSI_NUM, 12)
+    assert transformer.features_to_drop_ == [
+        var for var in EXPECTED_PSI_NUM if "drift" in var
+    ]
+    pd.testing.assert_frame_equal(
+        transformed, data.drop(transformer.features_to_drop_, axis=1)
+    )
 
 
 def test_raise_error_if_inf_in_df(df):
     """Test an error is raised for inf when missing values is set to raise."""
     data = df.copy()
-    data.iloc[15, data.columns.get_loc("var_3")] = np.nan
+    data.loc[15, "var_3"] = np.inf
 
     transformer = DropHighPSIFeatures(missing_values="raise")
 
@@ -413,24 +528,35 @@ def test_calculation_df_split_with_different_variable_types(df_mixed_types):
     results = {}
     cut_offs = {}
     for split_col in df_mixed_types.columns:
-        test = DropHighPSIFeatures(split_frac=0.5, split_col=split_col)
+        test = DropHighPSIFeatures(split_frac=0.5, split_col=split_col, variables="all")
         test.fit_transform(df_mixed_types)
         results[split_col] = test.psi_values_
         cut_offs[split_col] = test.cut_off_
 
-    assert results["A"] == pytest.approx({"B": 0.0}, 12)
-    assert results["B"] == pytest.approx({"A": 3.0375978817052403}, 12)
+    assert results["A"] == pytest.approx({"B": 0.0, "C": 0.1621860432432657}, 12)
+    assert results["B"] == pytest.approx(
+        {"A": 3.0375978817052403, "C": 8.515489752777954}, 12
+    )
     assert results["C"] == pytest.approx({"A": 2.27819841127893, "B": 0.0}, 12)
-    assert results["time"] == pytest.approx({"A": 8.283089355027482, "B": 0.0}, 12)
+    assert results["time"] == pytest.approx(
+        {"A": 8.283089355027482, "B": 0.0, "C": 0.1621860432432657}, 12
+    )
 
-    expected_cut_offs = {"A": 9.5, "B": 1.5, "C": "B", "time": date(2019, 1, 10)}
+    expected_cut_offs = {
+        "A": 9.5,
+        "B": 1.5,
+        "C": "B",
+        "time": np.datetime64(datetime(2019, 1, 10)),
+    }
 
     assert cut_offs == expected_cut_offs
 
-    # Test when no data frame with mixed data types when no split_col is provided.
-    test = DropHighPSIFeatures(split_frac=0.5)
+    # Test when no dataframe with mixed data types when no split_col is provided.
+    test = DropHighPSIFeatures(split_frac=0.5, variables="all")
     test.fit_transform(df_mixed_types)
-    assert test.psi_values_ == pytest.approx({"A": 8.283089355027482, "B": 0.0}, 12)
+    assert test.psi_values_ == pytest.approx(
+        {"A": 8.283089355027482, "B": 0.0, "C": 0.1621860432432657}, 12
+    )
 
 
 # =========== tests for user entered cut_off values ===========
@@ -441,7 +567,7 @@ type_test = [
     ("B", 1, [0, 3, 4, 7, 8, 11, 12, 15, 16, 19]),
     ("C", ["B"], [1, 5, 9, 13, 17]),
     ("C", "B", [0, 1, 4, 5, 8, 9, 12, 13, 16, 17]),
-    ("time", date(2019, 1, 4), [0, 1, 2, 3]),
+    ("time", datetime(2019, 1, 4), [0, 1, 2, 3]),
 ]
 
 
@@ -494,7 +620,6 @@ split_distinct_test = [
 @pytest.mark.parametrize("col, expected_index", split_distinct_test)
 def test_split_distinct(col, expected_index):
     """Test the cut off for different data types.
-
     For columns B, C and time we have 6 distinct values, 5 appearing 20 times and
     1 appearing 100 times. A 50% split based on the number of values will result
     in 2 groups of 3. One has 60 appearances (in total) and the other has 140.
@@ -640,48 +765,6 @@ def test_observation_frequency_per_bin():
     )
 
 
-def test_param_variable_definition(df):
-    """Test defining the subset of features through the variable argument.
-
-    Due to the small split_frac value, all variables will fail the PSI test, that is,
-    the variables in the list will show a high PSI value and this will be removed.
-
-    The aim of the test is to show that only those variables defined in the variable
-    argument are examined.
-    """
-    #
-    select = DropHighPSIFeatures(variables=["var_1", "var_3", "var_5"], split_frac=0.01)
-    transformed_df = select.fit_transform(df)
-
-    assert select.variables == ["var_1", "var_3", "var_5"]
-    assert select.variables_ == ["var_1", "var_3", "var_5"]
-    assert list(select.psi_values_.keys()) == ["var_1", "var_3", "var_5"]
-    assert select.features_to_drop_ == ["var_1", "var_3", "var_5"]
-
-    assert transformed_df.columns.to_list() == [
-        "var_0",
-        "var_2",
-        "var_4",
-        "drift_1",
-        "drift_2",
-    ]
-
-
-def test_transform_standard(df):
-    """Test the transform method in a standard approach."""
-    test = DropHighPSIFeatures()
-    test.fit(df)
-    transformed = test.transform(df)
-
-    # Check the features to drop
-    assert test.features_to_drop_ == ["drift_1", "drift_2"]
-
-    # Check the transformed dataframe
-    pd.testing.assert_frame_equal(
-        transformed, df[[col for col in df if "drift" not in col]]
-    )
-
-
 def test_transform_feature_to_drop_not_present(df):
     """Test transform when the feature to drop in not in the dataframe."""
     test = DropHighPSIFeatures()
@@ -694,7 +777,6 @@ def test_transform_feature_to_drop_not_present(df):
     data = data.drop("drift_1", axis=1)
 
     with pytest.raises(KeyError):
-        # Transform
         test.transform(data)