Skills › Scientific Agent Skills › etetoolkit › tree_operations.py
tree_operations.py
464 lines · skills/etetoolkit/scripts/tree_operations.py at 5f2bc03
1#!/usr/bin/env python3
2"""Validated command-line tree operations for ETE 4."""
4from __future__ import annotations
6import argparse
7import json
8import math
9import statistics
10import sys
11from collections import Counter
12from pathlib import Path
13from typing import Any
15try:
16 import ete4
17 from ete4 import Tree
18 from ete4.parser.newick import NewickError
19except ImportError as exc:
20 raise SystemExit(
21 'ETE 4 is required. Install it with: uv pip install "ete4==4.4.0"'
22 ) from exc
25ParserSpec = int | str
28class UserInputError(ValueError):
29 """An actionable problem with command-line input or tree content."""
32def parser_spec(value: str) -> ParserSpec:
33 """Parse numeric Newick parser IDs while retaining named aliases."""
34 text = value.strip()
35 try:
36 return int(text)
37 except ValueError:
38 if not text:
39 raise argparse.ArgumentTypeError("parser cannot be empty")
40 return text
43def comma_separated(value: str) -> list[str]:
44 """Parse a comma-separated property list."""
45 return [item.strip() for item in value.split(",") if item.strip()]
48def load_tree(path: Path, parser: ParserSpec) -> Tree:
49 """Load one Newick tree from a UTF-8 text file."""
50 if not path.is_file():
51 raise UserInputError(f"input tree does not exist or is not a file: {path}")
53 try:
54 with path.open(encoding="utf-8") as handle:
55 tree = Tree(handle, parser=parser)
56 for node in tree.traverse():
57 for prop in ("dist", "support"):
58 value = node.props.get(prop)
59 if value is not None and not math.isfinite(float(value)):
60 raise ValueError(f"nonfinite {prop} on node {node.name!r}")
61 return tree
62 except (OSError, ValueError, TypeError, NewickError) as exc:
63 raise UserInputError(
64 f"could not parse {path} with Newick parser {parser!r}: {exc}"
65 ) from exc
68def save_tree(
69 tree: Tree,
70 path: Path,
71 parser: ParserSpec,
72 props: list[str],
73) -> None:
74 """Serialize a tree after validating its destination directory."""
75 if not path.parent.is_dir():
76 raise UserInputError(f"output directory does not exist: {path.parent}")
78 try:
79 newick = tree.write(
80 parser=parser,
81 props=props,
82 format_root_node=True,
83 )
84 path.write_text(newick.rstrip("\n") + "\n", encoding="utf-8")
85 except (OSError, ValueError, TypeError, NewickError) as exc:
86 raise UserInputError(
87 f"could not write {path} with Newick parser {parser!r}: {exc}"
88 ) from exc
91def numeric_summary(values: list[float]) -> dict[str, float] | None:
92 """Return basic descriptive statistics for a numeric list."""
93 if not values:
94 return None
95 return {
96 "minimum": min(values),
97 "maximum": max(values),
98 "mean": statistics.fmean(values),
99 "median": statistics.median(values),
100 }
103def tree_stats(tree: Tree, source: Path) -> dict[str, Any]:
104 """Calculate structural, branch-length, and support diagnostics."""
105 nodes = list(tree.traverse())
106 leaves = list(tree.leaves())
107 internal = [node for node in nodes if not node.is_leaf]
108 names = [leaf.name for leaf in leaves]
109 duplicate_names = sorted(
110 name for name, count in Counter(names).items() if name is not None and count > 1
111 )
112 unnamed_leaf_ids = [list(leaf.id) for leaf in leaves if not leaf.name]
114 branch_lengths = [
115 float(node.dist) for node in nodes if not node.is_root and node.dist is not None
116 ]
117 support_values = [
118 float(node.support)
119 for node in internal
120 if not node.is_root and node.support is not None
121 ]
123 farthest_leaf, farthest_distance = tree.get_farthest_leaf()
125 return {
126 "source": str(source),
127 "ete_version": ete4.__version__,
128 "leaf_count": len(leaves),
129 "internal_node_count": len(internal),
130 "total_node_count": len(nodes),
131 "root_child_count": len(tree.children),
132 "polytomy_count": sum(len(node.children) > 2 for node in internal),
133 "unary_node_count": sum(len(node.children) == 1 for node in internal),
134 "duplicate_leaf_names": duplicate_names,
135 "unnamed_leaf_ids": unnamed_leaf_ids,
136 "branch_lengths": numeric_summary(branch_lengths),
137 "internal_support": numeric_summary(support_values),
138 "farthest_leaf": farthest_leaf.name,
139 "farthest_leaf_distance": float(farthest_distance),
140 }
143def print_stats(stats: dict[str, Any], as_json: bool) -> None:
144 """Print statistics as JSON or readable text."""
145 if as_json:
146 print(json.dumps(stats, indent=2, sort_keys=True))
147 return
149 print(f"File: {stats['source']}")
150 print(f"ETE version: {stats['ete_version']}")
151 print(f"Leaves: {stats['leaf_count']}")
152 print(f"Internal nodes: {stats['internal_node_count']}")
153 print(f"Total nodes: {stats['total_node_count']}")
154 print(f"Root children: {stats['root_child_count']}")
155 print(f"Polytomies: {stats['polytomy_count']}")
156 print(f"Unary nodes: {stats['unary_node_count']}")
157 print(f"Farthest leaf: {stats['farthest_leaf']!r}")
158 print(f"Farthest distance: {stats['farthest_leaf_distance']:.6g}")
160 for label, key in (
161 ("Branch lengths", "branch_lengths"),
162 ("Internal support", "internal_support"),
163 ):
164 summary = stats[key]
165 if summary is None:
166 print(f"{label}: none")
167 else:
168 print(
169 f"{label}: min={summary['minimum']:.6g}, "
170 f"median={summary['median']:.6g}, "
171 f"mean={summary['mean']:.6g}, "
172 f"max={summary['maximum']:.6g}"
173 )
175 print(f"Duplicate leaf names: {stats['duplicate_leaf_names'] or 'none'}")
176 print(f"Unnamed leaf IDs: {stats['unnamed_leaf_ids'] or 'none'}")
179def resolve_unique_node(tree: Tree, name: str) -> Tree:
180 """Resolve exactly one named node."""
181 matches = list(tree.search_nodes(name=name))
182 if not matches:
183 raise UserInputError(f"node not found: {name!r}")
184 if len(matches) > 1:
185 raise UserInputError(
186 f"node name is ambiguous ({len(matches)} matches): {name!r}"
187 )
188 return matches[0]
191def read_keep_names(values: list[str] | None, file_path: Path | None) -> list[str]:
192 """Read requested names from arguments or a one-name-per-line file."""
193 if file_path is not None:
194 if not file_path.is_file():
195 raise UserInputError(f"taxon file does not exist: {file_path}")
196 try:
197 names = [
198 line.strip()
199 for line in file_path.read_text(encoding="utf-8").splitlines()
200 if line.strip() and not line.lstrip().startswith("#")
201 ]
202 except OSError as exc:
203 raise UserInputError(
204 f"could not read taxon file {file_path}: {exc}"
205 ) from exc
206 else:
207 names = values or []
209 duplicates = sorted(name for name, count in Counter(names).items() if count > 1)
210 if not names:
211 raise UserInputError("at least one taxon must be requested")
212 if duplicates:
213 raise UserInputError(f"duplicate requested taxa: {duplicates}")
214 return names
217def validate_requested_names(tree: Tree, requested: list[str]) -> None:
218 """Require every requested name to resolve to exactly one tree node."""
219 counts = Counter(tree.leaf_names())
220 missing = sorted(name for name in requested if counts[name] == 0)
221 ambiguous = sorted(name for name in requested if counts[name] > 1)
222 if missing:
223 raise UserInputError(f"requested leaves are absent: {missing}")
224 if ambiguous:
225 raise UserInputError(f"requested leaf names are duplicated: {ambiguous}")
228def command_stats(args: argparse.Namespace) -> None:
229 tree = load_tree(args.input, args.parser)
230 print_stats(tree_stats(tree, args.input), args.json)
233def command_ascii(args: argparse.Namespace) -> None:
234 tree = load_tree(args.input, args.parser)
235 print(
236 tree.to_str(
237 show_internal=not args.no_internal,
238 compact=args.compact,
239 props=args.props,
240 )
241 )
244def command_leaves(args: argparse.Namespace) -> None:
245 tree = load_tree(args.input, args.parser)
246 for name in tree.leaf_names():
247 print("" if name is None else name)
250def command_convert(args: argparse.Namespace) -> None:
251 tree = load_tree(args.input, args.input_parser)
252 save_tree(tree, args.output, args.output_parser, args.props)
253 print(f"Wrote {args.output}")
256def command_reroot(args: argparse.Namespace) -> None:
257 tree = load_tree(args.input, args.parser)
258 if args.midpoint:
259 tree.set_midpoint_outgroup(topological=args.topological)
260 method = (
261 "topological midpoint" if args.topological else "branch-length midpoint"
262 )
263 else:
264 if args.topological:
265 raise UserInputError("--topological applies only to --midpoint")
266 outgroup = resolve_unique_node(tree, args.outgroup)
267 tree.set_outgroup(outgroup)
268 method = f"outgroup {args.outgroup!r}"
270 output_parser = args.parser if args.output_parser is None else args.output_parser
271 save_tree(tree, args.output, output_parser, args.props)
272 print(f"Rerooted with {method}; wrote {args.output}")
275def command_prune(args: argparse.Namespace) -> None:
276 tree = load_tree(args.input, args.parser)
277 names = read_keep_names(args.keep, args.keep_file)
278 validate_requested_names(tree, names)
279 # Resolve leaves, not arbitrary nodes sharing a leaf's name.
280 leaves_by_name = {leaf.name: leaf for leaf in tree.leaves()}
281 tree.prune(
282 [leaves_by_name[name] for name in names],
283 preserve_branch_length=args.preserve_branch_length,
284 )
285 output_parser = args.parser if args.output_parser is None else args.output_parser
286 save_tree(tree, args.output, output_parser, args.props)
287 print(f"Retained {len(names)} leaves; wrote {args.output}")
290def command_compare(args: argparse.Namespace) -> None:
291 tree_a = load_tree(args.tree_a, args.parser_a)
292 tree_b = load_tree(args.tree_b, args.parser_b)
294 for label, tree in (("tree_a", tree_a), ("tree_b", tree_b)):
295 names = list(tree.leaf_names())
296 unnamed_count = sum(not name for name in names)
297 if unnamed_count:
298 raise UserInputError(f"{label} has {unnamed_count} unnamed leaves")
299 counts = Counter(names)
300 duplicates = sorted(
301 name for name, count in counts.items() if name is not None and count > 1
302 )
303 if duplicates:
304 raise UserInputError(f"{label} has duplicate leaf names: {duplicates}")
306 if len(set(tree_a.leaf_names()) & set(tree_b.leaf_names())) < 2:
307 raise UserInputError("RF comparison requires at least two shared leaves")
309 (
310 rf,
311 max_rf,
312 common,
313 _edges_a,
314 _edges_b,
315 discarded_a,
316 discarded_b,
317 ) = tree_a.robinson_foulds(
318 tree_b,
319 unrooted_trees=args.unrooted,
320 min_support_t1=args.min_support_a,
321 min_support_t2=args.min_support_b,
322 )
324 result = {
325 "tree_a": str(args.tree_a),
326 "tree_b": str(args.tree_b),
327 "unrooted": args.unrooted,
328 "rf": rf,
329 "max_rf": max_rf,
330 # With no eligible splits there is no denominator, not perfect agreement.
331 "normalized_rf": rf / max_rf if max_rf > 0 else None,
332 "common_leaf_count": len(common),
333 "common_leaves": sorted(common),
334 "discarded_edge_count_a": len(discarded_a),
335 "discarded_edge_count_b": len(discarded_b),
336 }
337 print(json.dumps(result, indent=2, sort_keys=True))
340def build_parser() -> argparse.ArgumentParser:
341 parser = argparse.ArgumentParser(
342 description="Validated ETE 4 tree operations",
343 )
344 parser.add_argument(
345 "--version",
346 action="version",
347 version=f"%(prog)s (ETE {ete4.__version__})",
348 )
349 subparsers = parser.add_subparsers(dest="command", required=True)
351 stats = subparsers.add_parser("stats", help="report tree diagnostics")
352 stats.add_argument("input", type=Path)
353 stats.add_argument("--parser", type=parser_spec, default=0)
354 stats.add_argument("--json", action="store_true")
355 stats.set_defaults(handler=command_stats)
357 ascii_parser = subparsers.add_parser("ascii", help="print a terminal tree")
358 ascii_parser.add_argument("input", type=Path)
359 ascii_parser.add_argument("--parser", type=parser_spec, default=0)
360 ascii_parser.add_argument(
361 "--props",
362 type=comma_separated,
363 default=["name"],
364 help="comma-separated node properties to display (default: name)",
365 )
366 ascii_parser.add_argument("--compact", action="store_true")
367 ascii_parser.add_argument("--no-internal", action="store_true")
368 ascii_parser.set_defaults(handler=command_ascii)
370 leaves = subparsers.add_parser("leaves", help="print leaf names")
371 leaves.add_argument("input", type=Path)
372 leaves.add_argument("--parser", type=parser_spec, default=0)
373 leaves.set_defaults(handler=command_leaves)
375 convert = subparsers.add_parser("convert", help="convert Newick parsers")
376 convert.add_argument("input", type=Path)
377 convert.add_argument("output", type=Path)
378 convert.add_argument("--input-parser", type=parser_spec, default=0)
379 convert.add_argument("--output-parser", type=parser_spec, default=1)
380 convert.add_argument(
381 "--props",
382 type=comma_separated,
383 default=[],
384 help="comma-separated NHX properties to retain (default: none)",
385 )
386 convert.set_defaults(handler=command_convert)
388 reroot = subparsers.add_parser("reroot", help="reroot by outgroup or midpoint")
389 reroot.add_argument("input", type=Path)
390 reroot.add_argument("output", type=Path)
391 reroot.add_argument("--parser", type=parser_spec, default=0)
392 reroot.add_argument("--output-parser", type=parser_spec)
393 rooting = reroot.add_mutually_exclusive_group(required=True)
394 rooting.add_argument("--outgroup", help="unique node name to use as outgroup")
395 rooting.add_argument("--midpoint", action="store_true")
396 reroot.add_argument(
397 "--topological",
398 action="store_true",
399 help="use edge counts instead of branch lengths for midpoint rooting",
400 )
401 reroot.add_argument(
402 "--props",
403 type=comma_separated,
404 default=[],
405 help="comma-separated NHX properties to write (default: none)",
406 )
407 reroot.set_defaults(handler=command_reroot)
409 prune = subparsers.add_parser("prune", help="retain selected leaves")
410 prune.add_argument("input", type=Path)
411 prune.add_argument("output", type=Path)
412 prune.add_argument("--parser", type=parser_spec, default=0)
413 prune.add_argument("--output-parser", type=parser_spec)
414 selection = prune.add_mutually_exclusive_group(required=True)
415 selection.add_argument("--keep", nargs="+", help="leaf names to retain")
416 selection.add_argument(
417 "--keep-file",
418 type=Path,
419 help="UTF-8 file with one leaf name per line",
420 )
421 prune.add_argument(
422 "--preserve-branch-length",
423 action=argparse.BooleanOptionalAction,
424 default=True,
425 )
426 prune.add_argument(
427 "--props",
428 type=comma_separated,
429 default=[],
430 help="comma-separated NHX properties to write (default: none)",
431 )
432 prune.set_defaults(handler=command_prune)
434 compare = subparsers.add_parser(
435 "compare",
436 help="calculate Robinson-Foulds distance",
437 )
438 compare.add_argument("tree_a", type=Path)
439 compare.add_argument("tree_b", type=Path)
440 compare.add_argument("--parser-a", type=parser_spec, default=0)
441 compare.add_argument("--parser-b", type=parser_spec, default=0)
442 compare.add_argument("--unrooted", action="store_true")
443 compare.add_argument("--min-support-a", type=float, default=0.0)
444 compare.add_argument("--min-support-b", type=float, default=0.0)
445 compare.set_defaults(handler=command_compare)
447 return parser
450def main(argv: list[str] | None = None) -> int:
451 args = build_parser().parse_args(argv)
452 try:
453 args.handler(args)
454 except UserInputError as exc:
455 print(f"error: {exc}", file=sys.stderr)
456 return 2
457 except Exception as exc: # Preserve unexpected ETE failures for CLI users.
458 print(f"unexpected ETE error: {type(exc).__name__}: {exc}", file=sys.stderr)
459 return 1
460 return 0
463if __name__ == "__main__":
464 raise SystemExit(main())