Scikit-Learn Permutating and Updating Polars DataFrame

Question

I am trying to re-write the source code of scikit-learn permutation importance to achieve:

1. Compatibility with Polars
2. Compatibility with clusters of features

import polars as pl
import polars.selectors as cs
import numpy as np

from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split

X, y = make_classification(
    n_samples=1000,
    n_features=10,
    n_informative=3,
    n_redundant=0,
    n_repeated=0,
    n_classes=2,
    random_state=42,
    shuffle=False,
)
X_train, X_test, y_train, y_test = train_test_split(X, y, stratify=y, random_state=42)
feature_names = [f"feature_{i}" for i in range(X.shape[1])]

X_train_polars = pl.DataFrame(X_train, schema=feature_names)
X_test_polars = pl.DataFrame(X_test, schema=feature_names)
y_train_polars = pl.Series(y_train, schema=["target"])
y_test_polars = pl.Series(y_test, schema=["target"])

To get future importances for a cluster of feature, we need to permutate a cluster of features simutiousnly then pass into the scorer to compare with the baseline score.

However, I am struglling to replace multiple polars dataframe columns in case of examining clusters of features:

from sklearn.utils import check_random_state
random_state = check_random_state(42)
random_seed = random_state.randint(np.iinfo(np.int32).max + 1)

X_train_permuted = X_train_polars.clone()
shuffle_arr = np.array(X_train_permuted[:, ["feature_0", "feature_1"]])

random_state.shuffle(shuffle_arr)
X_train_permuted.replace_column( # This operation is in place
                0, 
                pl.Series(name="feature_0", values=shuffle_arr))

Normally the shuffle_arr would have a shape of (n_samples,) which can easily replace assosicated column in polars dataframe using polars.DataFrame.replace_column(). In this case, shuffle_arr has multi-dimensional shape of (n_samples, n_features in a cluster). What would be an efficient way to replace the assosicated columns?

Answer 1

TL;DR

X_train_permuted = (
    X_train_permuted.with_columns(
        pl.DataFrame(shuffle_arr, schema=features)
    )
)

---

Let's work with a simple example.

X_train_permuted

import polars as pl
import numpy as np

np.random.seed(0)

data = {f'feature_{i}': np.random.rand(4) for i in range(0,3)}

X_train_permuted = pl.DataFrame(data)

X_train_permuted

shape: (4, 3)
┌───────────┬───────────┬───────────┐
│ feature_0 ┆ feature_1 ┆ feature_2 │
│ ---       ┆ ---       ┆ ---       │
│ f64       ┆ f64       ┆ f64       │
╞═══════════╪═══════════╪═══════════╡
│ 0.548814  ┆ 0.423655  ┆ 0.963663  │
│ 0.715189  ┆ 0.645894  ┆ 0.383442  │
│ 0.602763  ┆ 0.437587  ┆ 0.791725  │
│ 0.544883  ┆ 0.891773  ┆ 0.528895  │
└───────────┴───────────┴───────────┘

Shuffle feature_0 and feature_1

Use a list to keep track of the features you are shuffling: features = ["feature_0", "feature_1"].

features = ["feature_0", "feature_1"]

shuffle_arr = np.array(X_train_permuted[:, features])

from sklearn.utils import check_random_state

random_state = check_random_state(42)
random_seed = random_state.randint(np.iinfo(np.int32).max + 1)

random_state.shuffle(shuffle_arr)

shuffle_arr

array([[0.71518937, 0.64589411],
       [0.60276338, 0.43758721],
       [0.5488135 , 0.4236548 ],
       [0.54488318, 0.891773  ]])

Replace associated columns in X_train_permuted with shuffle_arr values

• Use pl.DataFrame.with_columns and pass a pl.DataFrame with schema=features.

X_train_permuted = (
    X_train_permuted.with_columns(
        pl.DataFrame(shuffle_arr, schema=features)
    )
)

X_train_permuted

shape: (4, 3)
┌───────────┬───────────┬───────────┐
│ feature_0 ┆ feature_1 ┆ feature_2 │
│ ---       ┆ ---       ┆ ---       │
│ f64       ┆ f64       ┆ f64       │
╞═══════════╪═══════════╪═══════════╡
│ 0.715189  ┆ 0.645894  ┆ 0.963663  │
│ 0.602763  ┆ 0.437587  ┆ 0.383442  │
│ 0.548814  ┆ 0.423655  ┆ 0.791725  │
│ 0.544883  ┆ 0.891773  ┆ 0.528895  │
└───────────┴───────────┴───────────┘