Initial commit
This commit is contained in:
Zhongwei Li
214
skills/etetoolkit/scripts/quick_visualize.py
Executable file
214
skills/etetoolkit/scripts/quick_visualize.py
Executable file
@@ -0,0 +1,214 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Quick tree visualization script with common customization options.
|
||||
|
||||
Provides command-line interface for rapid tree visualization with
|
||||
customizable styles, layouts, and output formats.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
try:
|
||||
from ete3 import Tree, TreeStyle, NodeStyle
|
||||
except ImportError:
|
||||
print("Error: ete3 not installed. Install with: pip install ete3")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
def create_tree_style(args):
|
||||
"""Create TreeStyle based on arguments."""
|
||||
ts = TreeStyle()
|
||||
|
||||
# Basic display options
|
||||
ts.show_leaf_name = args.show_names
|
||||
ts.show_branch_length = args.show_lengths
|
||||
ts.show_branch_support = args.show_support
|
||||
ts.show_scale = args.show_scale
|
||||
|
||||
# Layout
|
||||
ts.mode = args.mode
|
||||
ts.rotation = args.rotation
|
||||
|
||||
# Circular tree options
|
||||
if args.mode == "c":
|
||||
ts.arc_start = args.arc_start
|
||||
ts.arc_span = args.arc_span
|
||||
|
||||
# Spacing
|
||||
ts.branch_vertical_margin = args.vertical_margin
|
||||
if args.scale_factor:
|
||||
ts.scale = args.scale_factor
|
||||
|
||||
# Title
|
||||
if args.title:
|
||||
from ete3 import TextFace
|
||||
title_face = TextFace(args.title, fsize=16, bold=True)
|
||||
ts.title.add_face(title_face, column=0)
|
||||
|
||||
return ts
|
||||
|
||||
|
||||
def apply_node_styling(tree, args):
|
||||
"""Apply styling to tree nodes."""
|
||||
for node in tree.traverse():
|
||||
nstyle = NodeStyle()
|
||||
|
||||
if node.is_leaf():
|
||||
# Leaf style
|
||||
nstyle["fgcolor"] = args.leaf_color
|
||||
nstyle["size"] = args.leaf_size
|
||||
else:
|
||||
# Internal node style
|
||||
nstyle["fgcolor"] = args.internal_color
|
||||
nstyle["size"] = args.internal_size
|
||||
|
||||
# Color by support if enabled
|
||||
if args.color_by_support and hasattr(node, 'support') and node.support:
|
||||
if node.support >= 0.9:
|
||||
nstyle["fgcolor"] = "darkgreen"
|
||||
elif node.support >= 0.7:
|
||||
nstyle["fgcolor"] = "orange"
|
||||
else:
|
||||
nstyle["fgcolor"] = "red"
|
||||
|
||||
node.set_style(nstyle)
|
||||
|
||||
|
||||
def visualize_tree(tree_file, output, args):
|
||||
"""Load tree, apply styles, and render."""
|
||||
try:
|
||||
tree = Tree(str(tree_file), format=args.format)
|
||||
except Exception as e:
|
||||
print(f"Error loading tree: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
# Apply styling
|
||||
apply_node_styling(tree, args)
|
||||
|
||||
# Create tree style
|
||||
ts = create_tree_style(args)
|
||||
|
||||
# Render
|
||||
try:
|
||||
# Determine output parameters based on format
|
||||
output_path = str(output)
|
||||
|
||||
render_args = {"tree_style": ts}
|
||||
|
||||
if args.width:
|
||||
render_args["w"] = args.width
|
||||
if args.height:
|
||||
render_args["h"] = args.height
|
||||
if args.units:
|
||||
render_args["units"] = args.units
|
||||
if args.dpi:
|
||||
render_args["dpi"] = args.dpi
|
||||
|
||||
tree.render(output_path, **render_args)
|
||||
print(f"Tree rendered successfully to: {output}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error rendering tree: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Quick tree visualization with ETE toolkit",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog="""
|
||||
Examples:
|
||||
# Basic visualization
|
||||
%(prog)s tree.nw output.pdf
|
||||
|
||||
# Circular tree
|
||||
%(prog)s tree.nw output.pdf --mode c
|
||||
|
||||
# Large tree with custom sizing
|
||||
%(prog)s tree.nw output.png --width 1200 --height 800 --units px --dpi 300
|
||||
|
||||
# Hide names, show support, color by support
|
||||
%(prog)s tree.nw output.pdf --no-names --show-support --color-by-support
|
||||
|
||||
# Custom title
|
||||
%(prog)s tree.nw output.pdf --title "Phylogenetic Tree of Species"
|
||||
|
||||
# Semicircular layout
|
||||
%(prog)s tree.nw output.pdf --mode c --arc-start -90 --arc-span 180
|
||||
"""
|
||||
)
|
||||
|
||||
parser.add_argument("input", help="Input tree file (Newick format)")
|
||||
parser.add_argument("output", help="Output image file (png, pdf, or svg)")
|
||||
|
||||
# Tree format
|
||||
parser.add_argument("--format", type=int, default=0,
|
||||
help="Newick format number (default: 0)")
|
||||
|
||||
# Display options
|
||||
display = parser.add_argument_group("Display options")
|
||||
display.add_argument("--no-names", dest="show_names", action="store_false",
|
||||
help="Don't show leaf names")
|
||||
display.add_argument("--show-lengths", action="store_true",
|
||||
help="Show branch lengths")
|
||||
display.add_argument("--show-support", action="store_true",
|
||||
help="Show support values")
|
||||
display.add_argument("--show-scale", action="store_true",
|
||||
help="Show scale bar")
|
||||
|
||||
# Layout options
|
||||
layout = parser.add_argument_group("Layout options")
|
||||
layout.add_argument("--mode", choices=["r", "c"], default="r",
|
||||
help="Tree mode: r=rectangular, c=circular (default: r)")
|
||||
layout.add_argument("--rotation", type=int, default=0,
|
||||
help="Tree rotation in degrees (default: 0)")
|
||||
layout.add_argument("--arc-start", type=int, default=0,
|
||||
help="Circular tree start angle (default: 0)")
|
||||
layout.add_argument("--arc-span", type=int, default=360,
|
||||
help="Circular tree arc span (default: 360)")
|
||||
|
||||
# Styling options
|
||||
styling = parser.add_argument_group("Styling options")
|
||||
styling.add_argument("--leaf-color", default="blue",
|
||||
help="Leaf node color (default: blue)")
|
||||
styling.add_argument("--leaf-size", type=int, default=6,
|
||||
help="Leaf node size (default: 6)")
|
||||
styling.add_argument("--internal-color", default="gray",
|
||||
help="Internal node color (default: gray)")
|
||||
styling.add_argument("--internal-size", type=int, default=4,
|
||||
help="Internal node size (default: 4)")
|
||||
styling.add_argument("--color-by-support", action="store_true",
|
||||
help="Color internal nodes by support value")
|
||||
|
||||
# Size and spacing
|
||||
size = parser.add_argument_group("Size and spacing")
|
||||
size.add_argument("--width", type=int, help="Output width")
|
||||
size.add_argument("--height", type=int, help="Output height")
|
||||
size.add_argument("--units", choices=["px", "mm", "in"],
|
||||
help="Size units (px, mm, in)")
|
||||
size.add_argument("--dpi", type=int, help="DPI for raster output")
|
||||
size.add_argument("--scale-factor", type=int,
|
||||
help="Branch length scale factor (pixels per unit)")
|
||||
size.add_argument("--vertical-margin", type=int, default=10,
|
||||
help="Vertical margin between branches (default: 10)")
|
||||
|
||||
# Other options
|
||||
parser.add_argument("--title", help="Tree title")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Validate output format
|
||||
output_path = Path(args.output)
|
||||
valid_extensions = {".png", ".pdf", ".svg"}
|
||||
if output_path.suffix.lower() not in valid_extensions:
|
||||
print(f"Error: Output must be PNG, PDF, or SVG file")
|
||||
sys.exit(1)
|
||||
|
||||
# Visualize
|
||||
visualize_tree(args.input, args.output, args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
229
skills/etetoolkit/scripts/tree_operations.py
Executable file
229
skills/etetoolkit/scripts/tree_operations.py
Executable file
@@ -0,0 +1,229 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Tree operations helper script for common ETE toolkit tasks.
|
||||
|
||||
Provides command-line interface for basic tree operations like:
|
||||
- Format conversion
|
||||
- Rooting (outgroup, midpoint)
|
||||
- Pruning
|
||||
- Basic statistics
|
||||
- ASCII visualization
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
try:
|
||||
from ete3 import Tree
|
||||
except ImportError:
|
||||
print("Error: ete3 not installed. Install with: pip install ete3")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
def load_tree(tree_file, format_num=0):
|
||||
"""Load tree from file."""
|
||||
try:
|
||||
return Tree(str(tree_file), format=format_num)
|
||||
except Exception as e:
|
||||
print(f"Error loading tree: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
def convert_format(tree_file, output, in_format=0, out_format=1):
|
||||
"""Convert tree between Newick formats."""
|
||||
tree = load_tree(tree_file, in_format)
|
||||
tree.write(outfile=str(output), format=out_format)
|
||||
print(f"Converted {tree_file} (format {in_format}) → {output} (format {out_format})")
|
||||
|
||||
|
||||
def reroot_tree(tree_file, output, outgroup=None, midpoint=False, format_num=0):
|
||||
"""Reroot tree by outgroup or midpoint."""
|
||||
tree = load_tree(tree_file, format_num)
|
||||
|
||||
if midpoint:
|
||||
midpoint_node = tree.get_midpoint_outgroup()
|
||||
tree.set_outgroup(midpoint_node)
|
||||
print(f"Rerooted tree using midpoint method")
|
||||
elif outgroup:
|
||||
try:
|
||||
outgroup_node = tree & outgroup
|
||||
tree.set_outgroup(outgroup_node)
|
||||
print(f"Rerooted tree using outgroup: {outgroup}")
|
||||
except Exception as e:
|
||||
print(f"Error: Could not find outgroup '{outgroup}': {e}")
|
||||
sys.exit(1)
|
||||
else:
|
||||
print("Error: Must specify either --outgroup or --midpoint")
|
||||
sys.exit(1)
|
||||
|
||||
tree.write(outfile=str(output), format=format_num)
|
||||
print(f"Saved rerooted tree to: {output}")
|
||||
|
||||
|
||||
def prune_tree(tree_file, output, keep_taxa, preserve_length=True, format_num=0):
|
||||
"""Prune tree to keep only specified taxa."""
|
||||
tree = load_tree(tree_file, format_num)
|
||||
|
||||
# Read taxa list
|
||||
taxa_file = Path(keep_taxa)
|
||||
if taxa_file.exists():
|
||||
with open(taxa_file) as f:
|
||||
taxa = [line.strip() for line in f if line.strip()]
|
||||
else:
|
||||
taxa = [t.strip() for t in keep_taxa.split(",")]
|
||||
|
||||
print(f"Pruning tree to {len(taxa)} taxa")
|
||||
|
||||
try:
|
||||
tree.prune(taxa, preserve_branch_length=preserve_length)
|
||||
tree.write(outfile=str(output), format=format_num)
|
||||
print(f"Pruned tree saved to: {output}")
|
||||
print(f"Retained {len(tree)} leaves")
|
||||
except Exception as e:
|
||||
print(f"Error pruning tree: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
def tree_stats(tree_file, format_num=0):
|
||||
"""Display tree statistics."""
|
||||
tree = load_tree(tree_file, format_num)
|
||||
|
||||
print(f"\n=== Tree Statistics ===")
|
||||
print(f"File: {tree_file}")
|
||||
print(f"Number of leaves: {len(tree)}")
|
||||
print(f"Total nodes: {len(list(tree.traverse()))}")
|
||||
|
||||
farthest_leaf, distance = tree.get_farthest_leaf()
|
||||
print(f"Tree depth: {distance:.4f}")
|
||||
print(f"Farthest leaf: {farthest_leaf.name}")
|
||||
|
||||
# Branch length statistics
|
||||
branch_lengths = [node.dist for node in tree.traverse() if not node.is_root()]
|
||||
if branch_lengths:
|
||||
print(f"\nBranch length statistics:")
|
||||
print(f" Mean: {sum(branch_lengths)/len(branch_lengths):.4f}")
|
||||
print(f" Min: {min(branch_lengths):.4f}")
|
||||
print(f" Max: {max(branch_lengths):.4f}")
|
||||
|
||||
# Support values
|
||||
supports = [node.support for node in tree.traverse() if not node.is_leaf() and hasattr(node, 'support')]
|
||||
if supports:
|
||||
print(f"\nSupport value statistics:")
|
||||
print(f" Mean: {sum(supports)/len(supports):.2f}")
|
||||
print(f" Min: {min(supports):.2f}")
|
||||
print(f" Max: {max(supports):.2f}")
|
||||
|
||||
print()
|
||||
|
||||
|
||||
def show_ascii(tree_file, format_num=0, show_internal=True):
|
||||
"""Display tree as ASCII art."""
|
||||
tree = load_tree(tree_file, format_num)
|
||||
print(tree.get_ascii(show_internal=show_internal))
|
||||
|
||||
|
||||
def list_leaves(tree_file, format_num=0):
|
||||
"""List all leaf names."""
|
||||
tree = load_tree(tree_file, format_num)
|
||||
for leaf in tree:
|
||||
print(leaf.name)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="ETE toolkit tree operations helper",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog="""
|
||||
Examples:
|
||||
# Convert format
|
||||
%(prog)s convert input.nw output.nw --in-format 0 --out-format 1
|
||||
|
||||
# Midpoint root
|
||||
%(prog)s reroot input.nw output.nw --midpoint
|
||||
|
||||
# Reroot with outgroup
|
||||
%(prog)s reroot input.nw output.nw --outgroup "Outgroup_species"
|
||||
|
||||
# Prune tree
|
||||
%(prog)s prune input.nw output.nw --keep-taxa "speciesA,speciesB,speciesC"
|
||||
|
||||
# Show statistics
|
||||
%(prog)s stats input.nw
|
||||
|
||||
# Display as ASCII
|
||||
%(prog)s ascii input.nw
|
||||
|
||||
# List all leaves
|
||||
%(prog)s leaves input.nw
|
||||
"""
|
||||
)
|
||||
|
||||
subparsers = parser.add_subparsers(dest="command", help="Command to execute")
|
||||
|
||||
# Convert command
|
||||
convert_parser = subparsers.add_parser("convert", help="Convert tree format")
|
||||
convert_parser.add_argument("input", help="Input tree file")
|
||||
convert_parser.add_argument("output", help="Output tree file")
|
||||
convert_parser.add_argument("--in-format", type=int, default=0, help="Input format (default: 0)")
|
||||
convert_parser.add_argument("--out-format", type=int, default=1, help="Output format (default: 1)")
|
||||
|
||||
# Reroot command
|
||||
reroot_parser = subparsers.add_parser("reroot", help="Reroot tree")
|
||||
reroot_parser.add_argument("input", help="Input tree file")
|
||||
reroot_parser.add_argument("output", help="Output tree file")
|
||||
reroot_parser.add_argument("--outgroup", help="Outgroup taxon name")
|
||||
reroot_parser.add_argument("--midpoint", action="store_true", help="Use midpoint rooting")
|
||||
reroot_parser.add_argument("--format", type=int, default=0, help="Newick format (default: 0)")
|
||||
|
||||
# Prune command
|
||||
prune_parser = subparsers.add_parser("prune", help="Prune tree to specified taxa")
|
||||
prune_parser.add_argument("input", help="Input tree file")
|
||||
prune_parser.add_argument("output", help="Output tree file")
|
||||
prune_parser.add_argument("--keep-taxa", required=True,
|
||||
help="Taxa to keep (comma-separated or file path)")
|
||||
prune_parser.add_argument("--no-preserve-length", action="store_true",
|
||||
help="Don't preserve branch lengths")
|
||||
prune_parser.add_argument("--format", type=int, default=0, help="Newick format (default: 0)")
|
||||
|
||||
# Stats command
|
||||
stats_parser = subparsers.add_parser("stats", help="Display tree statistics")
|
||||
stats_parser.add_argument("input", help="Input tree file")
|
||||
stats_parser.add_argument("--format", type=int, default=0, help="Newick format (default: 0)")
|
||||
|
||||
# ASCII command
|
||||
ascii_parser = subparsers.add_parser("ascii", help="Display tree as ASCII art")
|
||||
ascii_parser.add_argument("input", help="Input tree file")
|
||||
ascii_parser.add_argument("--format", type=int, default=0, help="Newick format (default: 0)")
|
||||
ascii_parser.add_argument("--no-internal", action="store_true",
|
||||
help="Don't show internal node names")
|
||||
|
||||
# Leaves command
|
||||
leaves_parser = subparsers.add_parser("leaves", help="List all leaf names")
|
||||
leaves_parser.add_argument("input", help="Input tree file")
|
||||
leaves_parser.add_argument("--format", type=int, default=0, help="Newick format (default: 0)")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if not args.command:
|
||||
parser.print_help()
|
||||
sys.exit(1)
|
||||
|
||||
# Execute command
|
||||
if args.command == "convert":
|
||||
convert_format(args.input, args.output, args.in_format, args.out_format)
|
||||
elif args.command == "reroot":
|
||||
reroot_tree(args.input, args.output, args.outgroup, args.midpoint, args.format)
|
||||
elif args.command == "prune":
|
||||
prune_tree(args.input, args.output, args.keep_taxa,
|
||||
not args.no_preserve_length, args.format)
|
||||
elif args.command == "stats":
|
||||
tree_stats(args.input, args.format)
|
||||
elif args.command == "ascii":
|
||||
show_ascii(args.input, args.format, not args.no_internal)
|
||||
elif args.command == "leaves":
|
||||
list_leaves(args.input, args.format)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user