如何修改permutation importance代码仅绘制前3个重要特征

你只需要调整排序索引的截取逻辑，仅保留重要性最高的3个特征的索引即可，修改后的完整代码如下：

import matplotlib.pyplot as plt
import numpy as np
from sklearn.datasets import fetch_openml
from sklearn.ensemble import RandomForestClassifier
from sklearn.impute import SimpleImputer
from sklearn.inspection import permutation_importance
from sklearn.compose import ColumnTransformer
from sklearn.model_selection import train_test_split
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder


result = permutation_importance(rf,
                                X_test,
                                y_test,
                                n_repeats=10,
                                random_state=42,
                                n_jobs=2)
# 取升序排序的最后3个（即重要性最高的3个），可选反转让重要性最高的在图最上方
sorted_idx = result.importances_mean.argsort()[-3:][::-1]
        

fig, ax = plt.subplots()
ax.boxplot(result.importances[sorted_idx].T,
           vert=False,
           labels=X_test.columns[sorted_idx])

ax.set_title("Permutation Importances (test set)")
fig.tight_layout()
plt.show()

核心修改说明 #

• argsort()默认返回升序排列的索引，因此result.importances_mean.argsort()得到的是按重要性从小到大排序的特征索引，取末尾[-3:]就能筛选出重要性最高的3个特征的索引
• 末尾加[::-1]是为了让箱线图从上到下按重要性从高到低展示，符合常规阅读习惯，不需要可以直接删掉这部分

内容的提问来源于stack exchange，提问作者user308827