Skills › Scientific Agent Skills › etetoolkit › quick_visualize.py
quick_visualize.py
485 lines · skills/etetoolkit/scripts/quick_visualize.py at 5f2bc03
1#!/usr/bin/env python3
2"""Interactive or static ETE 4 tree visualization."""
4from __future__ import annotations
6import argparse
7import ipaddress
8import math
9import sys
10import tempfile
11from pathlib import Path
13try:
14 import ete4
15 from ete4 import Tree
16 from ete4.parser.newick import NewickError
17except ImportError as exc:
18 raise SystemExit(
19 'ETE 4 is required. Install it with: uv pip install "ete4==4.4.0"'
20 ) from exc
23ParserSpec = int | str
26class UserInputError(ValueError):
27 """An actionable problem with command-line input."""
30def parser_spec(value: str) -> ParserSpec:
31 """Parse numeric Newick parser IDs while retaining named aliases."""
32 text = value.strip()
33 try:
34 return int(text)
35 except ValueError:
36 if not text:
37 raise argparse.ArgumentTypeError("parser cannot be empty")
38 return text
41def mode_spec(value: str) -> str:
42 """Normalize short and long layout mode names."""
43 aliases = {
44 "r": "rectangular",
45 "rectangular": "rectangular",
46 "c": "circular",
47 "circular": "circular",
48 }
49 try:
50 return aliases[value.lower()]
51 except KeyError as exc:
52 raise argparse.ArgumentTypeError(
53 "mode must be rectangular/r or circular/c"
54 ) from exc
57def load_tree(path: Path, parser: ParserSpec) -> Tree:
58 """Load a Newick tree from a UTF-8 file."""
59 if not path.is_file():
60 raise UserInputError(f"input tree does not exist or is not a file: {path}")
61 try:
62 with path.open(encoding="utf-8") as handle:
63 tree = Tree(handle, parser=parser)
64 for node in tree.traverse():
65 for prop in ("dist", "support"):
66 value = node.props.get(prop)
67 if value is not None and not math.isfinite(float(value)):
68 raise ValueError(f"nonfinite {prop} on node {node.name!r}")
69 return tree
70 except (OSError, ValueError, TypeError, NewickError) as exc:
71 raise UserInputError(
72 f"could not parse {path} with Newick parser {parser!r}: {exc}"
73 ) from exc
76def support_fraction(value: float | None, scale: str) -> float | None:
77 """Normalize a declared support convention; never guess per node."""
78 if value is None:
79 return None
80 if scale not in {"fraction", "percent"}:
81 raise UserInputError("support coloring requires --support-scale fraction or percent")
82 numeric = float(value)
83 maximum = 1 if scale == "fraction" else 100
84 if not math.isfinite(numeric) or not 0 <= numeric <= maximum:
85 raise UserInputError(f"support {value!r} is outside the {scale} range 0..{maximum}")
86 return numeric / maximum
89def support_color(node: Tree, args: argparse.Namespace) -> str:
90 """Map support to a color after normalization."""
91 support = support_fraction(node.support, args.support_scale)
92 if support is None:
93 return args.missing_support_color
94 if support >= args.high_support:
95 return args.high_support_color
96 if support >= args.moderate_support:
97 return args.moderate_support_color
98 return args.low_support_color
101def create_smartview_layout(args: argparse.Namespace):
102 """Create a current SmartView layout."""
103 try:
104 from ete4.smartview import Layout, PropFace, TextFace
105 except ImportError as exc:
106 raise UserInputError(
107 'SmartView is unavailable; reinstall with: uv pip install "ete4==4.4.0"'
108 ) from exc
110 def draw_tree(_tree):
111 tree_style = {
112 "shape": args.mode,
113 "node-height-min": args.collapse_pixels,
114 "content-height-min": args.content_pixels,
115 "show-popup-props": ["name", "dist", "support"],
116 }
117 if args.mode == "circular":
118 tree_style.update(
119 {
120 "angle-start": args.arc_start,
121 "angle-span": args.arc_span,
122 }
123 )
124 yield tree_style
125 if args.title:
126 yield TextFace(
127 args.title,
128 fs_min=8,
129 fs_max=22,
130 position="header",
131 )
133 def draw_node(node):
134 fill = (
135 support_color(node, args)
136 if args.color_by_support and not node.is_leaf
137 else args.leaf_color
138 if node.is_leaf
139 else args.internal_color
140 )
141 radius = args.leaf_size if node.is_leaf else args.internal_size
142 yield {
143 "dot": {
144 "shape": "circle",
145 "radius": radius,
146 "fill": fill,
147 }
148 }
150 if node.is_leaf and args.show_names:
151 yield PropFace(
152 "name",
153 fs_min=4,
154 fs_max=args.label_size,
155 position="right",
156 )
158 if args.show_support and not node.is_leaf and node.support is not None:
159 yield TextFace(
160 f"{node.support:g}",
161 fs_min=3,
162 fs_max=args.label_size,
163 style={"fill": "#555555"},
164 position="top",
165 )
167 if args.show_lengths and not node.is_root and node.dist is not None:
168 yield TextFace(
169 f"{node.dist:g}",
170 fs_min=3,
171 fs_max=args.label_size,
172 style={"fill": "#777777"},
173 position="bottom",
174 )
176 return Layout(
177 "quick visualization",
178 draw_tree=draw_tree,
179 draw_node=draw_node,
180 )
183def validate_bind_address(host: str, allow_remote: bool) -> None:
184 """Require explicit consent before binding beyond loopback."""
185 try:
186 address = ipaddress.ip_address(host)
187 except ValueError:
188 if host.lower() == "localhost":
189 return
190 if not allow_remote:
191 raise UserInputError("non-loopback host names require --allow-remote-bind")
192 return
194 if not address.is_loopback and not allow_remote:
195 raise UserInputError(
196 "refusing a non-loopback SmartView bind without --allow-remote-bind"
197 )
200def run_interactive_smartview(
201 tree: Tree,
202 layout,
203 args: argparse.Namespace,
204) -> None:
205 """Run the SmartView server until the user exits."""
206 validate_bind_address(args.host, args.allow_remote_bind)
207 try:
208 from ete4.smartview import explorer
209 except ImportError as exc:
210 raise UserInputError("could not import the SmartView explorer") from exc
212 tree.explore(
213 layouts=[layout],
214 host=args.host,
215 port=args.port,
216 open_browser=not args.no_browser,
217 )
218 print("SmartView is running. Press Enter or Ctrl-C to stop it.")
219 try:
220 input()
221 except (EOFError, KeyboardInterrupt):
222 pass
223 finally:
224 explorer.stop_server()
227def render_smartview(
228 tree: Tree,
229 layout,
230 output: Path,
231 args: argparse.Namespace,
232) -> None:
233 """Render a SmartView PNG screenshot."""
234 if output.suffix.lower() != ".png":
235 raise UserInputError(
236 "SmartView static output is PNG screenshot data; use a .png path "
237 "or select --engine treeview for PDF/SVG"
238 )
239 if not output.parent.is_dir():
240 raise UserInputError(f"output directory does not exist: {output.parent}")
242 try:
243 # ETE 4.4 warns and returns without output when selenium is absent.
244 from selenium import webdriver # noqa: F401
245 except (ImportError, ModuleNotFoundError) as exc:
246 raise UserInputError(
247 "SmartView static rendering needs the render-sm extra: "
248 'uv pip install "ete4[render-sm]==4.4.0"'
249 ) from exc
250 with tempfile.TemporaryDirectory(prefix="ete-render-", dir=output.parent) as temp:
251 temporary_output = Path(temp) / "tree.png"
252 tree.render_sm(
253 str(temporary_output), layouts=[layout], w=args.width, h=args.height,
254 )
255 if not temporary_output.is_file():
256 raise UserInputError("SmartView did not create a PNG output")
257 with temporary_output.open("rb") as handle:
258 if handle.read(8) != b"\x89PNG\r\n\x1a\n":
259 raise UserInputError("SmartView output is not a PNG image")
260 temporary_output.replace(output)
261 print(f"Wrote SmartView PNG: {output}")
264def create_treeview_style(tree: Tree, args: argparse.Namespace):
265 """Create a Qt treeview style and apply node styles."""
266 try:
267 from ete4.treeview import NodeStyle, TextFace, TreeStyle
268 except ImportError as exc:
269 raise UserInputError(
270 "Qt treeview output needs the treeview extra: "
271 'uv pip install "ete4[treeview]==4.4.0"'
272 ) from exc
274 for node in tree.traverse():
275 style = NodeStyle()
276 style["size"] = args.leaf_size if node.is_leaf else args.internal_size
277 style["fgcolor"] = (
278 support_color(node, args)
279 if args.color_by_support and not node.is_leaf
280 else args.leaf_color
281 if node.is_leaf
282 else args.internal_color
283 )
284 node.set_style(style)
286 tree_style = TreeStyle()
287 tree_style.mode = "c" if args.mode == "circular" else "r"
288 tree_style.show_leaf_name = args.show_names
289 tree_style.show_branch_support = args.show_support
290 tree_style.show_branch_length = args.show_lengths
291 tree_style.show_scale = args.show_scale
293 if args.mode == "circular":
294 tree_style.arc_start = args.arc_start
295 tree_style.arc_span = args.arc_span
297 if args.title:
298 tree_style.title.add_face(
299 TextFace(args.title, fsize=max(args.label_size, 12), bold=True),
300 column=0,
301 )
302 return tree_style
305def render_treeview(
306 tree: Tree,
307 output: Path,
308 args: argparse.Namespace,
309) -> None:
310 """Render PNG, PDF, or SVG through Qt treeview."""
311 if output.suffix.lower() not in {".png", ".pdf", ".svg"}:
312 raise UserInputError("Qt treeview output must end in .png, .pdf, or .svg")
313 if not output.parent.is_dir():
314 raise UserInputError(f"output directory does not exist: {output.parent}")
316 tree_style = create_treeview_style(tree, args)
317 render_args = {
318 "tree_style": tree_style,
319 "units": args.units,
320 "dpi": args.dpi,
321 }
322 if args.width is not None:
323 render_args["w"] = args.width
324 if args.height is not None:
325 render_args["h"] = args.height
327 tree.render(str(output), **render_args)
328 print(f"Wrote Qt treeview output: {output}")
331def choose_engine(args: argparse.Namespace) -> str:
332 """Choose a renderer from the requested engine and output suffix."""
333 if args.engine != "auto":
334 return args.engine
335 if args.output is None or args.output.suffix.lower() == ".png":
336 return "smartview"
337 if args.output.suffix.lower() in {".pdf", ".svg"}:
338 return "treeview"
339 raise UserInputError(
340 "cannot infer renderer from output suffix; use .png, .pdf, or .svg"
341 )
344def build_parser() -> argparse.ArgumentParser:
345 parser = argparse.ArgumentParser(
346 description=(
347 "Explore a tree with ETE 4 SmartView or render it with "
348 "SmartView/Qt treeview"
349 )
350 )
351 parser.add_argument("input", type=Path, help="Newick tree file")
352 parser.add_argument(
353 "output",
354 type=Path,
355 nargs="?",
356 help="optional .png, .pdf, or .svg output; omit for SmartView",
357 )
358 parser.add_argument(
359 "--version",
360 action="version",
361 version=f"%(prog)s (ETE {ete4.__version__})",
362 )
363 parser.add_argument("--parser", type=parser_spec, default=0)
364 parser.add_argument(
365 "--engine",
366 choices=["auto", "smartview", "treeview"],
367 default="auto",
368 )
369 parser.add_argument("--mode", type=mode_spec, default="rectangular")
370 parser.add_argument("--title")
372 display = parser.add_argument_group("display")
373 display.add_argument(
374 "--show-names",
375 action=argparse.BooleanOptionalAction,
376 default=True,
377 )
378 display.add_argument("--show-support", action="store_true")
379 display.add_argument("--show-lengths", action="store_true")
380 display.add_argument("--show-scale", action="store_true")
381 display.add_argument("--color-by-support", action="store_true")
382 display.add_argument("--label-size", type=int, default=12)
383 display.add_argument("--leaf-size", type=int, default=5)
384 display.add_argument("--internal-size", type=int, default=4)
385 display.add_argument("--leaf-color", default="#2166ac")
386 display.add_argument("--internal-color", default="#777777")
388 support = parser.add_argument_group("support colors")
389 support.add_argument(
390 "--support-scale", choices=["fraction", "percent"],
391 help="source support units; required with --color-by-support",
392 )
393 support.add_argument("--high-support", type=float, default=0.9)
394 support.add_argument("--moderate-support", type=float, default=0.7)
395 support.add_argument("--high-support-color", default="#1b7837")
396 support.add_argument("--moderate-support-color", default="#e08214")
397 support.add_argument("--low-support-color", default="#b2182b")
398 support.add_argument("--missing-support-color", default="#999999")
400 smartview = parser.add_argument_group("SmartView")
401 smartview.add_argument("--collapse-pixels", type=float, default=8)
402 smartview.add_argument("--content-pixels", type=float, default=4)
403 smartview.add_argument("--host", default="127.0.0.1")
404 smartview.add_argument("--port", type=int)
405 smartview.add_argument("--no-browser", action="store_true")
406 smartview.add_argument(
407 "--allow-remote-bind",
408 action="store_true",
409 help="allow SmartView to bind beyond loopback",
410 )
412 output = parser.add_argument_group("static output")
413 output.add_argument("--width", type=int)
414 output.add_argument("--height", type=int)
415 output.add_argument("--units", choices=["px", "mm", "in"], default="px")
416 output.add_argument("--dpi", type=int, default=300)
417 output.add_argument("--arc-start", type=int, default=0)
418 output.add_argument("--arc-span", type=int, default=360)
419 return parser
422def validate_args(args: argparse.Namespace) -> None:
423 """Validate numerical and renderer-specific options."""
424 if not 0 <= args.moderate_support <= args.high_support <= 1:
425 raise UserInputError(
426 "support thresholds must satisfy 0 <= moderate-support <= high-support <= 1"
427 )
428 if args.color_by_support and args.support_scale is None:
429 raise UserInputError("--color-by-support requires --support-scale fraction or percent")
430 for name in (
431 "label_size",
432 "leaf_size",
433 "internal_size",
434 "collapse_pixels",
435 "content_pixels",
436 ):
437 if getattr(args, name) < 0:
438 raise UserInputError(f"{name.replace('_', '-')} cannot be negative")
439 for name in ("width", "height", "dpi"):
440 value = getattr(args, name)
441 if value is not None and value <= 0:
442 raise UserInputError(f"{name} must be positive")
443 if args.port is not None and not 1 <= args.port <= 65535:
444 raise UserInputError("port must be between 1 and 65535")
445 if not -360 <= args.arc_start <= 360:
446 raise UserInputError("arc-start must be between -360 and 360 degrees")
447 if not 0 < args.arc_span <= 360:
448 raise UserInputError("arc-span must be greater than 0 and at most 360 degrees")
451def main(argv: list[str] | None = None) -> int:
452 args = build_parser().parse_args(argv)
453 try:
454 validate_args(args)
455 tree = load_tree(args.input, args.parser)
456 if args.color_by_support:
457 for node in tree.traverse():
458 if not node.is_leaf:
459 support_fraction(node.support, args.support_scale)
460 engine = choose_engine(args)
462 if engine == "smartview":
463 layout = create_smartview_layout(args)
464 if args.output is None:
465 run_interactive_smartview(tree, layout, args)
466 else:
467 render_smartview(tree, layout, args.output, args)
468 else:
469 if args.output is None:
470 raise UserInputError("Qt treeview mode requires an output path")
471 render_treeview(tree, args.output, args)
472 except UserInputError as exc:
473 print(f"error: {exc}", file=sys.stderr)
474 return 2
475 except Exception as exc: # Preserve unexpected ETE/renderer details.
476 print(
477 f"unexpected visualization error: {type(exc).__name__}: {exc}",
478 file=sys.stderr,
479 )
480 return 1
481 return 0
484if __name__ == "__main__":
485 raise SystemExit(main())