clustering_h_agglomerative.py 文件源码

def draw_binary_tree(Tree, correlation, c='none'):
    A = pg.AGraph(directed=True, strict=True)

    level = 0
    queue = [Tree.tree]
    while queue:
        node = queue.pop(0)
        #node_string = str(node.label) + '\n' + str(Tree.generate_bt_stats(node))
        node_string = str(node.ground_truth) + '\n' + str(Tree.generate_bt_stats(node))

        level += 1
        if node.get_left() is not None:
            queue.append(node.get_left())
            child_string = str(node.get_left().ground_truth) + '\n' + str(Tree.generate_bt_stats(node.get_left()))
            A.add_edge(node_string, child_string)
        if node.get_right() is not None:
            queue.append(node.get_right())
            child_string = str(node.get_right().ground_truth) + '\n' + str(Tree.generate_bt_stats(node.get_right()))
            A.add_edge(node_string, child_string)
        if level >= cfg.max_tree_size:
            break

    dot_path = util.generate_meta_path(Tree.system_name, 'dot', c)
    util.ensure_path_exists(dot_path)
    A.write('{}{} BT.dot'.format(dot_path, Tree.system_name))
    A.layout(prog='dot')
    A.draw('{}{} BT Score={}.png'.format(dot_path, Tree.system_name, correlation))
    print "[clustering] : Created n-ary tree at path {}.".format('{}{} BT Score={}.png'.format(dot_path,
                                                                                               Tree.system_name,
                                                                                               correlation))