Recursively add node or edge.
(root, total_count, parent=None, decision=None)
| 386 | raise ImportError('You must install graphviz to plot tree.') |
| 387 | |
| 388 | def add(root, total_count, parent=None, decision=None): |
| 389 | """Recursively add node or edge.""" |
| 390 | if 'split_index' in root: # non-leaf |
| 391 | l_dec = 'yes' |
| 392 | r_dec = 'no' |
| 393 | if root['decision_type'] == '<=': |
| 394 | lte_symbol = "≤" |
| 395 | operator = lte_symbol |
| 396 | elif root['decision_type'] == '==': |
| 397 | operator = "=" |
| 398 | else: |
| 399 | raise ValueError('Invalid decision type in tree model.') |
| 400 | name = 'split{0}'.format(root['split_index']) |
| 401 | if feature_names is not None: |
| 402 | label = '<B>{0}</B> {1} '.format(feature_names[root['split_feature']], operator) |
| 403 | else: |
| 404 | label = 'feature <B>{0}</B> {1} '.format(root['split_feature'], operator) |
| 405 | label += '<B>{0}</B>'.format(_float2str(root['threshold'], precision)) |
| 406 | for info in ['split_gain', 'internal_value', 'internal_weight', "internal_count", "data_percentage"]: |
| 407 | if info in show_info: |
| 408 | output = info.split('_')[-1] |
| 409 | if info in {'split_gain', 'internal_value', 'internal_weight'}: |
| 410 | label += '<br/>{0} {1}'.format(_float2str(root[info], precision), output) |
| 411 | elif info == 'internal_count': |
| 412 | label += '<br/>{0}: {1}'.format(output, root[info]) |
| 413 | elif info == "data_percentage": |
| 414 | label += '<br/>{0}% of data'.format(_float2str(root['internal_count'] / total_count * 100, 2)) |
| 415 | |
| 416 | fillcolor = "white" |
| 417 | style = "" |
| 418 | if constraints: |
| 419 | if constraints[root['split_feature']] == 1: |
| 420 | fillcolor = "#ddffdd" # light green |
| 421 | if constraints[root['split_feature']] == -1: |
| 422 | fillcolor = "#ffdddd" # light red |
| 423 | style = "filled" |
| 424 | label = "<" + label + ">" |
| 425 | graph.node(name, label=label, shape="rectangle", style=style, fillcolor=fillcolor) |
| 426 | add(root['left_child'], total_count, name, l_dec) |
| 427 | add(root['right_child'], total_count, name, r_dec) |
| 428 | else: # leaf |
| 429 | name = 'leaf{0}'.format(root['leaf_index']) |
| 430 | label = 'leaf {0}: '.format(root['leaf_index']) |
| 431 | label += '<B>{0}</B>'.format(_float2str(root['leaf_value'], precision)) |
| 432 | if 'leaf_weight' in show_info: |
| 433 | label += '<br/>{0} weight'.format(_float2str(root['leaf_weight'], precision)) |
| 434 | if 'leaf_count' in show_info: |
| 435 | label += '<br/>count: {0}'.format(root['leaf_count']) |
| 436 | if "data_percentage" in show_info: |
| 437 | label += '<br/>{0}% of data'.format(_float2str(root['leaf_count'] / total_count * 100, 2)) |
| 438 | label = "<" + label + ">" |
| 439 | graph.node(name, label=label) |
| 440 | if parent is not None: |
| 441 | graph.edge(parent, name, decision) |
| 442 | |
| 443 | graph = Digraph(**kwargs) |
| 444 | graph.attr("graph", nodesep="0.05", ranksep="0.3", rankdir="LR") |
no test coverage detected