Permutation importance selector

[juniorcl]

Add a selector based on permutation importance from scikit-learn.

[juniorcl]

Class permutation importance selector

import numpy as np
import pandas as pd

from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.inspection import permutation_importance
from sklearn.utils.validation import check_is_fitted


class PermutationImportanceSelector(BaseEstimator, TransformerMixin):
    """
    Feature selector baseado em Permutation Importance.

    Parameters
    ----------
    estimator : sklearn-like estimator
        Modelo já configurado (ex: LightGBM, RandomForest, LogisticRegression).

    top_n : int
        Número de variáveis mais importantes a serem selecionadas.

    scoring : str or callable, default=None
        Métrica usada no permutation importance.
        Se None, usa a métrica padrão do estimador.

    n_repeats : int, default=10
        Número de permutações por feature.

    random_state : int, default=None
        Controle de aleatoriedade.

    n_jobs : int, default=-1
        Número de threads.
    """

    def __init__(
        self,
        estimator,
        top_n=10,
        scoring=None,
        n_repeats=10,
        random_state=None,
        n_jobs=-1,
    ):
        self.estimator = estimator
        self.top_n = top_n
        self.scoring = scoring
        self.n_repeats = n_repeats
        self.random_state = random_state
        self.n_jobs = n_jobs

    def fit(self, X, y):
        # Converte para DataFrame se necessário
        if not isinstance(X, pd.DataFrame):
            X = pd.DataFrame(X)

        self.feature_names_in_ = X.columns.to_numpy()

        # Treina o modelo
        self.estimator_ = self.estimator
        self.estimator_.fit(X, y)

        # Calcula permutation importance
        result = permutation_importance(
            self.estimator_,
            X,
            y,
            scoring=self.scoring,
            n_repeats=self.n_repeats,
            random_state=self.random_state,
            n_jobs=self.n_jobs,
        )

        self.importances_ = pd.DataFrame(
            {
                "feature": self.feature_names_in_,
                "importance_mean": result.importances_mean,
                "importance_std": result.importances_std,
            }
        ).sort_values("importance_mean", ascending=False)

        # Seleciona top_n
        self.selected_features_ = (
            self.importances_
            .head(self.top_n)["feature"]
            .to_numpy()
        )

        return self

    def transform(self, X):
        check_is_fitted(self, "selected_features_")

        if not isinstance(X, pd.DataFrame):
            X = pd.DataFrame(X, columns=self.feature_names_in_)

        return X[self.selected_features_]

    def get_feature_names_out(self, input_features=None):
        check_is_fitted(self, "selected_features_")
        return self.selected_features_

    def get_support(self):
        """
        Retorna uma máscara booleana das features selecionadas.
        """
        check_is_fitted(self, "selected_features_")
        return np.isin(self.feature_names_in_, self.selected_features_)