Hendrik/xgboost
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

added the max_features parameter to the plot_importance function. (#1963)

Browse Source 
* added the max_features parameter to the plot_importance function.

* renamed max_features parameter to max_num_features for better understanding

* removed unwanted character in docstring
This commit is contained in:
Félix MIKAELIAN 2017-01-16 23:49:47 +01:00 committed by Tianqi Chen
parent 49ff7c1649
commit a7d2833766
1 changed files with 8 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
python-package/xgboost/plotting.py
View File

@ -14,7 +14,7 @@ from .sklearn import XGBModel
def plot_importance(booster, ax=None, height=0.2, def plot_importance(booster, ax=None, height=0.2,
xlim=None, ylim=None, title='Feature importance', xlim=None, ylim=None, title='Feature importance',
xlabel='F score', ylabel='Features', xlabel='F score', ylabel='Features',
importance_type='weight', importance_type='weight', max_num_features=None,
grid=True, **kwargs): grid=True, **kwargs):
"""Plot importance based on fitted trees. """Plot importance based on fitted trees.
@ -31,6 +31,8 @@ def plot_importance(booster, ax=None, height=0.2,
"gain" is the average gain of splits which use the feature "gain" is the average gain of splits which use the feature
"cover" is the average coverage of splits which use the feature "cover" is the average coverage of splits which use the feature
where coverage is defined as the number of samples affected by the split where coverage is defined as the number of samples affected by the split
max_num_features : int, default None
Maximum number of top features displayed on plot. If None, all features will be displayed.
height : float, default 0.2 height : float, default 0.2
Bar height, passed to ax.barh() Bar height, passed to ax.barh()
xlim : tuple, default None xlim : tuple, default None
@ -69,6 +71,10 @@ def plot_importance(booster, ax=None, height=0.2,
raise ValueError('Booster.get_score() results in empty') raise ValueError('Booster.get_score() results in empty')
tuples = [(k, importance[k]) for k in importance] tuples = [(k, importance[k]) for k in importance]
if max_num_features is not None:
tuples = sorted(tuples, key=lambda x: x[1])[-max_num_features:]
ylim = (-1, max_num_features)
else:
tuples = sorted(tuples, key=lambda x: x[1]) tuples = sorted(tuples, key=lambda x: x[1])
labels, values = zip(*tuples) labels, values = zip(*tuples)