Adds drop_original option to CombineWithReferenceFeatures, solves issue #331 (#350)

* adds option to drop original vars to CombineWithReferenceFeature transf

* changes wording of docstring

* makes an error msg more informative

* replaces pd.unique with set

* adds drop_original option to MathematicalCombination transf

Author: dodoarg

feature_engine/creation/combine_with_reference_feature.py

@@ -65,6 +65,9 @@ class CombineWithReferenceFeature(BaseEstimator, TransformerMixin):
         transformer will return an error if the training or the datasets to transform
         contain missing values.
 
+    drop_original: bool, default=False
+        If True, the original variables will be dropped from the dataframe
+        after their combination.
 
     Attributes
     ----------
@@ -98,6 +101,7 @@ def __init__(
         operations: List[str] = ["sub"],
         new_variables_names: Optional[List[str]] = None,
         missing_values: str = "ignore",
+        drop_original: bool = False,
     ) -> None:
 
         # check input types
@@ -158,11 +162,18 @@ def __init__(
         if missing_values not in ["raise", "ignore"]:
             raise ValueError("missing_values takes only values 'raise' or 'ignore'")
 
+        if not isinstance(drop_original, bool):
+            raise TypeError(
+                "drop_original takes only boolean values True and False. "
+                f"Got {drop_original} instead."
+            )
+
         self.reference_variables = reference_variables
         self.variables_to_combine = variables_to_combine
         self.new_variables_names = new_variables_names
         self.operations = operations
         self.missing_values = missing_values
+        self.drop_original = drop_original
 
     def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
         """
@@ -286,6 +297,12 @@ def transform(self, X: pd.DataFrame) -> pd.DataFrame:
         if self.new_variables_names:
             X.columns = original_col_names + self.new_variables_names
 
+        if self.drop_original:
+            X.drop(
+                columns=set(self.variables_to_combine + self.reference_variables),
+                inplace=True,
+            )
+
         return X
 
     def _more_tags(self):

feature_engine/creation/mathematical_combination.py

@@ -103,6 +103,7 @@ def __init__(
         math_operations: Optional[List[str]] = None,
         new_variables_names: Optional[List[str]] = None,
         missing_values: str = "raise",
+        drop_original: bool = False,
     ) -> None:
 
         # check input types
@@ -158,10 +159,17 @@ def __init__(
                     "combine."
                 )
 
+        if not isinstance(drop_original, bool):
+            raise TypeError(
+                "drop_original takes only boolean values True and False. "
+                f"Got {drop_original} instead."
+            )
+
         self.variables_to_combine = variables_to_combine
         self.new_variables_names = new_variables_names
         self.math_operations = math_operations
         self.missing_values = missing_values
+        self.drop_original = drop_original
 
     def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
         """
@@ -251,6 +259,9 @@ def transform(self, X: pd.DataFrame) -> pd.DataFrame:
         for new_variable_name, operation in self.combination_dict_.items():
             X[new_variable_name] = X[self.variables_to_combine].agg(operation, axis=1)
 
+        if self.drop_original:
+            X.drop(columns=self.variables_to_combine, inplace=True)
+
         return X
 
     def _more_tags(self):

tests/test_creation/test_combine_with_reference_feature.py

@@ -70,6 +70,15 @@ def test_error_if_new_variable_names_of_wrong_type():
         )
 
 
+def test_error_when_drop_original_not_bool():
+    with pytest.raises(TypeError):
+        CombineWithReferenceFeature(
+            variables_to_combine=["Age"],
+            reference_variables=["Marks"],
+            drop_original="not_a_bool"
+        )
+
+
 def test_error_when_variables_to_combine_not_numeric(df_vartypes):
     transformer = CombineWithReferenceFeature(
         variables_to_combine=["Name", "Age", "Marks"],
@@ -191,3 +200,27 @@ def test_user_enters_output_variable_names(df_vartypes):
 
     # transform params
     pd.testing.assert_frame_equal(X, ref)
+
+
+def test_drop_original_variables(df_vartypes):
+    transformer = CombineWithReferenceFeature(
+        variables_to_combine=["Age", "Marks"],
+        reference_variables=["Age", "Marks"],
+        drop_original=True
+    )
+
+    X = transformer.fit_transform(df_vartypes)
+
+    ref = pd.DataFrame.from_dict(
+        {
+            "Name": ["tom", "nick", "krish", "jack"],
+            "City": ["London", "Manchester", "Liverpool", "Bristol"],
+            "dob": pd.date_range("2020-02-24", periods=4, freq="T"),
+            "Age_sub_Age": [0, 0, 0, 0],
+            "Marks_sub_Age": [-19.1, -20.2, -18.3, -17.4],
+            "Age_sub_Marks": [19.1, 20.2, 18.3, 17.4],
+            "Marks_sub_Marks": [0.0, 0.0, 0.0, 0.0],
+        }
+    )
+
+    pd.testing.assert_frame_equal(X, ref)

tests/test_creation/test_mathematical_combination.py

@@ -119,6 +119,13 @@ def test_default_parameters(df_vartypes):
     pd.testing.assert_frame_equal(X, ref)
 
 
+def test_error_when_drop_original_not_bool():
+    with pytest.raises(TypeError):
+        MathematicalCombination(
+            variables_to_combine=["Age", "Marks"], drop_original="not_a_bool"
+        )
+
+
 def test_error_when_variables_to_combine_not_numeric(df_vartypes):
     transformer = MathematicalCombination(variables_to_combine=["Name", "Age", "Marks"])
     with pytest.raises(TypeError):
@@ -324,3 +331,25 @@ def test_no_error_when_null_values_in_variable(df_vartypes):
     )
     # transform params
     pd.testing.assert_frame_equal(X, ref)
+
+
+def test_drop_original_variables(df_vartypes):
+    transformer = MathematicalCombination(
+        variables_to_combine=["Age", "Marks"],
+        math_operations=["sum", "mean"],
+        drop_original=True,
+    )
+
+    X = transformer.fit_transform(df_vartypes)
+
+    ref = pd.DataFrame.from_dict(
+        {
+            "Name": ["tom", "nick", "krish", "jack"],
+            "City": ["London", "Manchester", "Liverpool", "Bristol"],
+            "dob": pd.date_range("2020-02-24", periods=4, freq="T"),
+            "sum(Age-Marks)": [20.9, 21.8, 19.7, 18.6],
+            "mean(Age-Marks)": [10.45, 10.9, 9.85, 9.3],
+        }
+    )
+
+    pd.testing.assert_frame_equal(X, ref)