feat: Make InformationFilter dataframe-agnostic

Author: FBruzzesi

sklego/preprocessing/projections.py

@@ -1,5 +1,5 @@
+import narwhals as nw
 import numpy as np
-import pandas as pd
 from sklearn.base import BaseEstimator, TransformerMixin
 from sklearn.utils import check_array
 from sklearn.utils.validation import check_is_fitted
@@ -181,16 +181,17 @@ def __init__(self, columns, alpha=1):
 
     def _check_coltype(self, X):
         """Check if the `columns` type(s) are compatible with `X` type."""
+        X_ = nw.from_native(X, strict=False, eager_only=True)
         for col in as_list(self.columns):
             if isinstance(col, str):
-                if isinstance(X, np.ndarray):
+                if isinstance(X_, np.ndarray):
                     raise ValueError(f"column {col} is a string but datatype receive is numpy.")
-                if isinstance(X, pd.DataFrame):
-                    if col not in X.columns:
-                        raise ValueError(f"column {col} is not in {X.columns}")
+                if isinstance(X_, nw.DataFrame):
+                    if col not in X_.columns:
+                        raise ValueError(f"column {col} is not in {X_.columns}")
             if isinstance(col, int):
-                if col not in range(np.atleast_2d(np.array(X)).shape[1]):
-                    raise ValueError(f"column {col} is out of bounds for input shape {X.shape}")
+                if col not in range(np.atleast_2d(np.array(X_)).shape[1]):
+                    raise ValueError(f"column {col} is out of bounds for input shape {X_.shape}")
 
     def _col_idx(self, X, name):
         """Get the column index of a column name."""

tests/test_preprocessing/test_informationfilter.py

@@ -1,5 +1,7 @@
+import narwhals as nw
 import numpy as np
 import pandas as pd
+import polars as pl
 import pytest
 from sklearn.datasets import fetch_openml
 from sklearn.linear_model import LinearRegression
@@ -49,57 +51,56 @@ def test_alpha_param1():
     assert np.isclose(ifilter.fit_transform(X), X_removed).all()
 
 
-def test_alpha_param2():
+@pytest.mark.parametrize("frame_func", [pd.DataFrame, pl.DataFrame])
+def test_alpha_param2(frame_func):
     X, y = fetch_openml(data_id=531, return_X_y=True, as_frame=False, parser="liac-arff")
-    df = pd.DataFrame(
-        X,
-        columns=[
-            "crim",
-            "zn",
-            "indus",
-            "chas",
-            "nox",
-            "rm",
-            "age",
-            "dis",
-            "rad",
-            "tax",
-            "ptratio",
-            "b",
-            "lstat",
-        ],
-    )
+    cols = [
+        "crim",
+        "zn",
+        "indus",
+        "chas",
+        "nox",
+        "rm",
+        "age",
+        "dis",
+        "rad",
+        "tax",
+        "ptratio",
+        "b",
+        "lstat",
+    ]
+    df = frame_func(dict(zip(cols, X.T)))
     ifilter = InformationFilter(columns=["b", "lstat"], alpha=0.0)
-    X_removed = df.drop(columns=["b", "lstat"]).values
+    X_removed = nw.from_native(df).drop(["b", "lstat"]).to_numpy()
     assert np.isclose(ifilter.fit_transform(df), X_removed).all()
 
 
-def test_output_orthogonal_pandas():
+@pytest.mark.parametrize("frame_func", [pd.DataFrame, pl.DataFrame])
+def test_output_orthogonal_frame(frame_func):
     X, y = fetch_openml(data_id=531, return_X_y=True, as_frame=False, parser="liac-arff")
-    df = pd.DataFrame(
-        X,
-        columns=[
-            "crim",
-            "zn",
-            "indus",
-            "chas",
-            "nox",
-            "rm",
-            "age",
-            "dis",
-            "rad",
-            "tax",
-            "ptratio",
-            "b",
-            "lstat",
-        ],
-    )
+    cols = [
+        "crim",
+        "zn",
+        "indus",
+        "chas",
+        "nox",
+        "rm",
+        "age",
+        "dis",
+        "rad",
+        "tax",
+        "ptratio",
+        "b",
+        "lstat",
+    ]
+    df = frame_func(dict(zip(cols, X.T)))
     X_fair = InformationFilter(columns=["b", "lstat"]).fit_transform(df)
     assert all([(c * df["b"]).sum() < 1e-5 for c in X_fair.T])
     assert all([(c * df["lstat"]).sum() < 1e-5 for c in X_fair.T])
 
 
-def test_output_orthogonal_general_cols():
+@pytest.mark.parametrize("frame_func", [pd.DataFrame, pl.DataFrame])
+def test_output_orthogonal_general_cols(frame_func):
     X, y = fetch_openml(data_id=531, return_X_y=True, as_frame=False, parser="liac-arff")
     cols = [
         "crim",
@@ -116,7 +117,7 @@ def test_output_orthogonal_general_cols():
         "b",
         "lstat",
     ]
-    df = pd.DataFrame(X, columns=cols)
+    df = frame_func(dict(zip(cols, X.T)))
     for col in cols:
         X_fair = InformationFilter(columns=col).fit_transform(df)
         assert all([(c * df[col]).sum() < 1e-5 for c in X_fair.T])