Преобразование дерева решений spark mllib в graphviz

Question

У меня есть дерево решений, построенное с помощью Apache Spark MLLIB.Выходной tree.txt с использованием myDecisionTreeModel.toDebugString () выглядит следующим образом:

If (feature 36 <= 0.0)
 If (feature 35 <= 5.0)
  If (feature 42 <= 61.0)
   If (feature 0 <= 3732128.0)
    If (feature 23 <= 2.0)
     Predict: 1.2779154046107128E-4
    Else (feature 23 > 2.0)
     Predict: 3.5523837168253053E-4
...
Else (feature 36 > 0.0)
 If (feature 23 <= 2.0)
  If (feature 41 <= 5.0)
etc.

Итак, я реализовал следующий код Python для его рендеринга с использованием graphviz:

import graphviz
from graphviz import Digraph

dot = Digraph(comment='decision tree')

def parse(from_node, subtree):
    if len(subtree) > 0:
        root = subtree.pop(0).rstrip("\n") 
        if "Predict" in root:
            dot.node(from_node, root)
        else:
            dot.node(from_node, root)
            dot.edge(from_node, str(2*int(from_node) + 2))
            dot.edge(from_node, str(2*int(from_node) + 1))
            #split the tree into two halves
            idx = 0
            tgt = len(root)-len(root.lstrip())
            for l in subtree:
                if (len(l)-len(l.lstrip()) == tgt):
                    break
                idx += 1
            subtreeRight = subtree[:idx]
            subtreeLeft = subtree[idx+1:]

            parse(str(2*int(from_node) + 1), subtreeRight)
            parse(str(2*int(from_node) + 2), subtreeLeft)


with open("tree.txt", "r") as f:
    lines = f.readlines()
    parse("0", lines)

Этовероятно, очень наивно, но я не знаю, как сделать это более эффективно: в основном я считаю количество начальных пробелов для разбиения моего дерева с учетом заданного условия.

Вывод выглядит нормально, но этоТрудно предсказать, будут ли правильно расположены правые и левые дети.Если я инвертирую две строки:

dot.edge(from_node, str(2*int(from_node) + 2))
dot.edge(from_node, str(2*int(from_node) + 1))

, кажется, что правый ребенок слева, а левый справа.Кроме того, дерево не очень хорошо сбалансировано.

Есть идеи, чтобы сделать его более похожим на дерево решений, и чтобы быть уверенным в порядке?

спасибо