diff --git a/feature_engine/_docstrings/init_parameters/imputers.py b/feature_engine/_docstrings/init_parameters/imputers.py new file mode 100644 index 000000000..a5df24884 --- /dev/null +++ b/feature_engine/_docstrings/init_parameters/imputers.py @@ -0,0 +1,4 @@ +_missing_only_docstring = """missing_only: bool, default=False + If `True` and `variables` is `None`, the transformer will select and impute + only the variables that show missing values during `fit()`. + """.rstrip() diff --git a/feature_engine/imputation/arbitrary_number.py b/feature_engine/imputation/arbitrary_number.py index 69d7c624b..914c5e2c2 100644 --- a/feature_engine/imputation/arbitrary_number.py +++ b/feature_engine/imputation/arbitrary_number.py @@ -28,6 +28,9 @@ from feature_engine._docstrings.init_parameters.all_transformers import ( _return_empty_docstring ) +from feature_engine._docstrings.init_parameters.imputers import ( + _missing_only_docstring, +) from feature_engine._docstrings.substitute import Substitution from feature_engine.dataframe_checks import check_X from feature_engine.imputation.base_imputer import BaseImputer @@ -41,6 +44,7 @@ imputer_dict_=_imputer_dict_docstring, variables_=_variables_attribute_docstring, return_empty=_return_empty_docstring, + missing_only=_missing_only_docstring, feature_names_in_=_feature_names_in_docstring, n_features_in_=_n_features_in_docstring, fit=_fit_not_learn_docstring, @@ -76,6 +80,8 @@ class ArbitraryNumberImputer(BaseImputer): The dictionary of variables and the arbitrary numbers for their imputation. If specified, it overrides the above parameters. + {missing_only} + Attributes ---------- @@ -126,6 +132,7 @@ def __init__( variables: Union[None, int, str, List[Union[str, int]]] = None, return_empty: bool = False, imputer_dict: Optional[dict] = None, + missing_only: bool = False, ) -> None: if isinstance(arbitrary_number, int) or isinstance(arbitrary_number, float): @@ -133,6 +140,8 @@ def __init__( else: raise ValueError("arbitrary_number must be numeric of type int or float") + super().__init__(missing_only) + _check_numerical_dict(imputer_dict) self.variables = _check_variables_input_value(variables) @@ -168,6 +177,8 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): else: if self.variables is None: self.variables_ = find_numerical_variables(X, self.return_empty) + if self.missing_only: + self.variables_ = self._filter_variables_with_na(X, self.variables_) else: self.variables_ = check_numerical_variables(X, self.variables) self.imputer_dict_ = {var: self.arbitrary_number for var in self.variables_} diff --git a/feature_engine/imputation/base_imputer.py b/feature_engine/imputation/base_imputer.py index f9c3a2fea..d20da72b7 100644 --- a/feature_engine/imputation/base_imputer.py +++ b/feature_engine/imputation/base_imputer.py @@ -12,6 +12,21 @@ class BaseImputer(TransformerMixin, BaseEstimator, GetFeatureNamesOutMixin): """shared set-up checks and methods across imputers""" + def __init__( + self, + missing_only: bool = False, + ) -> None: + if not isinstance(missing_only, bool): + raise ValueError( + f"missing_only must be a boolean. Got {missing_only} instead." + ) + + self.missing_only = missing_only + + def _filter_variables_with_na(self, X: pd.DataFrame, variables): + """Return variables that contain missing values.""" + return [var for var in variables if X[var].isnull().any()] + def _transform(self, X: pd.DataFrame) -> pd.DataFrame: """ Common checks before transforming data: diff --git a/feature_engine/imputation/categorical.py b/feature_engine/imputation/categorical.py index 1c989c437..4240f64a2 100644 --- a/feature_engine/imputation/categorical.py +++ b/feature_engine/imputation/categorical.py @@ -24,6 +24,9 @@ from feature_engine._docstrings.init_parameters.all_transformers import ( _return_empty_docstring ) +from feature_engine._docstrings.init_parameters.imputers import ( + _missing_only_docstring, +) from feature_engine._docstrings.substitute import Substitution from feature_engine.dataframe_checks import check_X from feature_engine.imputation.base_imputer import BaseImputer @@ -40,6 +43,7 @@ imputer_dict_=_imputer_dict_docstring, variables_=_variables_attribute_docstring, return_empty=_return_empty_docstring, + missing_only=_missing_only_docstring, feature_names_in_=_feature_names_in_docstring, n_features_in_=_n_features_in_docstring, transform=_transform_imputers_docstring, @@ -97,6 +101,8 @@ class CategoricalImputer(BaseImputer): type object or categorical. If True, the imputer will select all variables or accept all variables entered by the user, including those cast as numeric. + {missing_only} + Attributes ---------- {imputer_dict_} @@ -145,6 +151,7 @@ def __init__( return_empty: bool = False, return_object: bool = False, ignore_format: bool = False, + missing_only: bool = False, ) -> None: if imputation_method not in ["missing", "frequent"]: raise ValueError( @@ -154,6 +161,8 @@ def __init__( if not isinstance(ignore_format, bool): raise ValueError("ignore_format takes only booleans True and False") + super().__init__(missing_only) + self.imputation_method = imputation_method self.fill_value = fill_value self.variables = _check_variables_input_value(variables) @@ -190,6 +199,9 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): else: self.variables_ = check_categorical_variables(X, self.variables) + if self.variables is None and self.missing_only: + self.variables_ = self._filter_variables_with_na(X, self.variables_) + if self.imputation_method == "missing": self.imputer_dict_ = {var: self.fill_value for var in self.variables_} diff --git a/feature_engine/imputation/end_tail.py b/feature_engine/imputation/end_tail.py index cd895e971..ad9321aee 100644 --- a/feature_engine/imputation/end_tail.py +++ b/feature_engine/imputation/end_tail.py @@ -20,6 +20,9 @@ from feature_engine._docstrings.init_parameters.all_transformers import ( _variables_numerical_docstring, _return_empty_docstring ) +from feature_engine._docstrings.init_parameters.imputers import ( + _missing_only_docstring, +) from feature_engine._docstrings.methods import ( _fit_transform_docstring, _transform_imputers_docstring, @@ -36,6 +39,7 @@ @Substitution( variables=_variables_numerical_docstring, return_empty=_return_empty_docstring, + missing_only=_missing_only_docstring, imputer_dict_=_imputer_dict_docstring, variables_=_variables_attribute_docstring, feature_names_in_=_feature_names_in_docstring, @@ -105,6 +109,8 @@ class EndTailImputer(BaseImputer): {return_empty} + {missing_only} + Attributes ---------- {imputer_dict_} @@ -149,6 +155,7 @@ def __init__( fold: int = 3, variables: Union[None, int, str, List[Union[str, int]]] = None, return_empty: bool = False, + missing_only: bool = False, ) -> None: if imputation_method not in ["gaussian", "iqr", "max"]: @@ -162,6 +169,8 @@ def __init__( if fold <= 0: raise ValueError("fold takes only positive numbers") + super().__init__(missing_only) + self.imputation_method = imputation_method self.tail = tail self.fold = fold @@ -188,6 +197,8 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): # find or check for numerical variables if self.variables is None: self.variables_ = find_numerical_variables(X, self.return_empty) + if self.missing_only: + self.variables_ = self._filter_variables_with_na(X, self.variables_) else: self.variables_ = check_numerical_variables(X, self.variables) diff --git a/feature_engine/imputation/mean_median.py b/feature_engine/imputation/mean_median.py index 997ec2813..6d12876e5 100644 --- a/feature_engine/imputation/mean_median.py +++ b/feature_engine/imputation/mean_median.py @@ -20,6 +20,9 @@ from feature_engine._docstrings.init_parameters.all_transformers import ( _variables_numerical_docstring, _return_empty_docstring ) +from feature_engine._docstrings.init_parameters.imputers import ( + _missing_only_docstring, +) from feature_engine._docstrings.methods import ( _fit_transform_docstring, _transform_imputers_docstring, @@ -36,6 +39,7 @@ @Substitution( variables=_variables_numerical_docstring, return_empty=_return_empty_docstring, + missing_only=_missing_only_docstring, imputer_dict_=_imputer_dict_docstring, variables_=_variables_attribute_docstring, feature_names_in_=_feature_names_in_docstring, @@ -63,6 +67,8 @@ class MeanMedianImputer(BaseImputer): {return_empty} + {missing_only} + Attributes ---------- {imputer_dict_} @@ -108,11 +114,14 @@ def __init__( imputation_method: str = "median", variables: Union[None, int, str, List[Union[str, int]]] = None, return_empty: bool = False, + missing_only: bool = False, ) -> None: if imputation_method not in ["median", "mean"]: raise ValueError("imputation_method takes only values 'median' or 'mean'") + super().__init__(missing_only) + self.imputation_method = imputation_method self.variables = _check_variables_input_value(variables) @@ -138,6 +147,8 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): # find or check for numerical variables if self.variables is None: self.variables_ = find_numerical_variables(X, self.return_empty) + if self.missing_only: + self.variables_ = self._filter_variables_with_na(X, self.variables_) else: self.variables_ = check_numerical_variables(X, self.variables) diff --git a/feature_engine/imputation/random_sample.py b/feature_engine/imputation/random_sample.py index 9366c1120..568a33680 100644 --- a/feature_engine/imputation/random_sample.py +++ b/feature_engine/imputation/random_sample.py @@ -24,6 +24,9 @@ from feature_engine._docstrings.init_parameters.all_transformers import ( _return_empty_docstring ) +from feature_engine._docstrings.init_parameters.imputers import ( + _missing_only_docstring, +) from feature_engine._docstrings.substitute import Substitution from feature_engine.dataframe_checks import check_X from feature_engine.imputation.base_imputer import BaseImputer @@ -50,6 +53,7 @@ def _define_seed( @Substitution( variables_=_variables_attribute_docstring, return_empty=_return_empty_docstring, + missing_only=_missing_only_docstring, feature_names_in_=_feature_names_in_docstring, n_features_in_=_n_features_in_docstring, transform=_transform_imputers_docstring, @@ -100,6 +104,8 @@ class RandomSampleImputer(BaseImputer): observation, you can choose to combine those values as an addition or a multiplication. Can take the values 'add' or 'multiply'. + {missing_only} + Attributes ---------- X_: @@ -148,6 +154,7 @@ def __init__( random_state: Union[None, int, str, List[Union[str, int]]] = None, seed: str = "general", seeding_method: str = "add", + missing_only: bool = False, ) -> None: if seed not in ["general", "observation"]: @@ -168,6 +175,8 @@ def __init__( "or more variables which will be used to seed the imputer" ) + super().__init__(missing_only) + self.variables = _check_variables_input_value(variables) _check_return_empty_is_bool(return_empty) @@ -199,6 +208,8 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): # find variables to impute if self.variables is None: self.variables_ = find_all_variables(X, self.return_empty) + if self.missing_only: + self.variables_ = self._filter_variables_with_na(X, self.variables_) else: self.variables_ = check_all_variables(X, self.variables) diff --git a/tests/test_imputation/test_arbitrary_number_imputer.py b/tests/test_imputation/test_arbitrary_number_imputer.py index dd83dea68..a750df202 100644 --- a/tests/test_imputation/test_arbitrary_number_imputer.py +++ b/tests/test_imputation/test_arbitrary_number_imputer.py @@ -82,3 +82,32 @@ def test_dictionary_of_imputation_values(df_na): def imputer_error_when_dictionary_value_is_string(): with pytest.raises(ValueError): ArbitraryNumberImputer(imputer_dict={"Age": "arbitrary_number"}) + + +def test_missing_only_selects_numerical_variables_with_na(df_na): + imputer = ArbitraryNumberImputer(arbitrary_number=99, missing_only=True) + X = df_na.copy() + X["Var_No_Nulls"] = [333] * X.shape[0] + X_transformed = imputer.fit_transform(X) + + expected_results_df = df_na.copy() + expected_results_df["Age"] = expected_results_df["Age"].fillna(99) + expected_results_df["Marks"] = expected_results_df["Marks"].fillna(99) + expected_results_df["Var_No_Nulls"] = [333] * X.shape[0] + + assert imputer.variables_ == ["Age", "Marks"] + assert imputer.imputer_dict_ == {"Age": 99, "Marks": 99} + pd.testing.assert_frame_equal(X_transformed, expected_results_df) + + +def test_missing_only_ignored_when_imputer_dict_provided(df_na): + imputer = ArbitraryNumberImputer( + imputer_dict={"Age": -42, "Marks": -999}, + missing_only=True, + ) + X = df_na.copy() + X["Marks"] = [0.9] * X.shape[0] + imputer.fit(X) + + assert imputer.variables_ == ["Age", "Marks"] + assert imputer.imputer_dict_ == {"Age": -42, "Marks": -999} diff --git a/tests/test_imputation/test_categorical_imputer.py b/tests/test_imputation/test_categorical_imputer.py index 182e8826b..c2e1211c5 100644 --- a/tests/test_imputation/test_categorical_imputer.py +++ b/tests/test_imputation/test_categorical_imputer.py @@ -305,3 +305,39 @@ def test_error_when_ignore_format_is_not_boolean(ignore_format): # check that error message matches assert str(record.value) == msg + + +def test_missing_only_selects_categorical_variables_with_na(df_na): + imputer = CategoricalImputer( + imputation_method="missing", + missing_only=True, + ) + X = df_na.copy() + X["Cat_No_Nulls"] = ["pasta"] * X.shape[0] + X_transformed = imputer.fit_transform(X) + + expected_results_df = df_na.copy() + expected_results_df["Name"] = expected_results_df["Name"].fillna("Missing") + expected_results_df["City"] = expected_results_df["City"].fillna("Missing") + expected_results_df["Studies"] = expected_results_df["Studies"].fillna("Missing") + expected_results_df["Cat_No_Nulls"] = ["pasta"] * X.shape[0] + + assert imputer.variables_ == ["Name", "City", "Studies"] + assert "Cat_No_Nulls" not in imputer.imputer_dict_ + pd.testing.assert_frame_equal(X_transformed, expected_results_df) + + +def test_missing_only_with_ignore_format_selects_all_variables_with_na(df_na): + imputer = CategoricalImputer( + imputation_method="missing", + fill_value="Missing", + ignore_format=True, + missing_only=True, + ) + X = df_na.copy() + X["Var_No_Nulls"] = [333] * X.shape[0] + imputer.fit(X) + + assert imputer.variables_ == ["Name", "City", "Studies", "Age", "Marks"] + assert "Var_No_Nulls" not in imputer.variables_ + assert "dob" not in imputer.variables_ diff --git a/tests/test_imputation/test_end_tail_imputer.py b/tests/test_imputation/test_end_tail_imputer.py index 88998d658..9eb00b875 100644 --- a/tests/test_imputation/test_end_tail_imputer.py +++ b/tests/test_imputation/test_end_tail_imputer.py @@ -101,3 +101,24 @@ def test_error_when_tail_is_string(): def test_error_when_fold_is_1(): with pytest.raises(ValueError): EndTailImputer(fold=-1) + + +def test_missing_only_selects_numerical_variables_with_na(df_na): + imputer = EndTailImputer( + imputation_method="iqr", + tail="right", + fold=1.5, + missing_only=True, + ) + X = df_na.copy() + X["Var_No_Nulls"] = [1984] * X.shape[0] + X_transformed = imputer.fit_transform(X) + + expected_results_df = df_na.copy() + expected_results_df["Age"] = expected_results_df["Age"].fillna(65.5) + expected_results_df["Marks"] = expected_results_df["Marks"].fillna(1.0625) + expected_results_df["Var_No_Nulls"] = [1984] * X.shape[0] + + assert imputer.variables_ == ["Age", "Marks"] + assert "Var_No_Nulls" not in imputer.imputer_dict_ + pd.testing.assert_frame_equal(X_transformed, expected_results_df) diff --git a/tests/test_imputation/test_mean_median_imputer.py b/tests/test_imputation/test_mean_median_imputer.py index b065c2190..f95591790 100644 --- a/tests/test_imputation/test_mean_median_imputer.py +++ b/tests/test_imputation/test_mean_median_imputer.py @@ -62,3 +62,61 @@ def test_median_imputation_when_user_enters_single_variables(df_na): def test_error_with_wrong_imputation_method(): with pytest.raises(ValueError): MeanMedianImputer(imputation_method="arbitrary") + + +def test_missing_only_selects_numerical_variables_with_na(df_na): + imputer = MeanMedianImputer( + imputation_method="mean", + missing_only=True, + ) + X = df_na.copy() + X["Var_No_Nulls"] = [333] * X.shape[0] + X_transformed = imputer.fit_transform(X) + + expected_results_df = df_na.copy() + expected_results_df["Age"] = expected_results_df["Age"].fillna(28.714285714285715) + expected_results_df["Marks"] = expected_results_df["Marks"].fillna( + 0.6833333333333332 + ) + expected_results_df["Var_No_Nulls"] = [333] * X.shape[0] + + assert imputer.variables_ == ["Age", "Marks"] + assert "Var_No_Nulls" not in imputer.imputer_dict_ + pd.testing.assert_frame_equal(X_transformed, expected_results_df) + + +def test_missing_only_with_median(df_na): + imputer = MeanMedianImputer( + imputation_method="median", + missing_only=True, + ) + X = df_na.copy() + X["Var_No_Nulls"] = [333] * X.shape[0] + X_transformed = imputer.fit_transform(X) + + expected_results_df = df_na.copy() + expected_results_df["Age"] = expected_results_df["Age"].fillna(23.0) + expected_results_df["Marks"] = expected_results_df["Marks"].fillna(0.75) + expected_results_df["Var_No_Nulls"] = [333] * X.shape[0] + + assert imputer.variables_ == ["Age", "Marks"] + pd.testing.assert_frame_equal(X_transformed, expected_results_df) + + +def test_missing_only_ignored_when_variables_provided(df_na): + imputer = MeanMedianImputer( + imputation_method="median", + variables=["Age", "Marks"], + missing_only=True, + ) + X = df_na.copy() + X["Marks"] = [0.9] * X.shape[0] + imputer.fit(X) + + assert imputer.variables_ == ["Age", "Marks"] + assert set(imputer.imputer_dict_.keys()) == {"Age", "Marks"} + + +def test_error_when_missing_only_not_bool(): + with pytest.raises(ValueError): + MeanMedianImputer(missing_only="yes") diff --git a/tests/test_imputation/test_random_sample_imputer.py b/tests/test_imputation/test_random_sample_imputer.py index cd296b7c8..bff74c6be 100644 --- a/tests/test_imputation/test_random_sample_imputer.py +++ b/tests/test_imputation/test_random_sample_imputer.py @@ -307,3 +307,24 @@ def test_variables_cast_as_category(df_na): # test transform output pd.testing.assert_frame_equal(X_transformed, ref, check_dtype=False) + + +def test_missing_only_selects_variables_with_na(df_na): + imputer = RandomSampleImputer( + variables=None, + random_state=47, + seed="general", + seeding_method="add", + missing_only=True, + ) + X = df_na.copy() + X["Var_No_Nulls"] = ["pasta"] * X.shape[0] + imputer.fit(X) + + assert imputer.variables_ == ["Name", "City", "Studies", "Age", "Marks"] + assert "Var_No_Nulls" not in imputer.variables_ + assert "dob" not in imputer.variables_ + + X_transformed = imputer.transform(X) + assert X_transformed[imputer.variables_].isnull().sum().sum() == 0 + assert X_transformed["Var_No_Nulls"].tolist() == ["pasta"] * X.shape[0]