对进化树着色

from ete3 import Tree, TreeStyle, NodeStyle
#安装模块： pip3 install  PyQt5 ete3
#如果在ubuntu系统上报错qt.qpa.plugin: Could not load the Qt platform plugin "xcb"，可安装：sudo apt-get install libxcb-xinerama0

# 加载进化树文件
tree = Tree("XXX.nwk")

# 创建一个 TreeStyle 对象
ts = TreeStyle()

# 定义着色规则
color_group1 = "red"  # -group1- 枝条的颜色
color_group2 = "yellow"  # -group2- 枝条的颜色
color_group3 = "purple" 

# 遍历树的每个节点
for node in tree.traverse():
    if "-group1-" in node.name:
        # 如果节点名包含 -group1-，则将其枝条着色为红色
        node_style = NodeStyle()
        node_style["fgcolor"] = color_group1
        node.set_style(node_style)
    elif "-group2-" in node.name:
        # 如果节点名包含 -group2-，则将其枝条着色为蓝色
        node_style = NodeStyle()
        node_style["bgcolor"] = color_group2
        node.set_style(node_style)
    elif "-group3-" in node.name:
        node_style = NodeStyle()
        node_style["bgcolor"] = color_group3
        node.set_style(node_style)


# 将进化树显示为圆形
'''
ts.mode = "c"
ts.arc_start = 0  # 起始角度
ts.arc_span = 360  # 弧度范围
'''
# 设置树样式
ts.show_leaf_name = True

# 输出着色后的进化树
#tree.show(tree_style=ts)

# 渲染进化树
tree.render("tree.png", w=800, dpi=800, tree_style=ts)