Function _to_graphviz

python-package/lightgbmmt/plotting.py:377–468

Convert specified tree to graphviz instance. See: - https://graphviz.readthedocs.io/en/stable/api.html#digraph

(tree_info, show_info, feature_names, precision=3, constraints=None, **kwargs)

def _to_graphviz(tree_info, show_info, feature_names, precision=3, constraints=None, **kwargs):
    """Convert specified tree to graphviz instance.

    See:
      - https://graphviz.readthedocs.io/en/stable/api.html#digraph
    """
    if GRAPHVIZ_INSTALLED:
        from graphviz import Digraph
    else:
        raise ImportError('You must install graphviz to plot tree.')

    def add(root, total_count, parent=None, decision=None):
        """Recursively add node or edge."""
        if 'split_index' in root:  # non-leaf
            l_dec = 'yes'
            r_dec = 'no'
            if root['decision_type'] == '<=':
                lte_symbol = "&#8804;"
                operator = lte_symbol
            elif root['decision_type'] == '==':
                operator = "="
            else:
                raise ValueError('Invalid decision type in tree model.')
            name = 'split{0}'.format(root['split_index'])
            if feature_names is not None:
                label = '<B>{0}</B> {1} '.format(feature_names[root['split_feature']], operator)
            else:
                label = 'feature <B>{0}</B> {1} '.format(root['split_feature'], operator)
            label += '<B>{0}</B>'.format(_float2str(root['threshold'], precision))
            for info in ['split_gain', 'internal_value', 'internal_weight', "internal_count", "data_percentage"]:
                if info in show_info:
                    output = info.split('_')[-1]
                    if info in {'split_gain', 'internal_value', 'internal_weight'}:
                        label += '<br/>{0} {1}'.format(_float2str(root[info], precision), output)
                    elif info == 'internal_count':
                        label += '<br/>{0}: {1}'.format(output, root[info])
                    elif info == "data_percentage":
                        label += '<br/>{0}% of data'.format(_float2str(root['internal_count'] / total_count * 100, 2))

            fillcolor = "white"
            style = ""
            if constraints:
                if constraints[root['split_feature']] == 1:
                    fillcolor = "#ddffdd"  # light green
                if constraints[root['split_feature']] == -1:
                    fillcolor = "#ffdddd"  # light red
                style = "filled"
            label = "<" + label + ">"
            graph.node(name, label=label, shape="rectangle", style=style, fillcolor=fillcolor)
            add(root['left_child'], total_count, name, l_dec)
            add(root['right_child'], total_count, name, r_dec)
        else:  # leaf
            name = 'leaf{0}'.format(root['leaf_index'])
            label = 'leaf {0}: '.format(root['leaf_index'])
            label += '<B>{0}</B>'.format(_float2str(root['leaf_value'], precision))
            if 'leaf_weight' in show_info:
                label += '<br/>{0} weight'.format(_float2str(root['leaf_weight'], precision))
            if 'leaf_count' in show_info: