[python-package] adding max_category_values parameter to create_tree_digraph method (fixes #5687)

python-package/lightgbm/plotting.py

@@ -455,6 +455,7 @@ def _to_graphviz(
     orientation: str,
     constraints: Optional[List[int]],
     example_case: Optional[Union[np.ndarray, pd_DataFrame]],
+    max_category_values: int,
     **kwargs: Any
 ) -> Any:
     """Convert specified tree to graphviz instance.
@@ -477,6 +478,7 @@ def add(
         """Recursively add node or edge."""
         fillcolor = 'white'
         style = ''
+        tooltip = None
         if highlight:
             color = 'blue'
             penwidth = '3'
@@ -487,6 +489,7 @@ def add(
             shape = "rectangle"
             l_dec = 'yes'
             r_dec = 'no'
+            threshold = root['threshold']
             if root['decision_type'] == '<=':
                 operator = "&#8804;"
             elif root['decision_type'] == '==':
@@ -513,7 +516,13 @@ def add(
                         missing_type_str=root['missing_type'],
                         default_left=root['default_left']
                     )
-            label += f"<B>{_float2str(root['threshold'], precision)}</B>"
+            if root['decision_type'] == '==':
+                category_values = root['threshold'].split('||')
+                if len(category_values) > max_category_values:
+                    tooltip = root['threshold']
+                    threshold = '||'.join(category_values[:2]) + '||...||' + category_values[-1]
+
+            label += f"<B>{_float2str(threshold, precision)}</B>"
             for info in ['split_gain', 'internal_value', 'internal_weight', "internal_count", "data_percentage"]:
                 if info in show_info:
                     output = info.split('_')[-1]
@@ -557,7 +566,7 @@ def add(
             if "data_percentage" in show_info:
                 label += f"<br/>{_float2str(root['leaf_count'] / total_count * 100, 2)}% of data"
             label = f"<{label}>"
-        graph.node(name, label=label, shape=shape, style=style, fillcolor=fillcolor, color=color, penwidth=penwidth)
+        graph.node(name, label=label, shape=shape, style=style, fillcolor=fillcolor, color=color, penwidth=penwidth, tooltip=tooltip)
         if parent is not None:
             graph.edge(parent, name, decision, color=color, penwidth=penwidth)
 
@@ -603,6 +612,7 @@ def create_tree_digraph(
     precision: Optional[int] = 3,
     orientation: str = 'horizontal',
     example_case: Optional[Union[np.ndarray, pd_DataFrame]] = None,
+    max_category_values: int = 10,
     **kwargs: Any
 ) -> Any:
     """Create a digraph representation of specified tree.
@@ -646,6 +656,22 @@ def create_tree_digraph(
     example_case : numpy 2-D array, pandas DataFrame or None, optional (default=None)
         Single row with the same structure as the training data.
         If not None, the plot will highlight the path that sample takes through the tree.
+    max_category_values : int, optional (default=10)
+        The maximum number of category values to display in tree nodes, if the number of thresholds is greater than this value, thresholds will be collapsed and displayed on the label tooltip instead.
+
+        .. warning::
+
+            Consider wrapping the SVG string of the tree graph with ``IPython.display.HTML`` when running on JupyterLab to get the `tooltip <https://graphviz.org/docs/attrs/tooltip>`_ working right.
+
+            Example:
+
+            .. code-block:: python
+
+                from IPython.display import HTML
+
+                graph = lgb.create_tree_digraph(clf, max_category_values=5)
+                HTML(graph._repr_image_svg_xml())
+
     **kwargs
         Other parameters passed to ``Digraph`` constructor.
         Check https://graphviz.readthedocs.io/en/stable/api.html#digraph for the full list of supported parameters.
@@ -699,6 +725,7 @@ def create_tree_digraph(
         orientation=orientation,
         constraints=monotone_constraints,
         example_case=example_case,
+        max_category_values=max_category_values,
         **kwargs
     )
 

tests/python_package_test/test_plotting.py

@@ -1,5 +1,6 @@
 # coding: utf-8
 import numpy as np
+import pandas as pd
 import pytest
 from sklearn.model_selection import train_test_split
 
@@ -21,6 +22,20 @@ def breast_cancer_split():
                             test_size=0.1, random_state=1)
 
 
+@pytest.fixture(scope="module")
+def categorical_data():
+    def _categorical_data(category_values_lower_bound, category_values_upper_bound):
+        X, y = load_breast_cancer(return_X_y=True)
+        X_df = pd.DataFrame()
+        rnd = np.random.RandomState(0)
+        n_cat_values = rnd.randint(category_values_lower_bound, category_values_upper_bound, size=X.shape[1])
+        for i in range(X.shape[1]):
+            bins = np.linspace(0, 1, num=n_cat_values[i]+1)
+            X_df[f"cat_col_{i}"] = pd.qcut(X[:, i], q=bins, labels=range(n_cat_values[i])).as_unordered()
+        return train_test_split(X_df, y, test_size=0.1, random_state=1)
+    return _categorical_data
+
+
 @pytest.fixture(scope="module")
 def train_data(breast_cancer_split):
     X_train, _, y_train, _ = breast_cancer_split
@@ -188,6 +203,37 @@ def test_create_tree_digraph(breast_cancer_split):
     assert 'count' not in graph_body
 
 
+@pytest.mark.skipif(not GRAPHVIZ_INSTALLED, reason='graphviz is not installed')
+def test_tree_with_categories_below_max_category_values(categorical_data):
+    X_train, _, y_train, _ = categorical_data(2, 10)
+    gbm = lgb.LGBMClassifier(n_estimators=10, num_leaves=3, verbose=-1)
+    gbm.fit(X_train, y_train)
+
+    with pytest.raises(IndexError):
+        lgb.create_tree_digraph(gbm, tree_index=83)
+
+    graph = lgb.create_tree_digraph(gbm, tree_index=3,
+                                    show_info=['split_gain', 'internal_value', 'internal_weight'],
+                                    name='Tree4', node_attr={'color': 'red'},
+                                    max_category_values=10)
+    graph.render(view=False)
+    assert isinstance(graph, graphviz.Digraph)
+    assert graph.name == 'Tree4'
+    assert len(graph.node_attr) == 1
+    assert graph.node_attr['color'] == 'red'
+    assert len(graph.graph_attr) == 0
+    assert len(graph.edge_attr) == 0
+    graph_body = ''.join(graph.body)
+    assert 'leaf' in graph_body
+    assert 'gain' in graph_body
+    assert 'value' in graph_body
+    assert 'weight' in graph_body
+    assert 'data' not in graph_body
+    assert 'count' not in graph_body
+    assert '||...||' not in graph_body
+
+
+
 @pytest.mark.parametrize('use_missing', [True, False])
 @pytest.mark.parametrize('zero_as_missing', [True, False])
 def test_numeric_split_direction(use_missing, zero_as_missing):