[ENH] Add keywords parameter to StringSimilarityEncoder

feature_engine/encoding/similarity_encoder.py

@@ -1,5 +1,5 @@
 from difflib import SequenceMatcher
-from typing import List, Optional, Union
+from typing import Optional, Union, List
 
 import numpy as np
 import pandas as pd
@@ -102,14 +102,22 @@ class StringSimilarityEncoder(CategoricalInitMixin, CategoricalMethodsMixin):
         categories to encode. In this case, similarity variables will be created
         only for those popular categories.
 
-    missing_values : str, default='impute'
+    missing_values: str, default='impute'
         Indicates if missing values should be ignored, raised or imputed. If 'raise' the
         transformer will return an error if the datasets to `fit` or `transform`
         contain missing values. If 'ignore', missing data will be ignored when learning
         parameters or performing the transformation. If 'impute', the transformer will
         replace missing values with an empty string, '', and then return the similarity
         measures.
 
+    keywords: dict, default=None
+        Dictionary with a set of keywords to be used to create the similarity variables.
+        The format should be: dict(feature: [keyword1, keyword2, ...]). The encoder will
+        use these keywords to create the similarity variables. The dictionary can be
+        defined for all the features to encode, or only for a subset of them. In this
+        case, for the features not specified in the dictionary, the encoder will
+        identify the categories from the data.
+
     {variables}
 
     {ignore_format}
@@ -175,7 +183,8 @@ class StringSimilarityEncoder(CategoricalInitMixin, CategoricalMethodsMixin):
 
     def __init__(
         self,
-        top_categories: Union[None, int] = None,
+        top_categories: Optional[int] = None,
+        keywords: Optional[dict] = None,
         missing_values: str = "impute",
         variables: Union[None, int, str, List[Union[str, int]]] = None,
         ignore_format: bool = False,
@@ -189,9 +198,19 @@ def __init__(
                 "missing_values should be one of 'raise', 'impute' or 'ignore'."
                 f" Got {missing_values!r} instead."
             )
+        if keywords and not isinstance(keywords, dict):
+            raise ValueError(
+                f"keywords should be a dictionary or None. Got {keywords!r} instead."
+            )
+        if keywords and not all(isinstance(item, list) for item in keywords.values()):
+            raise ValueError(
+                "The items in keywords should be lists."
+                f" Got {keywords.values()!r} instead."
+            )
         super().__init__(variables, ignore_format)
         self.top_categories = top_categories
         self.missing_values = missing_values
+        self.keywords = keywords
 
     def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
         """
@@ -213,11 +232,22 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
         X = check_X(X)
         self._check_or_select_variables(X)
         self._get_feature_names_in(X)
+        if self.keywords:
+            if not all(item in self.variables_ for item in self.keywords.keys()):
+                raise ValueError(
+                    "There are variables in keywords that are not present "
+                    "in the dataset."
+                )
         self.encoder_dict_ = {}
 
+        if self.keywords:
+            self.encoder_dict_.update(self.keywords)
+            cols_to_iterate = [x for x in self.variables_ if x not in self.keywords]
+        else:
+            cols_to_iterate = self.variables_
         if self.missing_values == "raise":
             _check_contains_na(X, self.variables_)
-            for var in self.variables_:
+            for var in cols_to_iterate:
                 self.encoder_dict_[var] = (
                     X[var]
                     .astype(str)
@@ -226,7 +256,7 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
                     .index.tolist()
                 )
         elif self.missing_values == "impute":
-            for var in self.variables_:
+            for var in cols_to_iterate:
                 self.encoder_dict_[var] = (
                     X[var]
                     .astype(str)
@@ -236,11 +266,12 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
                     .index.tolist()
                 )
         elif self.missing_values == "ignore":
-            for var in self.variables_:
+            for var in cols_to_iterate:
                 self.encoder_dict_[var] = (
                     X[var]
                     .astype(str)
                     .value_counts(dropna=True)
+                    .drop("nan", errors="ignore")
                     .head(self.top_categories)
                     .index.tolist()
                 )
@@ -288,7 +319,7 @@ def transform(self, X: pd.DataFrame) -> pd.DataFrame:
 
         return X.drop(self.variables_, axis=1)
 
-    def _get_new_features_name(self) -> List:
+    def _get_new_features_name(self) -> List[str]:
         """Return names of the created features."""
         feature_names = []
         for feature in self.variables_:
@@ -300,7 +331,7 @@ def _get_new_features_name(self) -> List:
 
         return feature_names
 
-    def _add_new_feature_names(self, feature_names) -> List:
+    def _add_new_feature_names(self, feature_names: List[str]) -> List[str]:
         """Creates new features names and removes original categorical variables."""
         feature_names = feature_names + self._get_new_features_name()
         feature_names = [f for f in feature_names if f not in self.variables_]

tests/test_encoding/test_similarity_encoder.py

@@ -111,12 +111,22 @@ def test_nan_behaviour_impute(df_enc_big_na):
     encoder = StringSimilarityEncoder(missing_values="impute")
     X = encoder.fit_transform(df_enc_big_na)
     assert (X.isna().sum() == 0).all(axis=None)
+    assert encoder.encoder_dict_ == {
+        "var_A": ["B", "D", "G", "A", "C", "E", "F", ""],
+        "var_B": ["A", "D", "B", "G", "C", "E", "F"],
+        "var_C": ["C", "D", "B", "G", "A", "E", "F"],
+    }
 
 
 def test_nan_behaviour_ignore(df_enc_big_na):
     encoder = StringSimilarityEncoder(missing_values="ignore")
     X = encoder.fit_transform(df_enc_big_na)
     assert (X.isna().any(1) == df_enc_big_na.isna().any(1)).all()
+    assert encoder.encoder_dict_ == {
+        "var_A": ["B", "D", "G", "A", "C", "E", "F"],
+        "var_B": ["A", "D", "B", "G", "C", "E", "F"],
+        "var_C": ["C", "D", "B", "G", "A", "E", "F"],
+    }
 
 
 def test_inverse_transform_error(df_enc_big):
@@ -213,3 +223,165 @@ def test_get_feature_names_out_na(df_enc_big_na):
     }
     assert tr.get_feature_names_out(input_features=None) == out
     assert tr.get_feature_names_out(input_features=input_features) == out
+
+
+@pytest.mark.parametrize("keywords", ["hello", 0.5, [1]])
+def test_keywords_bad_type(keywords):
+    with pytest.raises(ValueError):
+        StringSimilarityEncoder(keywords=keywords)
+
+
+@pytest.mark.parametrize("item", ["hello", 0.5, 1])
+def test_keywords_bad_items(item):
+    with pytest.raises(ValueError):
+        StringSimilarityEncoder(keywords={"var_A": item})
+
+
+@pytest.mark.parametrize("key", ["hello", 0.5, 1])
+def test_keywords_bad_keys(df_enc_big, key):
+    encoder = StringSimilarityEncoder(keywords={key: ["A"]})
+    with pytest.raises(ValueError):
+        encoder.fit(df_enc_big)
+
+
+def test_encode_partial_keywords():
+    df = pd.DataFrame(
+        {
+            "var_A": ["A"] * 5
+            + ["B"] * 11
+            + ["C"] * 4
+            + ["D"] * 9
+            + ["E"] * 2
+            + ["F"] * 2
+            + ["G"] * 7,
+            "var_B": ["A"] * 11
+            + ["B"] * 7
+            + ["C"] * 4
+            + ["D"] * 9
+            + ["E"] * 2
+            + ["F"] * 2
+            + ["G"] * 5,
+            "var_C": ["A"] * 4
+            + ["B"] * 5
+            + ["C"] * 11
+            + ["D"] * 9
+            + ["E"] * 2
+            + ["F"] * 2
+            + ["G"] * 7,
+        }
+    )
+
+    encoder = StringSimilarityEncoder(top_categories=2, keywords={"var_A": ["XYZ"]})
+    X = encoder.fit_transform(df)
+
+    # test init params
+    assert encoder.top_categories == 2
+    # test fit attr
+    transf = {
+        "var_A_XYZ": 0,
+        "var_B_A": 11,
+        "var_B_D": 9,
+        "var_C_D": 9,
+        "var_C_C": 11,
+    }
+
+    # test fit attr
+    assert encoder.variables_ == ["var_A", "var_B", "var_C"]
+    assert encoder.n_features_in_ == 3
+    assert encoder.encoder_dict_ == {
+        "var_A": ["XYZ"],
+        "var_B": ["A", "D"],
+        "var_C": ["C", "D"],
+    }
+    # test transform output
+    for col in transf.keys():
+        assert X[col].sum() == transf[col]
+    assert "var_B" not in X.columns
+    assert "var_B_F" not in X.columns
+
+
+def test_encode_complete_keywords():
+    df = pd.DataFrame(
+        {
+            "var_A": ["A"] * 5
+            + ["B"] * 11
+            + ["C"] * 4
+            + ["D"] * 9
+            + ["E"] * 2
+            + ["F"] * 2
+            + ["G"] * 7,
+            "var_B": ["A"] * 11
+            + ["B"] * 7
+            + ["C"] * 4
+            + ["D"] * 9
+            + ["E"] * 2
+            + ["F"] * 2
+            + ["G"] * 5,
+            "var_C": ["A"] * 4
+            + ["B"] * 5
+            + ["C"] * 11
+            + ["D"] * 9
+            + ["E"] * 2
+            + ["F"] * 2
+            + ["G"] * 7,
+        }
+    )
+
+    encoder = StringSimilarityEncoder(
+        keywords={"var_A": ["X"], "var_B": ["Y"], "var_C": ["Z"]}
+    )
+    X = encoder.fit_transform(df)
+
+    # test fit attr
+    transf = {
+        "var_A_X": 0,
+        "var_B_Y": 0,
+        "var_C_Z": 0,
+    }
+
+    # test fit attr
+    assert encoder.variables_ == ["var_A", "var_B", "var_C"]
+    assert encoder.n_features_in_ == 3
+    assert encoder.encoder_dict_ == {
+        "var_A": ["X"],
+        "var_B": ["Y"],
+        "var_C": ["Z"],
+    }
+    # test transform output
+    for col in transf.keys():
+        assert X[col].sum() == transf[col]
+    assert "var_B" not in X.columns
+    assert "var_B_F" not in X.columns
+
+
+def test_get_feature_names_out_w_keywords(df_enc_big_na):
+    input_features = df_enc_big_na.columns.tolist()
+
+    tr = StringSimilarityEncoder(keywords={"var_A": ["XYZ"]})
+    tr.fit(df_enc_big_na)
+
+    out = [
+        "var_A_XYZ",
+        "var_B_A",
+        "var_B_D",
+        "var_B_B",
+        "var_B_G",
+        "var_B_C",
+        "var_B_E",
+        "var_B_F",
+        "var_C_C",
+        "var_C_D",
+        "var_C_B",
+        "var_C_G",
+        "var_C_A",
+        "var_C_E",
+        "var_C_F",
+    ]
+
+    assert tr.encoder_dict_ == {
+        "var_A": ["XYZ"],
+        "var_B": ["A", "D", "B", "G", "C", "E", "F"],
+        "var_C": ["C", "D", "B", "G", "A", "E", "F"],
+    }
+    assert tr.get_feature_names_out(input_features=None) == out
+    assert tr.get_feature_names_out(input_features=input_features) == out