diff --git a/docs/user_guide/transformation/YeoJohnsonTransformer.rst b/docs/user_guide/transformation/YeoJohnsonTransformer.rst index 160e6f795..7c6257356 100644 --- a/docs/user_guide/transformation/YeoJohnsonTransformer.rst +++ b/docs/user_guide/transformation/YeoJohnsonTransformer.rst @@ -201,6 +201,43 @@ values, using the `inverse_transform` method. test_unt = tf.inverse_transform(test_t) +With polars +----------- + +:class:`YeoJohnsonTransformer()` works in the same way with a polars dataframe: + +.. code:: python + + import polars as pl + from feature_engine.transformation import YeoJohnsonTransformer + + df = pl.DataFrame({ + "var_1": [-4.0, -1.0, 0.0, 3.0, 10.0], + "var_2": [1.0, 8.0, 27.0, 64.0, 125.0], + }) + + tf = YeoJohnsonTransformer(variables=None) + tf.fit(df) + Xt = tf.transform(df) + + print(Xt) + +.. code:: text + + shape: (5, 2) + ┌───────────┬──────────┐ + │ var_1 ┆ var_2 │ + │ --- ┆ --- │ + │ f64 ┆ f64 │ + ╞═══════════╪══════════╡ + │ -5.520511 ┆ 0.740576 │ + │ -1.129157 ┆ 2.723463 │ + │ 0.0 ┆ 4.64057 │ + │ 2.323426 ┆ 6.353727 │ + │ 6.13494 ┆ 7.905058 │ + └───────────┴──────────┘ + + Additional resources -------------------- diff --git a/feature_engine/transformation/yeojohnson.py b/feature_engine/transformation/yeojohnson.py index 82fa53dac..54da3001b 100644 --- a/feature_engine/transformation/yeojohnson.py +++ b/feature_engine/transformation/yeojohnson.py @@ -3,9 +3,10 @@ from typing import List, Optional, Union +import narwhals as nw import numpy as np -import pandas as pd import scipy.stats as stats +from narwhals.typing import IntoDataFrame, IntoSeries from feature_engine._base_transformers.base_numerical import BaseNumericalTransformer from feature_engine._check_init_parameters.check_init_input_params import ( @@ -108,11 +109,35 @@ class YeoJohnsonTransformer(BaseNumericalTransformer): >>> X = yjt.transform(X) >>> X.head() x - 0 -267042.906453 - 1 -444357.138990 - 2 -221626.115742 - 3 -23647.632651 - 4 -467264.993249 + 0 -267042.661354 + 1 -444356.715596 + 2 -221625.915167 + 3 -23647.614887 + 4 -467264.546413 + + With polars: + + >>> import numpy as np + >>> import polars as pl + >>> from feature_engine.transformation import YeoJohnsonTransformer + >>> np.random.seed(42) + >>> X = pl.DataFrame({"x": list(np.random.lognormal(size=6) - 10)}) + >>> yjt = YeoJohnsonTransformer() + >>> yjt.fit(X) + >>> yjt.transform(X) + shape: (6, 1) + ┌────────────────┐ + │ x │ + │ --- │ + │ f64 │ + ╞════════════════╡ + │ -467714.164249 │ + │ -795057.401919 │ + │ -385148.281012 │ + │ -37417.353351 │ + │ -837807.71099 │ + │ -837800.580457 │ + └────────────────┘ """ def __init__( @@ -125,27 +150,31 @@ def __init__( self.variables = _check_variables_input_value(variables) self.return_empty = return_empty - def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): + def fit(self, X: IntoDataFrame, y: Optional[IntoSeries] = None): """ Learn the optimal lambda for the Yeo-Johnson transformation. Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features] + X: dataframe of shape = [n_samples, n_features] The training input samples. Can be the entire dataframe, not just the variables to transform. - y: pandas Series, default=None + y: Series, default=None It is not needed in this transformer. You can pass y or None. """ # check input dataframe X, variables_ = self._fit_setup(X) - lambda_dict_ = {} + values = nw.from_native(X, eager_only=True).select(variables_).to_numpy() + values = values.astype(float) - for var in variables_: - _, lambda_dict_[var] = stats.yeojohnson(X[var]) + # scipy searches the optimal lambda one column at a time, there is no + # vectorized multi-column form of the search. + lambda_dict_ = {} + for i, var in enumerate(variables_): + _, lambda_dict_[var] = stats.yeojohnson(values[:, i]) self.variables_ = variables_ self.lambda_dict_ = lambda_dict_ @@ -153,55 +182,77 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): return self - def transform(self, X: pd.DataFrame) -> pd.DataFrame: + def transform(self, X: IntoDataFrame) -> IntoDataFrame: """ Apply the Yeo-Johnson transformation. Parameters ---------- - X: pandas DataFrame of shape = [n_samples, n_features] + X: dataframe of shape = [n_samples, n_features] The data to be transformed. Returns ------- - X: pandas dataframe + X_new: dataframe The dataframe with the transformed variables. """ # check input dataframe and if class was fitted - X = self._check_transform_input_and_state(X) - for feature in self.variables_: - X[feature] = stats.yeojohnson(X[feature], lmbda=self.lambda_dict_[feature]) + + nw_X = nw.from_native(X, eager_only=True) + values = nw_X.select(self.variables_).to_numpy().astype(float) + + # transform + result = np.empty_like(values) + for i, var in enumerate(self.variables_): + result[:, i] = stats.yeojohnson(values[:, i], lmbda=self.lambda_dict_[var]) + + new_series = [ + nw.new_series(var, result[:, i], backend=nw_X.implementation) + for i, var in enumerate(self.variables_) + ] + X = nw_X.with_columns(*new_series).to_native() return X - def inverse_transform(self, X: pd.DataFrame) -> pd.DataFrame: + def inverse_transform(self, X: IntoDataFrame) -> IntoDataFrame: """ Convert the data back to the original representation. Parameters ---------- - X: pandas DataFrame of shape = [n_samples, n_features] + X: dataframe of shape = [n_samples, n_features] The data to be transformed. Returns ------- - X_tr: pandas dataframe + X_tr: dataframe The dataframe with the transformed variables. """ # check input dataframe and if class was fitted X = self._check_transform_input_and_state(X) - for feature in self.variables_: - X[feature] = self._inverse_transform_series( - X[feature], lmbda=self.lambda_dict_[feature] + nw_X = nw.from_native(X, eager_only=True) + values = nw_X.select(self.variables_).to_numpy().astype(float) + + # inverse_transform + result = np.empty_like(values) + for i, var in enumerate(self.variables_): + result[:, i] = self._inverse_transform_array( + values[:, i], lmbda=self.lambda_dict_[var] ) + new_series = [ + nw.new_series(var, result[:, i], backend=nw_X.implementation) + for i, var in enumerate(self.variables_) + ] + X = nw_X.with_columns(*new_series).to_native() + return X - def _inverse_transform_series(self, X: pd.Series, lmbda: float) -> pd.Series: - x_inv = pd.Series(np.zeros_like(X), index=X.index) + def _inverse_transform_array(self, X: np.ndarray, lmbda: float) -> np.ndarray: + x_inv = np.zeros_like(X) pos = X >= 0 # when x >= 0 diff --git a/tests/test_transformation/test_yeojohnson_transformer.py b/tests/test_transformation/test_yeojohnson_transformer.py index f4eb32f93..b411bddcf 100644 --- a/tests/test_transformation/test_yeojohnson_transformer.py +++ b/tests/test_transformation/test_yeojohnson_transformer.py @@ -1,170 +1,194 @@ +import narwhals as nw import numpy as np import pandas as pd +import polars as pl import pytest from sklearn.exceptions import NotFittedError from feature_engine.transformation import YeoJohnsonTransformer - -def test_automatically_select_variables(df_vartypes): - # test case 1: automatically select variables +DATA = { + "Name": ["tom", "nick", "krish", "jack"], + "City": ["London", "Manchester", "Liverpool", "Bristol"], + "Age": [20, 21, 19, 18], + "Marks": [0.9, 0.8, 0.7, 0.6], +} +DATA_NA = { + "Name": ["tom", "nick", "krish", "jack"], + "City": ["London", "Manchester", "Liverpool", "Bristol"], + "Age": [20.0, 21.0, 19.0, np.nan], + "Marks": [0.9, 0.8, 0.7, np.nan], +} + + +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_automatically_select_variables_and_inverse_transform(make_df): + X = make_df(DATA) transformer = YeoJohnsonTransformer(variables=None) - X = transformer.fit_transform(df_vartypes) - - # expected result - transf_df = df_vartypes.copy() - transf_df["Age"] = [10.167, 10.5406, 9.78774, 9.40229] - transf_df["Marks"] = [0.804449, 0.722367, 0.638807, 0.553652] + Xt = transformer.fit_transform(X) # test init params assert transformer.variables is None - # test fit attr + # test fit attrs assert transformer.variables_ == ["Age", "Marks"] - assert transformer.n_features_in_ == 5 + assert transformer.n_features_in_ == 4 + # test transform output - pd.testing.assert_frame_equal(X, transf_df) + result = nw.from_native(Xt, eager_only=True).to_dict(as_series=False) + assert result["Age"] == pytest.approx( + [10.167048, 10.540602, 9.787738, 9.402289], abs=1e-5 + ) + assert result["Marks"] == pytest.approx( + [0.804449, 0.722367, 0.638807, 0.553652], abs=1e-5 + ) + + # test inverse_transform, including non-transformed columns + Xit = transformer.inverse_transform(Xt) + result_it = nw.from_native(Xit, eager_only=True).to_dict(as_series=False) + assert [round(v) for v in result_it["Age"]] == DATA["Age"] + assert [round(v, 1) for v in result_it["Marks"]] == DATA["Marks"] + assert result_it["Name"] == DATA["Name"] + assert result_it["City"] == DATA["City"] -def test_transformer_on_integer_variables(): - df = pd.DataFrame( +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_transformer_on_integer_variables(make_df): + X = make_df( { "var1": [0, 1, 0, 2, 3, 4, 5, 6, 8, 10], "var2": [12, 11, 10, 15, 13, 12, 11, 10, 10, 20], } ) - dft = pd.DataFrame( - { - "var1": { - 0: 0.0, - 1: 0.7871467037957388, - 2: 0.0, - 3: 1.34716625120788, - 4: 1.797027857352365, - 5: 2.1794549065159363, - 6: 2.5155129679774246, - 7: 2.817344570368886, - 8: 3.346739213848269, - 9: 3.8051709334268566, - }, - "var2": { - 0: 0.2891005444159968, - 1: 0.2890875957028113, - 2: 0.2890687942494933, - 3: 0.2891213447054929, - 4: 0.2891097235906253, - 5: 0.2891005444159968, - 6: 0.2890875957028113, - 7: 0.2890687942494933, - 8: 0.2890687942494933, - 9: 0.28913341330818815, - }, - } + Xt = YeoJohnsonTransformer().fit_transform(X) + result = nw.from_native(Xt, eager_only=True).to_dict(as_series=False) + + assert result["var1"] == pytest.approx( + [ + 0.0, + 0.787147, + 0.0, + 1.347166, + 1.797028, + 2.179455, + 2.515513, + 2.817345, + 3.346739, + 3.805171, + ], + abs=1e-5, + ) + assert result["var2"] == pytest.approx( + [ + 0.289101, + 0.289088, + 0.289069, + 0.289121, + 0.289110, + 0.289101, + 0.289088, + 0.289069, + 0.289069, + 0.289133, + ], + abs=1e-5, ) - - X_tr = YeoJohnsonTransformer().fit_transform(df) - pd.testing.assert_frame_equal(X_tr, dft) -def test_fit_raises_error_if_na_in_df(df_na): - # test case 2: when dataset contains na, fit method +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_fit_raises_error_if_na_in_df(make_df): + X = make_df(DATA_NA) with pytest.raises(ValueError): transformer = YeoJohnsonTransformer() - transformer.fit(df_na) + transformer.fit(X) -def test_transform_raises_error_if_na_in_df(df_vartypes, df_na): - # test case 3: when dataset contains na, transform method +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_transform_raises_error_if_na_in_df(make_df): + X = make_df(DATA) + X_na = make_df(DATA_NA) + transformer = YeoJohnsonTransformer() + transformer.fit(X) with pytest.raises(ValueError): - transformer = YeoJohnsonTransformer() - transformer.fit(df_vartypes) - transformer.transform(df_na[["Name", "City", "Age", "Marks", "dob"]]) + transformer.transform(X_na) -def test_non_fitted_error(df_vartypes): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_non_fitted_error(make_df): + X = make_df(DATA) with pytest.raises(NotFittedError): transformer = YeoJohnsonTransformer() - transformer.transform(df_vartypes) - - -def test_inverse_transform_automatically_select_only_transformed_columns(df_vartypes): - X = df_vartypes.copy(deep=True) - transformer = YeoJohnsonTransformer(variables=None) - X_trans = transformer.fit_transform(X) + transformer.transform(X) - X_inverse = transformer.inverse_transform(X_trans) - X_inverse["Age"] = X_inverse["Age"].round(0).astype(int) - pd.testing.assert_frame_equal(X, X_inverse, check_dtype=False) - - -def test_inverse_with_X_negative_and_positive(): - X = pd.DataFrame( +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_inverse_with_x_negative_and_positive(make_df): + X = make_df( { - "var1": np.arange(-20, 0), - "var2": np.arange(0, 20), - "var3": np.arange(-10, 10), + "var1": list(np.arange(-20, 0)), + "var2": list(np.arange(0, 20)), + "var3": list(np.arange(-10, 10)), } ) transformer = YeoJohnsonTransformer(variables=None) - X_trans = transformer.fit_transform(X) - - X_inverse = transformer.inverse_transform(X_trans) - X_inverse = X_inverse.round(0).astype(int) + Xt = transformer.fit_transform(X) + Xi = transformer.inverse_transform(Xt) + result = nw.from_native(Xi, eager_only=True).to_dict(as_series=False) - pd.testing.assert_frame_equal(X, X_inverse, check_dtype=False) + assert [round(v) for v in result["var1"]] == list(np.arange(-20, 0)) + assert [round(v) for v in result["var2"]] == list(np.arange(0, 20)) + assert [round(v) for v in result["var3"]] == list(np.arange(-10, 10)) -def test_inverse_with_with_non_linear_index(): +def test_inverse_with_non_linear_index(): + # pandas-specific: exercises index-preserving behaviour, which has no + # polars equivalent (polars has no row index). X = pd.DataFrame( { "var1": np.arange(-20, 0), "var2": np.arange(0, 20), "var3": np.arange(-10, 10), }, - index=[13, 15, 12, 11, 17, 9, 4, 0, 1, 14, 18, 2, 3, 6, 5, 7, 8, 2, 16, 10] + index=[13, 15, 12, 11, 17, 9, 4, 0, 1, 14, 18, 2, 3, 6, 5, 7, 8, 2, 16, 10], ) transformer = YeoJohnsonTransformer(variables=None) - X_trans = transformer.fit_transform(X) + Xt = transformer.fit_transform(X) - X_inverse = transformer.inverse_transform(X_trans) - X_inverse = X_inverse.round(0).astype(int) + Xi = transformer.inverse_transform(Xt) + Xi = Xi.round(0).astype(int) - pd.testing.assert_frame_equal(X, X_inverse, check_dtype=False) + pd.testing.assert_frame_equal(X, Xi, check_dtype=False) -def test_lambda_equals_lambda_equal_0(): - X = pd.DataFrame( - { - "var1": np.arange(0, 20), - "var2": np.arange(20, 40), - } - ) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_lambda_equal_0(make_df): + X = make_df({"var1": list(np.arange(0, 20)), "var2": list(np.arange(20, 40))}) transformer = YeoJohnsonTransformer(variables=None) transformer = transformer.fit(X) - transformer.lambda_dict_ = {"var1": 0, "var2": 0} - X_trans = transformer.transform(X) - X_inverse = transformer.inverse_transform(X_trans) - X_inverse = X_inverse.round(0).astype(int) + Xt = transformer.transform(X) + Xi = transformer.inverse_transform(Xt) + result = nw.from_native(Xi, eager_only=True).to_dict(as_series=False) - pd.testing.assert_frame_equal(X, X_inverse, check_dtype=False) + assert [round(v) for v in result["var1"]] == list(np.arange(0, 20)) + assert [round(v) for v in result["var2"]] == list(np.arange(20, 40)) -def test_lambda_equals_lambda_equal_2(): - X = pd.DataFrame({"var1": np.arange(-21, -1), "var2": np.arange(-41, -21)}) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_lambda_equal_2(make_df): + X = make_df({"var1": list(np.arange(-21, -1)), "var2": list(np.arange(-41, -21))}) transformer = YeoJohnsonTransformer(variables=None) transformer = transformer.fit(X) - transformer.lambda_dict_ = {"var1": 2, "var2": 2} - X_trans = transformer.transform(X) - X_inverse = transformer.inverse_transform(X_trans) - X_inverse = X_inverse.round(0).astype(int) + Xt = transformer.transform(X) + Xi = transformer.inverse_transform(Xt) + result = nw.from_native(Xi, eager_only=True).to_dict(as_series=False) - pd.testing.assert_frame_equal(X, X_inverse, check_dtype=False) + assert [round(v) for v in result["var1"]] == list(np.arange(-21, -1)) + assert [round(v) for v in result["var2"]] == list(np.arange(-41, -21))