Function print_tree_graphviz

visualize_mcts.py:311–432

使用 Graphviz 生成带箭头的树形图 需要安装: pip install graphviz 系统需要安装 Graphviz: https://graphviz.org/download/

(nodes: Dict[str, NodeInfo], root_id: str, output_file: str = "mcts_tree")

def print_tree_graphviz(nodes: Dict[str, NodeInfo], root_id: str, output_file: str = "mcts_tree"):
    """
    使用 Graphviz 生成带箭头的树形图
    需要安装: pip install graphviz
    系统需要安装 Graphviz: https://graphviz.org/download/
    """
    try:
        from graphviz import Digraph
    except ImportError:
        print("Error: graphviz not installed. Install with: pip install graphviz")
        print("Also ensure Graphviz executables are installed on your system.")
        print("  Ubuntu/Debian: sudo apt-get install graphviz")
        print("  macOS: brew install graphviz")
        print("  Windows: Download from https://graphviz.org/download/")
        return
    
    # 创建有向图
    g = Digraph(
        name='MCTS_Tree',
        comment='MCTS Search Tree Visualization',
        format='png',  # 可以改为 'pdf', 'svg' 等
        graph_attr={
            'rankdir': 'TB',  # Top to Bottom (从上到下)
            'splines': 'ortho',  # 正交边
            'nodesep': '0.5',
            'ranksep': '1.0',
            'bgcolor': '#f5f5f5'
        },
        node_attr={
            'shape': 'box',
            'style': 'rounded,filled',
            'fontsize': '10',
            'fontname': 'Arial'
        },
        edge_attr={
            'arrowhead': 'vee',
            'color': '#666666',
            'penwidth': '1.5'
        }
    )
    
    # 添加节点和边的递归函数
    def add_node_recursive(node_id: str):
        if node_id not in nodes:
            return
        
        node = nodes[node_id]
        
        # 构建节点标签
        status_sym = "✓" if node.run_status == "OK" else "✗" if node.run_status == "FAIL" else "?"
        stage_name = {"root": "R", "draft": "D", "improve": "I", "debug": "DB"}.get(node.stage, node.stage[:1].upper())
        
        label_lines = [
            f"{status_sym} {stage_name} {node.id[:8]}",
        ]
        
        if node.metric is not None:
            label_lines.append(f"metric: {node.metric:.4f}")