From c5e73f4d72bd0995f1c25c07db2dd3812dd086d5 Mon Sep 17 00:00:00 2001 From: Andrew Date: Wed, 23 Mar 2016 07:50:18 -0700 Subject: [PATCH] issue 6580: add check for equal color bounds to prevent division by zero error. Add corresponding unit test. --- sklearn/tree/export.py | 9 ++++++--- sklearn/tree/tests/test_export.py | 17 +++++++++++++++++ 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/sklearn/tree/export.py b/sklearn/tree/export.py index ff1355db478ee..551bf2fe4daab 100644 --- a/sklearn/tree/export.py +++ b/sklearn/tree/export.py @@ -157,9 +157,12 @@ def get_color(value): else: # Regression tree or multi-output color = list(colors['rgb'][0]) - alpha = int(np.round(255 * ((value - colors['bounds'][0]) / - (colors['bounds'][1] - - colors['bounds'][0])), 0)) + if colors['bounds'][0] == colors['bounds'][1]: + alpha = 255 + else: + alpha = int(np.round(255 * ((value - colors['bounds'][0]) / + (colors['bounds'][1] - + colors['bounds'][0])), 0)) # Return html color code in #RRGGBBAA format color.append(alpha) diff --git a/sklearn/tree/tests/test_export.py b/sklearn/tree/tests/test_export.py index 8a7421923b152..e0dc801979af6 100644 --- a/sklearn/tree/tests/test_export.py +++ b/sklearn/tree/tests/test_export.py @@ -17,6 +17,7 @@ X = [[-2, -1], [-1, -1], [-1, -2], [1, 1], [1, 2], [2, 1]] y = [-1, -1, -1, 1, 1, 1] y2 = [[-1, 1], [-1, 1], [-1, 1], [1, 2], [1, 2], [1, 3]] +y3 = [0, 0, 0, 0, 0, 0] w = [1, 1, 1, .5, .5, .5] @@ -235,3 +236,19 @@ def test_friedman_mse_in_graphviz(): for finding in finditer("\[.*?samples.*?\]", dot_data.getvalue()): assert_in("friedman_mse", finding.group()) + + +def test_graphviz_single_node_tree(): + dtc = DecisionTreeClassifier() + dtc.fit(X, y3) + + out = StringIO() + export_graphviz(dtc, out_file=out, filled=True) + + contents1 = out.getvalue() + contents2 = 'digraph Tree {\n' \ + 'node [shape=box, style="filled", color="black"] ;\n' \ + '0 [label="gini = 0.0\\nsamples = 6\\nvalue = 6.0", fillcolor="#e58139ff"] ;\n' \ + '}' + + assert_equal(contents1, contents2)