forked from admin/PlantAI-rag-new
61 lines
1.7 KiB
Python
Executable File
61 lines
1.7 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""
|
|
使用 matplotlib + Bio.Phylo 绘制进化树 (无头模式)
|
|
"""
|
|
import sys
|
|
import os
|
|
os.environ['QT_QPA_PLATFORM'] = 'offscreen' # 无头模式
|
|
os.environ['MPLBACKEND'] = 'Agg' # 使用Agg后端
|
|
|
|
from pathlib import Path
|
|
import matplotlib
|
|
matplotlib.use('Agg') # 无头后端
|
|
import matplotlib.pyplot as plt
|
|
from Bio import Phylo
|
|
|
|
def plot_tree(treefile, output_dir):
|
|
"""绘制进化树并输出PDF和PNG"""
|
|
treefile = Path(treefile)
|
|
output_dir = Path(output_dir)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# 读取树 (Newick格式)
|
|
tree = Phylo.read(str(treefile), "newick")
|
|
|
|
# 输出文件名
|
|
base_name = treefile.stem.replace('.iqtree', '').replace('.treefile', '')
|
|
pdf_file = output_dir / f"{base_name}.pdf"
|
|
png_file = output_dir / f"{base_name}.png"
|
|
|
|
# 创建图形
|
|
fig, ax = plt.subplots(figsize=(14, 10))
|
|
|
|
# 绘制进化树
|
|
Phylo.draw(tree, axes=ax, do_show=False,
|
|
branch_labels=lambda c: c.confidence if c.confidence else None)
|
|
|
|
# 设置标题
|
|
ax.set_title("Interspecies Phylogenetic Tree", fontsize=14, fontweight='bold')
|
|
|
|
# 保存PDF
|
|
plt.savefig(str(pdf_file), format='pdf', bbox_inches='tight', dpi=300)
|
|
print(f"PDF: {pdf_file}")
|
|
|
|
# 保存PNG
|
|
plt.savefig(str(png_file), format='png', bbox_inches='tight', dpi=300)
|
|
print(f"PNG: {png_file}")
|
|
|
|
plt.close()
|
|
print("完成")
|
|
return True
|
|
|
|
if __name__ == "__main__":
|
|
if len(sys.argv) < 2:
|
|
print("用法: python plot_tree.py <treefile> [output_dir]")
|
|
sys.exit(1)
|
|
|
|
treefile = sys.argv[1]
|
|
output_dir = sys.argv[2] if len(sys.argv) >= 3 else Path(treefile).parent
|
|
|
|
plot_tree(treefile, output_dir)
|