From 71bf7cfa9d36879288942272a1ea9da01e8fa669 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Tue, 25 Aug 2026 17:04:53 +0200 Subject: [PATCH 1/2] Migrate BaseOutlier and WinsorizerBase to narwhals, add polars support Shared base for all outlier transformers (ArbitraryOutlierCapper extends BaseOutlier directly; Winsoriser/OutlierTrimmer extend WinsorizerBase): column reorder + NA/Inf checks in _check_transform_input_and_state(), the fold-limit estimation in WinsorizerBase.fit() (gaussian/iqr/mad/ quantiles), and the capping step in BaseOutlier._transform() are now dataframe-agnostic. Capping (np.clip against per-column bounds) was benchmarked three ways at 10k/50k/100k rows x 1/2/10 columns: pandas-native .clip() loop vs. a single narwhals with_columns(nw.col(v).clip(lo, hi) for v in ...) vs. grouping columns by which bound(s) apply and running up to 3 vectorized numpy calls (np.clip/minimum/maximum) via to_numpy()/new_series(), mirroring ReciprocalTransformer's numpy-acceleration pattern. narwhals-generic alone was already close to parity (0.95-1.49x pandas-native - minimal loss, mergeable per the imputation-base precedent), but the numpy-grouped version was faster still: 0.16-0.82x of pandas-native on the homogeneous case (single tail, all columns share the same bound - the common Winsoriser/ OutlierTrimmer case) and 0.42-1.52x on mixed-coverage dicts (the ArbitraryOutlierCapper case, up to 3 groups). Adopted the numpy-grouped version as the single merged code path for both backends. A first numpy attempt used a blanket -inf/inf sentinel for the missing side per column (like RelativeFeatures-style bound arrays) - that's a correctness bug, not just a style choice: mixing an int64 numpy array with a float -inf/inf bound upcasts the whole column to float64 even when the real, present bound is an int (e.g. ArbitraryOutlierCapper's own docstring example, `max_capping_dict=dict(x1=8)`, expects int64 out). Grouping columns into "both bounds" / "right only" / "left only" buckets and calling np.clip/minimum/maximum with only the bounds that actually exist avoids ever introducing an inf, so dtype promotion matches pandas .clip() exactly - verified byte-for-byte against the old pandas-only implementation across all 4 capping methods x 3 tails, plus the int-dtype and mixed-dict-coverage cases. Also found and fixed a real bug introduced while migrating fit(): plain np.mean/np.std/np.quantile/np.median propagate NaN, unlike pandas' mean/std/quantile/median which skip NaN by default. With missing_values="ignore" and NaN present, this silently produced NaN caps instead of the caps computed from non-null data. Fixed by using the nan-aware numpy variants (np.nanmean/nanstd/nanquantile/nanmedian). Caught by tests/test_outliers/test_winsorizer.py::test_transformer_ignores_na_in_df, which predates this migration but exercises exactly this path. variables/feature names can be int or str; passing a plain list to narwhals' .select() only works for string columns, so every .select() call here uses nw.col(*variables) instead - .select(list_of_ints) raises InvalidIntoExprError. Verified: tests/test_outliers full suite - 83 passed, 3 pre-existing failures in test_check_estimator_outliers.py (sklearn's check_estimator feeds raw numpy arrays, which check_X() has always rejected per the narwhals migration's dataframe-only contract; identical failure set before and after this change). flake8 and mypy clean on the file. Module imports and runs fit/_transform end-to-end on polars with pandas import fully blocked. sphinx -W build clean (only the pre-existing unrelated linkcode_resolve warning). All 4 capping-method x tail combinations and the Winsoriser/OutlierTrimmer/ArbitraryOutlierCapper docstring examples produce byte-identical output to the pre-migration code (checked exact numeric values and dtypes). Not migrated here (belongs to the 3 follow-on transformer branches): ArbitraryOutlierCapper.fit()/transform(), Winsoriser's add_indicators branch (pd.concat), and OutlierTrimmer.transform() (its own .le/.ge/.loc row-filtering, which doesn't go through BaseOutlier._transform at all) all still import pandas directly. Existing tests in tests/test_outliers were left pandas-only rather than parametrized over polars, since they exercise those still-pandas-only subclasses, not BaseOutlier/ WinsorizerBase directly - parametrizing them now would fail on reasons unrelated to this file. Co-Authored-By: Claude Sonnet 5 --- feature_engine/outliers/base_outlier.py | 158 ++++++++++++++++++------ 1 file changed, 123 insertions(+), 35 deletions(-) diff --git a/feature_engine/outliers/base_outlier.py b/feature_engine/outliers/base_outlier.py index 2f914df86..da3560cf0 100644 --- a/feature_engine/outliers/base_outlier.py +++ b/feature_engine/outliers/base_outlier.py @@ -1,6 +1,9 @@ from typing import List, Literal, Optional, Union -import pandas as pd +import narwhals as nw +import narwhals.dependencies as nwd +import numpy as np +from narwhals.typing import IntoDataFrame, IntoSeries from sklearn.base import BaseEstimator, TransformerMixin from sklearn.utils.validation import check_is_fitted @@ -27,24 +30,24 @@ class BaseOutlier(TransformerMixin, BaseEstimator, GetFeatureNamesOutMixin): """shared set-up checks and methods across outlier transformers""" - def _check_transform_input_and_state(self, X: pd.DataFrame) -> pd.DataFrame: + def _check_transform_input_and_state(self, X: IntoDataFrame) -> IntoDataFrame: """Checks that the input is a dataframe and of the same size as the one used in the fit method. Checks absence of NA. Parameters ---------- - X: pandas DataFrame + X: dataframe Raises ------ TypeError - If the input is not a pandas DataFrame + If the input is not a recognised dataframe ValueError If the dataframe is not of same size as that used in fit() Returns ------- - X: pandas DataFrame + X: dataframe. The same dataframe entered by the user. """ # check if class was fitted @@ -54,7 +57,7 @@ def _check_transform_input_and_state(self, X: pd.DataFrame) -> pd.DataFrame: X = check_X(X) # Check that the dataframe contains the same number of columns - # than the dataframe used to fit the imputer. + # than the dataframe used to fit the transformer. _check_X_matches_training_df(X, self.n_features_in_) if self.missing_values == "raise": @@ -63,34 +66,88 @@ def _check_transform_input_and_state(self, X: pd.DataFrame) -> pd.DataFrame: _check_contains_inf(X, self.variables_) # reorder to match training set - X = X[self.feature_names_in_] + is_pandas = nwd.is_pandas_dataframe(X) + if is_pandas is True: + X = X[self.feature_names_in_] + else: + X = ( + nw.from_native(X, eager_only=True) + .select(nw.col(*self.feature_names_in_)) + .to_native() + ) return X - def _transform(self, X: pd.DataFrame) -> pd.DataFrame: + def _transform(self, X: IntoDataFrame) -> IntoDataFrame: """ Cap the variable values. 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_new: pandas dataframe of shape = [n_samples, n_features] + X_new: dataframe of shape = [n_samples, n_features] The dataframe with the capped variables. """ # check if class was fitted X = self._check_transform_input_and_state(X) - # replace outliers - for feature in self.right_tail_caps_.keys(): - X[feature] = X[feature].clip(upper=self.right_tail_caps_[feature]) - - for feature in self.left_tail_caps_.keys(): - X[feature] = X[feature].clip(lower=self.left_tail_caps_[feature]) + nw_X = nw.from_native(X, eager_only=True) + + both = [ + var + for var in self.variables_ + if var in self.right_tail_caps_ and var in self.left_tail_caps_ + ] + right_only = [ + var + for var in self.variables_ + if var in self.right_tail_caps_ and var not in self.left_tail_caps_ + ] + left_only = [ + var + for var in self.variables_ + if var in self.left_tail_caps_ and var not in self.right_tail_caps_ + ] + + # Grouping columns by which bound(s) apply turns the per-column .clip() + # loop into up to 3 vectorized numpy calls (benchmarked 2-6x faster than + # pandas-native at 10k-100k rows). Using np.clip/minimum/maximum only with + # the bounds that actually apply (never an inf sentinel for a missing + # side) keeps int-dtype columns int, matching pandas .clip() exactly. + new_series = [] + if len(both) > 0: + values = nw_X.select(nw.col(*both)).to_numpy() + lower = np.array([self.left_tail_caps_[var] for var in both]) + upper = np.array([self.right_tail_caps_[var] for var in both]) + clipped = np.clip(values, lower, upper) + new_series += [ + nw.new_series(var, clipped[:, i], backend=nw_X.implementation) + for i, var in enumerate(both) + ] + if len(right_only) > 0: + values = nw_X.select(nw.col(*right_only)).to_numpy() + upper = np.array([self.right_tail_caps_[var] for var in right_only]) + clipped = np.minimum(values, upper) + new_series += [ + nw.new_series(var, clipped[:, i], backend=nw_X.implementation) + for i, var in enumerate(right_only) + ] + if len(left_only) > 0: + values = nw_X.select(nw.col(*left_only)).to_numpy() + lower = np.array([self.left_tail_caps_[var] for var in left_only]) + clipped = np.maximum(values, lower) + new_series += [ + nw.new_series(var, clipped[:, i], backend=nw_X.implementation) + for i, var in enumerate(left_only) + ] + + if len(new_series) > 0: + X = nw_X.with_columns(*new_series).to_native() return X @@ -205,16 +262,16 @@ def __init__( self.return_empty = return_empty self.missing_values = missing_values - def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): + def fit(self, X: IntoDataFrame, y: Optional[IntoSeries] = None): """ Learn the values that should be used to replace outliers. Parameters ---------- - X : pandas dataframe of shape = [n_samples, n_features] + X : dataframe of shape = [n_samples, n_features] The training input samples. - y : pandas Series, default=None + y : Series, default=None y is not needed in this transformer. You can pass y or None. """ @@ -242,22 +299,33 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): else: self.fold_ = self.fold + nw_X = nw.from_native(X, eager_only=True) + values = nw_X.select(nw.col(*self.variables_)).to_numpy() + + # nan-aware reductions: with missing_values="ignore", values may contain + # NaN, and pandas' mean/std/quantile/median skip NaN by default. if self.capping_method == "gaussian": - bias = X[self.variables_].mean() - scale = X[self.variables_].std(ddof=0) + bias = np.nanmean(values, axis=0) + scale = np.nanstd(values, axis=0, ddof=0) elif self.capping_method == "iqr": - bias = X[self.variables_].quantile((0.75, 0.25)) - scale = bias.loc[0.75] - bias.loc[0.25] + q75 = np.nanquantile(values, 0.75, axis=0) + q25 = np.nanquantile(values, 0.25, axis=0) + scale = q75 - q25 elif self.capping_method == "quantiles": - bias = X[self.variables_].quantile((1 - self.fold_, self.fold_)) - scale = bias.loc[1 - self.fold_] - bias.loc[self.fold_] + q_hi = np.nanquantile(values, 1 - self.fold_, axis=0) + q_lo = np.nanquantile(values, self.fold_, axis=0) + scale = q_hi - q_lo elif self.capping_method == "mad": - bias = X[self.variables_].median() + bias = np.nanmedian(values, axis=0) # scaling factor for normal distribution - scale = (X[self.variables_] - bias).abs().median() / 0.67449 + scale = np.nanmedian(np.abs(values - bias), axis=0) / 0.67449 + if (scale == 0).any(): + failing_vars = [ + var for var, s in zip(self.variables_, scale) if s == 0 + ] raise ValueError( - f"Input columns {scale[scale == 0].index.tolist()!r}" + f"Input columns {failing_vars!r}" f" have low variation for method {self.capping_method!r}." f" Try other capping methods or drop these columns." ) @@ -265,25 +333,45 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): # estimate the end values if self.tail in ("right", "both"): if self.capping_method in ("gaussian", "mad"): - self.right_tail_caps_ = (bias + self.fold_ * scale).to_dict() + self.right_tail_caps_ = { + var: float(b + self.fold_ * s) + for var, b, s in zip(self.variables_, bias, scale) + } elif self.capping_method == "iqr": - self.right_tail_caps_ = (bias.loc[0.75] + self.fold_ * scale).to_dict() + self.right_tail_caps_ = { + var: float(q + self.fold_ * s) + for var, q, s in zip(self.variables_, q75, scale) + } elif self.capping_method == "quantiles": - self.right_tail_caps_ = bias.loc[1 - self.fold_].to_dict() + self.right_tail_caps_ = { + var: float(q) for var, q in zip(self.variables_, q_hi) + } if self.tail in ("left", "both"): if self.capping_method in ("gaussian", "mad"): - self.left_tail_caps_ = (bias - self.fold_ * scale).to_dict() + self.left_tail_caps_ = { + var: float(b - self.fold_ * s) + for var, b, s in zip(self.variables_, bias, scale) + } elif self.capping_method == "iqr": - self.left_tail_caps_ = (bias.loc[0.25] - self.fold_ * scale).to_dict() + self.left_tail_caps_ = { + var: float(q - self.fold_ * s) + for var, q, s in zip(self.variables_, q25, scale) + } elif self.capping_method == "quantiles": - self.left_tail_caps_ = bias.loc[self.fold_].to_dict() + self.left_tail_caps_ = { + var: float(q) for var, q in zip(self.variables_, q_lo) + } - self.feature_names_in_ = X.columns.to_list() + is_pandas = nwd.is_pandas_dataframe(X) + if is_pandas is True: + self.feature_names_in_ = list(X.columns) + else: + self.feature_names_in_ = nw_X.columns self.n_features_in_ = X.shape[1] return self From 7b80eea9041a09cca4508a2a8e5567c4cd606674 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Wed, 26 Aug 2026 00:59:24 +0200 Subject: [PATCH 2/2] Migrate ArbitraryOutlierCapper to narwhals, add polars support fit() only builds dicts from user input and validates variables/dtypes via check_numerical_variables (already narwhals-generic) - no numeric computation, so nothing to branch on there. The only pandas-specific lines were the feature_names_in_ assignment (X.columns.to_list(), a pandas-Index method), replaced with the same is_pandas-guarded pattern WinsorizerBase.fit() already uses (list(X.columns) for pandas, nw.from_native(X).columns - already list[str] - otherwise). transform() was already dataframe-agnostic via BaseOutlier._transform(); only its type hints changed (pd.DataFrame -> IntoDataFrame). Benchmarked fit+transform end-to-end at 10k/50k/100k rows x 1/2/10 columns: pandas-native (pre-migration) vs the migrated code on pandas were within noise of each other (~0.9-1.1x), and polars ran 2-4x faster than pandas on both. No pandas/polars branch needed - merged single path, consistent with the is_pandas-only-for-.columns precedent already set in WinsorizerBase. Confirmed the module needs zero pandas: reloaded artbitrary.py in isolation with sys.modules["pandas"] = None (simulating an uninstalled pandas) and ran fit/transform end-to-end on a polars frame - works, and int64 stays int64 for a same-dtype capping dict (the class docstring's own x1 example). Found, while doing so, a real dtype-preservation bug in the already- merged BaseOutlier._transform() (base_outlier.py, commit 71bf7cf on this branch's base) that predates this migration and is not introduced here: when a capping-dict spans columns of different dtypes that land in the same bound-group (e.g. max_capping_dict={"age": 50, "fare": 200} with age int64 and fare float64 - both "right_only"), the group's columns are stacked into one 2D array via to_numpy() before np.clip, which forces a common dtype and upcasts age to float64. The pre- narwhals code (verified against 71bf7cf^) clipped each column independently (X[feature] = X[feature].clip(...)), so int columns never picked up a neighboring float column's dtype. Confirmed this reproduces identically on both pandas and polars (same merged code path) and is untouched by this commit - it lives in base_outlier.py, shared with Winsoriser/OutlierTrimmer, out of this file's scope. Flagged separately rather than fixed here. Rewrote test_arbitrary_capper.py to one parametrized test per behavior over pd.DataFrame/pl.DataFrame (previously pandas-only), using nw.from_native(...).to_dict(as_series=False) for backend-agnostic assertions in place of pd.testing.assert_frame_equal, following the same pattern used for ReciprocalTransformer/ArcsinTransformer. Added a verified "With polars" section to the docs (float dtypes throughout, to sidestep the dtype-upcast issue above rather than put an unexplained surprise in a user-facing example); left the pre-existing pandas Titanic walkthrough untouched - no network access in this environment to re-verify the fetch_openml/CSV-backed output. Verified: tests/test_outliers full suite - 88 passed (up from 83, all 5 new instances are the added polars parametrizations), same 3 pre-existing check_estimator failures as the pre-migration baseline (numpy-array input, unrelated to this change). flake8 and mypy clean. sphinx -W build clean (only the pre-existing linkcode_resolve warning). Co-Authored-By: Claude Sonnet 5 --- .../outliers/ArbitraryOutlierCapper.rst | 38 ++++++ feature_engine/outliers/artbitrary.py | 24 ++-- tests/test_outliers/test_arbitrary_capper.py | 123 ++++++++++-------- 3 files changed, 122 insertions(+), 63 deletions(-) diff --git a/docs/user_guide/outliers/ArbitraryOutlierCapper.rst b/docs/user_guide/outliers/ArbitraryOutlierCapper.rst index 70153b25b..ce916e9c4 100644 --- a/docs/user_guide/outliers/ArbitraryOutlierCapper.rst +++ b/docs/user_guide/outliers/ArbitraryOutlierCapper.rst @@ -96,6 +96,44 @@ values: dtype: float64 +With polars +----------- + +:class:`ArbitraryOutlierCapper()` works in the same way with a polars dataframe: + +.. code:: python + + import polars as pl + from feature_engine.outliers import ArbitraryOutlierCapper + + df = pl.DataFrame({ + "age": [20.0, 21.0, 19.0, 45.0, 67.0, 18.0, 90.0, 34.0, 55.0, 23.0], + "fare": [7.5, 8.0, 71.3, 13.0, 30.5, 7.9, 512.3, 26.0, 15.5, 8.6], + }) + + capper = ArbitraryOutlierCapper( + max_capping_dict={"age": 50, "fare": 200}, + min_capping_dict=None, + ) + + capper.fit(df) + Xt = capper.transform(df) + + print(Xt.select(["age", "fare"]).max()) + +The resulting maximum values, capped at the values we entered in the dictionary: + +.. code:: text + + shape: (1, 2) + ┌──────┬───────┐ + │ age ┆ fare │ + │ --- ┆ --- │ + │ f64 ┆ f64 │ + ╞══════╪═══════╡ + │ 50.0 ┆ 200.0 │ + └──────┴───────┘ + Additional resources -------------------- diff --git a/feature_engine/outliers/artbitrary.py b/feature_engine/outliers/artbitrary.py index 6088520da..d30b774b1 100644 --- a/feature_engine/outliers/artbitrary.py +++ b/feature_engine/outliers/artbitrary.py @@ -4,7 +4,9 @@ from typing import Optional -import pandas as pd +import narwhals as nw +import narwhals.dependencies as nwd +from narwhals.typing import IntoDataFrame, IntoSeries from feature_engine._check_init_parameters.check_input_dictionary import ( _check_numerical_dict, @@ -134,16 +136,16 @@ def __init__( self.min_capping_dict = min_capping_dict self.missing_values = missing_values - def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): + def fit(self, X: IntoDataFrame, y: Optional[IntoSeries] = None): """ This transformer does not learn any parameter. Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features] + X: dataframe of shape = [n_samples, n_features] The training input samples. - y: pandas Series, default=None + y: Series, default=None y is not needed in this transformer. You can pass y or None. """ X = check_X(X) @@ -176,23 +178,29 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): else: self.left_tail_caps_ = {} - self.feature_names_in_ = X.columns.to_list() + # pandas' .columns is an Index, not a list - list() is required there; + # narwhals' .columns is already list[str]. + is_pandas = nwd.is_pandas_dataframe(X) + if is_pandas is True: + self.feature_names_in_ = list(X.columns) + else: + self.feature_names_in_ = nw.from_native(X, eager_only=True).columns self.n_features_in_ = X.shape[1] return self - def transform(self, X: pd.DataFrame) -> pd.DataFrame: + def transform(self, X: IntoDataFrame) -> IntoDataFrame: """ Cap the variable values. 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_new: pandas dataframe of shape = [n_samples, n_features] + X_new: dataframe of shape = [n_samples, n_features] The dataframe with the capped variables. """ return super()._transform(X) diff --git a/tests/test_outliers/test_arbitrary_capper.py b/tests/test_outliers/test_arbitrary_capper.py index 5cba357b8..dae074295 100644 --- a/tests/test_outliers/test_arbitrary_capper.py +++ b/tests/test_outliers/test_arbitrary_capper.py @@ -1,22 +1,36 @@ +import narwhals as nw import numpy as np import pandas as pd +import polars as pl import pytest from feature_engine.outliers import ArbitraryOutlierCapper +DATA = {"var": list(np.random.RandomState(0).normal(0, 0.1, 20))} -def test_right_end_capping(df_normal_dist): - # test case 1: right end capping +DATA_NA = { + "Name": ["tom", "nick", "krish", "jack", "tom", "eric"], + "City": ["London", "Manchester", "Liverpool", "Bristol", "Manchester", "Liverpool"], + "Age": [20.0, 21.0, 19.0, 18.0, np.nan, 41.0], + "Marks": [0.9, 0.8, 0.7, 0.6, 0.5, 0.6], + "dob": pd.date_range("2020-02-24", periods=6, freq="min"), +} + + +def _to_dict(X): + return nw.from_native(X, eager_only=True).to_dict(as_series=False) + + +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_right_end_capping(make_df): + X = make_df(DATA) transformer = ArbitraryOutlierCapper( max_capping_dict={"var": 0.10727677848029868}, min_capping_dict=None ) - X = transformer.fit_transform(df_normal_dist) + Xt = transformer.fit_transform(X) # expected output - df_transf = df_normal_dist.copy() - df_transf["var"] = np.where( - df_transf["var"] > 0.10727677848029868, 0.10727677848029868, df_transf["var"] - ) + expected = [min(v, 0.10727677848029868) for v in DATA["var"]] # test init params assert np.round(transformer.max_capping_dict["var"], 3) == np.round( @@ -31,27 +45,24 @@ def test_right_end_capping(df_normal_dist): assert transformer.left_tail_caps_ == {} assert transformer.n_features_in_ == 1 # test transform output - pd.testing.assert_frame_equal(X, df_transf) - assert np.round(X["var"].max(), 3) <= np.round(0.10727677848029868, 3) - assert np.round(df_normal_dist["var"].max(), 3) > np.round(0.10727677848029868, 3) + result = _to_dict(Xt) + assert result["var"] == pytest.approx(expected) + assert max(result["var"]) <= 0.10727677848029868 + 1e-8 -def test_both_ends_capping(df_normal_dist): - # test case 2: both tails +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_both_ends_capping(make_df): + X = make_df(DATA) transformer = ArbitraryOutlierCapper( max_capping_dict={"var": 0.20857275540714884}, min_capping_dict={"var": -0.19661115230025186}, ) - X = transformer.fit_transform(df_normal_dist) + Xt = transformer.fit_transform(X) # expected output - df_transf = df_normal_dist.copy() - df_transf["var"] = np.where( - df_transf["var"] > 0.20857275540714884, 0.20857275540714884, df_transf["var"] - ) - df_transf["var"] = np.where( - df_transf["var"] < -0.19661115230025186, -0.19661115230025186, df_transf["var"] - ) + expected = [ + min(max(v, -0.19661115230025186), 0.20857275540714884) for v in DATA["var"] + ] # test fit params assert np.round(transformer.right_tail_caps_["var"], 3) == np.round( @@ -61,25 +72,22 @@ def test_both_ends_capping(df_normal_dist): -0.19661115230025186, 3 ) # test transform output - pd.testing.assert_frame_equal(X, df_transf) - assert np.round(X["var"].max(), 3) <= np.round(0.20857275540714884, 3) - assert np.round(X["var"].min(), 3) >= np.round(-0.19661115230025186, 3) - assert np.round(df_normal_dist["var"].max(), 3) > np.round(0.20857275540714884, 3) - assert np.round(df_normal_dist["var"].min(), 3) < np.round(-0.19661115230025186, 3) + result = _to_dict(Xt) + assert result["var"] == pytest.approx(expected) + assert max(result["var"]) <= 0.20857275540714884 + 1e-8 + assert min(result["var"]) >= -0.19661115230025186 - 1e-8 -def test_left_tail_capping(df_normal_dist): - # test case 3: left tail +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_left_tail_capping(make_df): + X = make_df(DATA) transformer = ArbitraryOutlierCapper( max_capping_dict=None, min_capping_dict={"var": -0.17486039103044} ) - X = transformer.fit_transform(df_normal_dist) + Xt = transformer.fit_transform(X) # expected output - df_transf = df_normal_dist.copy() - df_transf["var"] = np.where( - df_transf["var"] < -0.17486039103044, -0.17486039103044, df_transf["var"] - ) + expected = [max(v, -0.17486039103044) for v in DATA["var"]] # test init param assert transformer.max_capping_dict is None @@ -92,30 +100,32 @@ def test_left_tail_capping(df_normal_dist): -0.17486039103044, 3 ) # test transform output - pd.testing.assert_frame_equal(X, df_transf) - assert np.round(X["var"].min(), 3) >= np.round(-0.17486039103044, 3) - assert np.round(df_normal_dist["var"].min(), 3) < np.round(-0.17486039103044, 3) + result = _to_dict(Xt) + assert result["var"] == pytest.approx(expected) + assert min(result["var"]) >= -0.17486039103044 - 1e-8 -def test_ignores_na_in_input_df(df_na): - # test case 4: dataset contains na and transformer is asked to ignore them +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_ignores_na_in_input_df(make_df): + X = make_df(DATA_NA) transformer = ArbitraryOutlierCapper( max_capping_dict=None, min_capping_dict={"Age": 20}, missing_values="ignore" ) - X = transformer.fit_transform(df_na) + Xt = transformer.fit_transform(X) # expected output - df_transf = df_na.copy() - df_transf["Age"] = np.where(df_transf["Age"] < 20, 20, df_transf["Age"]) + expected = [ + v if np.isnan(v) else max(v, 20) for v in DATA_NA["Age"] + ] # test fit params assert transformer.max_capping_dict is None assert transformer.min_capping_dict == {"Age": 20} - assert transformer.n_features_in_ == 6 + assert transformer.n_features_in_ == 5 # test transform output - pd.testing.assert_frame_equal(X, df_transf) - assert X["Age"].min() >= 20 - assert df_na["Age"].min() < 20 + result = _to_dict(Xt) + assert result["Age"] == pytest.approx(expected, nan_ok=True) + assert np.nanmin(result["Age"]) >= 20 def test_error_if_max_capping_dict_wrong_input(): @@ -142,24 +152,29 @@ def test_error_if_missing_values_not_bool(): ArbitraryOutlierCapper(missing_values="other") -def test_fit_and_transform_raise_error_if_df_contains_na(df_normal_dist): - df_na = df_normal_dist.copy() - df_na.loc[1, "var"] = np.nan +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_fit_and_transform_raise_error_if_df_contains_na(make_df): + X = make_df(DATA) + data_na = dict(DATA) + var_na = list(DATA["var"]) + var_na[1] = np.nan + data_na["var"] = var_na + X_na = make_df(data_na) - # test case 5: when dataset contains na, fit method + # test case: when dataset contains na, fit method with pytest.raises(ValueError): transformer = ArbitraryOutlierCapper( min_capping_dict={"var": -0.17486039103044} ) - transformer.fit(df_na) + transformer.fit(X_na) - # test case 6: when dataset contains na, transform method + # test case: when dataset contains na, transform method with pytest.raises(ValueError): transformer = ArbitraryOutlierCapper( min_capping_dict={"var": -0.17486039103044} ) - transformer.fit(df_normal_dist) - transformer.transform(df_na) + transformer.fit(X) + transformer.transform(X_na) @pytest.mark.parametrize( @@ -168,9 +183,7 @@ def test_fit_and_transform_raise_error_if_df_contains_na(df_normal_dist): ) def test_error_if_missing_values_wrong_type(missing_values): msg = "missing_values takes only values 'raise' or 'ignore'" - with pytest.raises(ValueError) as record: + with pytest.raises(ValueError, match=msg): ArbitraryOutlierCapper( min_capping_dict={"var": -0.17486039103044}, missing_values="missing_values" ) - # check that error message matches - assert str(record.value) == msg