Skip to content

GitLab

Unverified Commit b6386842 authored by Thomas J. Fan's avatar Thomas J. Fan Committed by GitHub
Browse files

[python-package] migrate test_plotting.py to pytest (#3811) 

* TST Migrates test_plotting.py to pytest

* STY Fixes linting
parent bf22a25d
Hide whitespace changes
Inline Side-by-side
Showing with 184 additions and 161 deletions
tests/python_package_test/test_plotting.py
View file @ b6386842
# coding: utf-8 # coding: utf-8
import unittest
import lightgbm as lgb import lightgbm as lgb
from lightgbm.compat import MATPLOTLIB_INSTALLED, GRAPHVIZ_INSTALLED from lightgbm.compat import MATPLOTLIB_INSTALLED, GRAPHVIZ_INSTALLED
from sklearn.model_selection import train_test_split from sklearn.model_selection import train_test_split
import pytest
if MATPLOTLIB_INSTALLED: if MATPLOTLIB_INSTALLED:
import matplotlib import matplotlib
...@@ -14,164 +13,188 @@ if GRAPHVIZ_INSTALLED: ...@@ -14,164 +13,188 @@ if GRAPHVIZ_INSTALLED:
from .utils import load_breast_cancer from .utils import load_breast_cancer
class TestBasic(unittest.TestCase): @pytest.fixture(scope="module")
def breast_cancer_split():
return train_test_split(*load_breast_cancer(return_X_y=True),
test_size=0.1, random_state=1)
@pytest.fixture(scope="module")
def train_data(breast_cancer_split):
X_train, _, y_train, _ = breast_cancer_split
return lgb.Dataset(X_train, y_train)
def setUp(self): @pytest.fixture
self.X_train, self.X_test, self.y_train, self.y_test = train_test_split(*load_breast_cancer(return_X_y=True), def params():
test_size=0.1, random_state=1) return {"objective": "binary",
self.train_data = lgb.Dataset(self.X_train, self.y_train)
self.params = {
"objective": "binary",
"verbose": -1, "verbose": -1,
"num_leaves": 3 "num_leaves": 3}
}
@unittest.skipIf(not MATPLOTLIB_INSTALLED, 'matplotlib is not installed') @pytest.mark.skipif(not MATPLOTLIB_INSTALLED, reason='matplotlib is not installed')
def test_plot_importance(self): def test_plot_importance(params, breast_cancer_split, train_data):
gbm0 = lgb.train(self.params, self.train_data, num_boost_round=10) X_train, _, y_train, _ = breast_cancer_split
ax0 = lgb.plot_importance(gbm0)
self.assertIsInstance(ax0, matplotlib.axes.Axes) gbm0 = lgb.train(params, train_data, num_boost_round=10)
self.assertEqual(ax0.get_title(), 'Feature importance') ax0 = lgb.plot_importance(gbm0)
self.assertEqual(ax0.get_xlabel(), 'Feature importance') assert isinstance(ax0, matplotlib.axes.Axes)
self.assertEqual(ax0.get_ylabel(), 'Features') assert ax0.get_title() == 'Feature importance'
self.assertLessEqual(len(ax0.patches), 30) assert ax0.get_xlabel() == 'Feature importance'
assert ax0.get_ylabel() == 'Features'
gbm1 = lgb.LGBMClassifier(n_estimators=10, num_leaves=3, silent=True) assert len(ax0.patches) <= 30
gbm1.fit(self.X_train, self.y_train)
gbm1 = lgb.LGBMClassifier(n_estimators=10, num_leaves=3, silent=True)
ax1 = lgb.plot_importance(gbm1, color='r', title='t', xlabel='x', ylabel='y') gbm1.fit(X_train, y_train)
self.assertIsInstance(ax1, matplotlib.axes.Axes)
self.assertEqual(ax1.get_title(), 't') ax1 = lgb.plot_importance(gbm1, color='r', title='t', xlabel='x', ylabel='y')
self.assertEqual(ax1.get_xlabel(), 'x') assert isinstance(ax1, matplotlib.axes.Axes)
self.assertEqual(ax1.get_ylabel(), 'y') assert ax1.get_title() == 't'
self.assertLessEqual(len(ax1.patches), 30) assert ax1.get_xlabel() == 'x'
for patch in ax1.patches: assert ax1.get_ylabel() == 'y'
self.assertTupleEqual(patch.get_facecolor(), (1., 0, 0, 1.)) # red assert len(ax1.patches) <= 30
for patch in ax1.patches:
ax2 = lgb.plot_importance(gbm0, color=['r', 'y', 'g', 'b'], assert patch.get_facecolor() == (1., 0, 0, 1.) # red
title=None, xlabel=None, ylabel=None)
self.assertIsInstance(ax2, matplotlib.axes.Axes) ax2 = lgb.plot_importance(gbm0, color=['r', 'y', 'g', 'b'],
self.assertEqual(ax2.get_title(), '') title=None, xlabel=None, ylabel=None)
self.assertEqual(ax2.get_xlabel(), '') assert isinstance(ax2, matplotlib.axes.Axes)
self.assertEqual(ax2.get_ylabel(), '') assert ax2.get_title() == ''
self.assertLessEqual(len(ax2.patches), 30) assert ax2.get_xlabel() == ''
self.assertTupleEqual(ax2.patches[0].get_facecolor(), (1., 0, 0, 1.)) # r assert ax2.get_ylabel() == ''
self.assertTupleEqual(ax2.patches[1].get_facecolor(), (.75, .75, 0, 1.)) # y assert len(ax2.patches) <= 30
self.assertTupleEqual(ax2.patches[2].get_facecolor(), (0, .5, 0, 1.)) # g assert ax2.patches[0].get_facecolor() == (1., 0, 0, 1.) # r
self.assertTupleEqual(ax2.patches[3].get_facecolor(), (0, 0, 1., 1.)) # b assert ax2.patches[1].get_facecolor() == (.75, .75, 0, 1.) # y
assert ax2.patches[2].get_facecolor() == (0, .5, 0, 1.) # g
@unittest.skipIf(not MATPLOTLIB_INSTALLED, 'matplotlib is not installed') assert ax2.patches[3].get_facecolor() == (0, 0, 1., 1.) # b
def test_plot_split_value_histogram(self):
gbm0 = lgb.train(self.params, self.train_data, num_boost_round=10)
ax0 = lgb.plot_split_value_histogram(gbm0, 27) @pytest.mark.skipif(not MATPLOTLIB_INSTALLED, reason='matplotlib is not installed')
self.assertIsInstance(ax0, matplotlib.axes.Axes) def test_plot_split_value_histogram(params, breast_cancer_split, train_data):
self.assertEqual(ax0.get_title(), 'Split value histogram for feature with index 27') X_train, _, y_train, _ = breast_cancer_split
self.assertEqual(ax0.get_xlabel(), 'Feature split value')
self.assertEqual(ax0.get_ylabel(), 'Count') gbm0 = lgb.train(params, train_data, num_boost_round=10)
self.assertLessEqual(len(ax0.patches), 2) ax0 = lgb.plot_split_value_histogram(gbm0, 27)
assert isinstance(ax0, matplotlib.axes.Axes)
gbm1 = lgb.LGBMClassifier(n_estimators=10, num_leaves=3, silent=True) assert ax0.get_title() == 'Split value histogram for feature with index 27'
gbm1.fit(self.X_train, self.y_train) assert ax0.get_xlabel() == 'Feature split value'
assert ax0.get_ylabel() == 'Count'
ax1 = lgb.plot_split_value_histogram(gbm1, gbm1.booster_.feature_name()[27], figsize=(10, 5), assert len(ax0.patches) <= 2
title='Histogram for feature @index/name@ @feature@',
xlabel='x', ylabel='y', color='r') gbm1 = lgb.LGBMClassifier(n_estimators=10, num_leaves=3, silent=True)
self.assertIsInstance(ax1, matplotlib.axes.Axes) gbm1.fit(X_train, y_train)
self.assertEqual(ax1.get_title(),
'Histogram for feature name {}'.format(gbm1.booster_.feature_name()[27])) ax1 = lgb.plot_split_value_histogram(gbm1, gbm1.booster_.feature_name()[27], figsize=(10, 5),
self.assertEqual(ax1.get_xlabel(), 'x') title='Histogram for feature @index/name@ @feature@',
self.assertEqual(ax1.get_ylabel(), 'y') xlabel='x', ylabel='y', color='r')
self.assertLessEqual(len(ax1.patches), 2) assert isinstance(ax1, matplotlib.axes.Axes)
for patch in ax1.patches: title = 'Histogram for feature name {}'.format(gbm1.booster_.feature_name()[27])
self.assertTupleEqual(patch.get_facecolor(), (1., 0, 0, 1.)) # red assert ax1.get_title() == title
assert ax1.get_xlabel() == 'x'
ax2 = lgb.plot_split_value_histogram(gbm0, 27, bins=10, color=['r', 'y', 'g', 'b'], assert ax1.get_ylabel() == 'y'
title=None, xlabel=None, ylabel=None) assert len(ax1.patches) <= 2
self.assertIsInstance(ax2, matplotlib.axes.Axes) for patch in ax1.patches:
self.assertEqual(ax2.get_title(), '') assert patch.get_facecolor() == (1., 0, 0, 1.) # red
self.assertEqual(ax2.get_xlabel(), '')
self.assertEqual(ax2.get_ylabel(), '') ax2 = lgb.plot_split_value_histogram(gbm0, 27, bins=10, color=['r', 'y', 'g', 'b'],
self.assertEqual(len(ax2.patches), 10) title=None, xlabel=None, ylabel=None)
self.assertTupleEqual(ax2.patches[0].get_facecolor(), (1., 0, 0, 1.)) # r assert isinstance(ax2, matplotlib.axes.Axes)
self.assertTupleEqual(ax2.patches[1].get_facecolor(), (.75, .75, 0, 1.)) # y assert ax2.get_title() == ''
self.assertTupleEqual(ax2.patches[2].get_facecolor(), (0, .5, 0, 1.)) # g assert ax2.get_xlabel() == ''
self.assertTupleEqual(ax2.patches[3].get_facecolor(), (0, 0, 1., 1.)) # b assert ax2.get_ylabel() == ''
assert len(ax2.patches) == 10
self.assertRaises(ValueError, lgb.plot_split_value_histogram, gbm0, 0) # was not used in splitting assert ax2.patches[0].get_facecolor() == (1., 0, 0, 1.) # r
assert ax2.patches[1].get_facecolor() == (.75, .75, 0, 1.) # y
@unittest.skipIf(not MATPLOTLIB_INSTALLED or not GRAPHVIZ_INSTALLED, 'matplotlib or graphviz is not installed') assert ax2.patches[2].get_facecolor() == (0, .5, 0, 1.) # g
def test_plot_tree(self): assert ax2.patches[3].get_facecolor() == (0, 0, 1., 1.) # b
gbm = lgb.LGBMClassifier(n_estimators=10, num_leaves=3, silent=True)
gbm.fit(self.X_train, self.y_train, verbose=False) with pytest.raises(ValueError):
lgb.plot_split_value_histogram(gbm0, 0) # was not used in splitting
self.assertRaises(IndexError, lgb.plot_tree, gbm, tree_index=83)
ax = lgb.plot_tree(gbm, tree_index=3, figsize=(15, 8), show_info=['split_gain']) @pytest.mark.skipif(not MATPLOTLIB_INSTALLED or not GRAPHVIZ_INSTALLED,
self.assertIsInstance(ax, matplotlib.axes.Axes) reason='matplotlib or graphviz is not installed')
w, h = ax.axes.get_figure().get_size_inches() def test_plot_tree(breast_cancer_split):
self.assertEqual(int(w), 15) X_train, _, y_train, _ = breast_cancer_split
self.assertEqual(int(h), 8) gbm = lgb.LGBMClassifier(n_estimators=10, num_leaves=3, silent=True)
gbm.fit(X_train, y_train, verbose=False)
@unittest.skipIf(not GRAPHVIZ_INSTALLED, 'graphviz is not installed')
def test_create_tree_digraph(self): with pytest.raises(IndexError):
constraints = [-1, 1] * int(self.X_train.shape[1] / 2) lgb.plot_tree(gbm, tree_index=83)
gbm = lgb.LGBMClassifier(n_estimators=10, num_leaves=3, silent=True, monotone_constraints=constraints)
gbm.fit(self.X_train, self.y_train, verbose=False) ax = lgb.plot_tree(gbm, tree_index=3, figsize=(15, 8), show_info=['split_gain'])
assert isinstance(ax, matplotlib.axes.Axes)
self.assertRaises(IndexError, lgb.create_tree_digraph, gbm, tree_index=83) w, h = ax.axes.get_figure().get_size_inches()
assert int(w) == 15
graph = lgb.create_tree_digraph(gbm, tree_index=3, assert int(h) == 8
show_info=['split_gain', 'internal_value', 'internal_weight'],
name='Tree4', node_attr={'color': 'red'})
graph.render(view=False) @pytest.mark.skipif(not GRAPHVIZ_INSTALLED, reason='graphviz is not installed')
self.assertIsInstance(graph, graphviz.Digraph) def test_create_tree_digraph(breast_cancer_split):
self.assertEqual(graph.name, 'Tree4') X_train, _, y_train, _ = breast_cancer_split
self.assertEqual(graph.filename, 'Tree4.gv')
self.assertEqual(len(graph.node_attr), 1) constraints = [-1, 1] * int(X_train.shape[1] / 2)
self.assertEqual(graph.node_attr['color'], 'red') gbm = lgb.LGBMClassifier(n_estimators=10, num_leaves=3, silent=True, monotone_constraints=constraints)
self.assertEqual(len(graph.graph_attr), 0) gbm.fit(X_train, y_train, verbose=False)
self.assertEqual(len(graph.edge_attr), 0)
graph_body = ''.join(graph.body) with pytest.raises(IndexError):
self.assertIn('leaf', graph_body) lgb.create_tree_digraph(gbm, tree_index=83)
self.assertIn('gain', graph_body)
self.assertIn('value', graph_body) graph = lgb.create_tree_digraph(gbm, tree_index=3,
self.assertIn('weight', graph_body) show_info=['split_gain', 'internal_value', 'internal_weight'],
self.assertIn('#ffdddd', graph_body) name='Tree4', node_attr={'color': 'red'})
self.assertIn('#ddffdd', graph_body) graph.render(view=False)
self.assertNotIn('data', graph_body) assert isinstance(graph, graphviz.Digraph)
self.assertNotIn('count', graph_body) assert graph.name == 'Tree4'
assert graph.filename == 'Tree4.gv'
@unittest.skipIf(not MATPLOTLIB_INSTALLED, 'matplotlib is not installed') assert len(graph.node_attr) == 1
def test_plot_metrics(self): assert graph.node_attr['color'] == 'red'
test_data = lgb.Dataset(self.X_test, self.y_test, reference=self.train_data) assert len(graph.graph_attr) == 0
self.params.update({"metric": {"binary_logloss", "binary_error"}}) assert len(graph.edge_attr) == 0
graph_body = ''.join(graph.body)
evals_result0 = {} assert 'leaf' in graph_body
gbm0 = lgb.train(self.params, self.train_data, assert 'gain' in graph_body
valid_sets=[self.train_data, test_data], assert 'value' in graph_body
valid_names=['v1', 'v2'], assert 'weight' in graph_body
num_boost_round=10, assert '#ffdddd' in graph_body
evals_result=evals_result0, assert '#ddffdd' in graph_body
verbose_eval=False) assert 'data' not in graph_body
ax0 = lgb.plot_metric(evals_result0) assert 'count' not in graph_body
self.assertIsInstance(ax0, matplotlib.axes.Axes)
self.assertEqual(ax0.get_title(), 'Metric during training')
self.assertEqual(ax0.get_xlabel(), 'Iterations') @pytest.mark.skipif(not MATPLOTLIB_INSTALLED, reason='matplotlib is not installed')
self.assertIn(ax0.get_ylabel(), {'binary_logloss', 'binary_error'}) def test_plot_metrics(params, breast_cancer_split, train_data):
ax0 = lgb.plot_metric(evals_result0, metric='binary_error') X_train, X_test, y_train, y_test = breast_cancer_split
ax0 = lgb.plot_metric(evals_result0, metric='binary_logloss', dataset_names=['v2']) test_data = lgb.Dataset(X_test, y_test, reference=train_data)
params.update({"metric": {"binary_logloss", "binary_error"}})
evals_result1 = {}
gbm1 = lgb.train(self.params, self.train_data, evals_result0 = {}
num_boost_round=10, lgb.train(params, train_data,
evals_result=evals_result1, valid_sets=[train_data, test_data],
verbose_eval=False) valid_names=['v1', 'v2'],
self.assertRaises(ValueError, lgb.plot_metric, evals_result1) num_boost_round=10,
evals_result=evals_result0,
gbm2 = lgb.LGBMClassifier(n_estimators=10, num_leaves=3, silent=True) verbose_eval=False)
gbm2.fit(self.X_train, self.y_train, eval_set=[(self.X_test, self.y_test)], verbose=False) ax0 = lgb.plot_metric(evals_result0)
ax2 = lgb.plot_metric(gbm2, title=None, xlabel=None, ylabel=None) assert isinstance(ax0, matplotlib.axes.Axes)
self.assertIsInstance(ax2, matplotlib.axes.Axes) assert ax0.get_title() == 'Metric during training'
self.assertEqual(ax2.get_title(), '') assert ax0.get_xlabel() == 'Iterations'
self.assertEqual(ax2.get_xlabel(), '') assert ax0.get_ylabel() in {'binary_logloss', 'binary_error'}
self.assertEqual(ax2.get_ylabel(), '') ax0 = lgb.plot_metric(evals_result0, metric='binary_error')
ax0 = lgb.plot_metric(evals_result0, metric='binary_logloss', dataset_names=['v2'])
evals_result1 = {}
lgb.train(params, train_data,
num_boost_round=10,
evals_result=evals_result1,
verbose_eval=False)
with pytest.raises(ValueError):
lgb.plot_metric(evals_result1)
gbm2 = lgb.LGBMClassifier(n_estimators=10, num_leaves=3, silent=True)
gbm2.fit(X_train, y_train, eval_set=[(X_test, y_test)], verbose=False)
ax2 = lgb.plot_metric(gbm2, title=None, xlabel=None, ylabel=None)
assert isinstance(ax2, matplotlib.axes.Axes)
assert ax2.get_title() == ''
assert ax2.get_xlabel() == ''
assert ax2.get_ylabel() == ''
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment