matplotlib: Plot Feature Importance with feature names

add-semi-colons picture add-semi-colons · Jun 13, 2017 · Viewed 36.4k times · Source

In R there are pre-built functions to plot feature importance of Random Forest model. But in python such method seems to be missing. I search for a method in matplotlib.

model.feature_importances gives me following:

array([  2.32421835e-03,   7.21472336e-04,   2.70491223e-03,
         3.34521084e-03,   4.19443238e-03,   1.50108737e-03,
         3.29160540e-03,   4.82320256e-01,   3.14117333e-03])

Then using following plotting function:

>> pyplot.bar(range(len(model.feature_importances_)), model.feature_importances_)
>> pyplot.show()

I get a barplot but I would like to get barplot with labels while importance showing horizontally in a sorted fashion. I am also exploring seaborn and was not able to find a method.

Answer

fordy picture fordy · Mar 7, 2018

Quick answer for data scientists that ain't got no time to waste:

Load the feature importances into a pandas series indexed by your column names, then use its plot method. For a classifier model trained using X:

feat_importances = pd.Series(model.feature_importances_, index=X.columns)
feat_importances.nlargest(20).plot(kind='barh')

Slightly more detailed answer with a full example:

Assuming you trained your model with data contained in a pandas dataframe, this is fairly painless if you load the feature importance into a panda's series, then you can leverage its indexing to get the variable names displayed easily. The plot argument kind='barh' gives us a horizontal bar chart, but you could easily substitute this argument for kind='bar' for a traditional bar chart with the feature names along the x-axis if you prefer.

nlargest(n) is a pandas Series method which will return a subset of the series with the largest n values. This is useful if you've got lots of features in your model and you only want to plot the most important.

A quick complete example using the classic Kaggle Titanic dataset...

import pandas as pd
from sklearn.ensemble import RandomForestClassifier
%matplotlib inline            # don't forget this if you're using jupyter!

X = pd.read_csv("titanic_train.csv")
X = X[['Pclass', 'Age', 'Fare', 'Parch', 'SibSp', 'Survived']].dropna()
y = X.pop('Survived')

model = RandomForestClassifier()
model.fit(X, y)

(pd.Series(model.feature_importances_, index=X.columns)
   .nlargest(4)
   .plot(kind='barh'))        # some method chaining, because it's sexy!

Which will give you this:

sklearn random forest feature importances