diff --git a/python/CHANGELOG.rst b/python/CHANGELOG.rst index 1e21d6af02..9815792123 100644 --- a/python/CHANGELOG.rst +++ b/python/CHANGELOG.rst @@ -59,6 +59,11 @@ which assigns a string to node and mutation symbols, commonly shown on mouseover. This can reduce label clutter while retaining useful info (:user:`hyanwong`, :pr:`3007`) +- Added (currently undocumented) use of the `order` parameter in ``Tree.draw_svg()`` to + pass a subset of nodes, so subtrees can be visually collapsed. Additionally, an option + ``pack_untracked_polytomies`` allows large polytomies involving untracked samples to + be summarised as a dotted line (:user:`hyanwong`, :issue:`3011` :pr:`3010`, :pr:`3012`) + -------------------- [0.5.8] - 2024-06-27 -------------------- diff --git a/python/tests/data/svg/internal_sample_ts.svg b/python/tests/data/svg/internal_sample_ts.svg index 4843cd70c4..58e9185b17 100644 --- a/python/tests/data/svg/internal_sample_ts.svg +++ b/python/tests/data/svg/internal_sample_ts.svg @@ -1,7 +1,7 @@ - + diff --git a/python/tests/data/svg/tree.svg b/python/tests/data/svg/tree.svg index 242cec60b3..1eae56aa85 100644 --- a/python/tests/data/svg/tree.svg +++ b/python/tests/data/svg/tree.svg @@ -1,7 +1,7 @@ - + diff --git a/python/tests/data/svg/tree_both_axes.svg b/python/tests/data/svg/tree_both_axes.svg index 5b7e4950aa..863bdf09b8 100644 --- a/python/tests/data/svg/tree_both_axes.svg +++ b/python/tests/data/svg/tree_both_axes.svg @@ -1,7 +1,7 @@ - + diff --git a/python/tests/data/svg/tree_muts.svg b/python/tests/data/svg/tree_muts.svg index d5e0ecaaa7..1651e84776 100644 --- a/python/tests/data/svg/tree_muts.svg +++ b/python/tests/data/svg/tree_muts.svg @@ -1,7 +1,7 @@ - + diff --git a/python/tests/data/svg/tree_muts_all_edge.svg b/python/tests/data/svg/tree_muts_all_edge.svg index 368d06020d..2777e7115b 100644 --- a/python/tests/data/svg/tree_muts_all_edge.svg +++ b/python/tests/data/svg/tree_muts_all_edge.svg @@ -1,7 +1,7 @@ - + diff --git a/python/tests/data/svg/tree_poly.svg b/python/tests/data/svg/tree_poly.svg new file mode 100644 index 0000000000..f1997c8b5b --- /dev/null +++ b/python/tests/data/svg/tree_poly.svg @@ -0,0 +1,219 @@ + + + + + + + + + + + + + 0 + + + + + 1 + + + + + 2 + + + + + + 3 + + + + + 4 + + + + + 5 + + + + + 6 + + + + 30 + + + + 31 + + + + + + 7 + + + + + 8 + + + + + 9 + + + + + + 10 + + + + + 11 + + + + + 12 + + + + + 13 + + + + 32 + + + + 33 + + + + + + 14 + + + + + 15 + + + + + 16 + + + + + + 17 + + + + + 18 + + + + + 19 + + + + + 20 + + + + 34 + + + + 35 + + + + + + + 21 + + + + + 22 + + + + 36 + + + + + + 23 + + + + + 24 + + + + 37 + + + + + + 25 + + + + + 26 + + + + 38 + + + + + + 27 + + + + + 28 + + + + + 29 + + + + 39 + + + + 40 + + + 41 + + + + diff --git a/python/tests/data/svg/tree_poly_tracked.svg b/python/tests/data/svg/tree_poly_tracked.svg new file mode 100644 index 0000000000..517d508e21 --- /dev/null +++ b/python/tests/data/svg/tree_poly_tracked.svg @@ -0,0 +1,103 @@ + + + + + + + + + + + + + + 20 + + + + +3/𝟑 + + + + 34 + + + + +3/𝟑 + + + + 35 + + + + + + + +2 + 36 + + + + + + 23 + + + + + 24 + + + + 37 + + + + + + 25 + + + + + 26 + + + + 38 + + + + + + 27 + + + + + 28 + + + + + 29 + + + + 39 + + + + 40 + + + + +14/𝟐 + + + 41 + + + + diff --git a/python/tests/data/svg/tree_poly_tracked_collapse.svg b/python/tests/data/svg/tree_poly_tracked_collapse.svg new file mode 100644 index 0000000000..437c2c355b --- /dev/null +++ b/python/tests/data/svg/tree_poly_tracked_collapse.svg @@ -0,0 +1,90 @@ + + + + + + + + + + + + + + 20 + + + + +3/𝟑 + + + + 34 + + + + +3/𝟑 + + + + 35 + + + + + + + +2 + 36 + + + + + + 23 + + + + + 24 + + + + 37 + + + + + + 25 + + + + + 26 + + + + 38 + + + + + + +3 + 39 + + + + 40 + + + + +14/𝟐 + + + 41 + + + + diff --git a/python/tests/data/svg/tree_simple_collapsed.svg b/python/tests/data/svg/tree_simple_collapsed.svg index 9299313dc3..67446bb6dd 100644 --- a/python/tests/data/svg/tree_simple_collapsed.svg +++ b/python/tests/data/svg/tree_simple_collapsed.svg @@ -1,7 +1,7 @@ - + diff --git a/python/tests/data/svg/tree_subtree.svg b/python/tests/data/svg/tree_subtree.svg index caa32f80df..593b7a1c64 100644 --- a/python/tests/data/svg/tree_subtree.svg +++ b/python/tests/data/svg/tree_subtree.svg @@ -1,7 +1,7 @@ - + diff --git a/python/tests/data/svg/tree_subtrees_with_collapsed.svg b/python/tests/data/svg/tree_subtrees_with_collapsed.svg index bee71475ac..5bd9d76b74 100644 --- a/python/tests/data/svg/tree_subtrees_with_collapsed.svg +++ b/python/tests/data/svg/tree_subtrees_with_collapsed.svg @@ -1,7 +1,7 @@ - + diff --git a/python/tests/data/svg/tree_timed_muts.svg b/python/tests/data/svg/tree_timed_muts.svg index 27e1c8aaee..ca187728ca 100644 --- a/python/tests/data/svg/tree_timed_muts.svg +++ b/python/tests/data/svg/tree_timed_muts.svg @@ -1,7 +1,7 @@ - + diff --git a/python/tests/data/svg/tree_x_axis.svg b/python/tests/data/svg/tree_x_axis.svg index 4a6cca8e30..3c4f2e7bd1 100644 --- a/python/tests/data/svg/tree_x_axis.svg +++ b/python/tests/data/svg/tree_x_axis.svg @@ -1,7 +1,7 @@ - + diff --git a/python/tests/data/svg/tree_y_axis_rank.svg b/python/tests/data/svg/tree_y_axis_rank.svg index abba87a033..79af1d5072 100644 --- a/python/tests/data/svg/tree_y_axis_rank.svg +++ b/python/tests/data/svg/tree_y_axis_rank.svg @@ -1,7 +1,7 @@ - + diff --git a/python/tests/data/svg/ts.svg b/python/tests/data/svg/ts.svg index 13bc8db1cb..0323aa1c99 100644 --- a/python/tests/data/svg/ts.svg +++ b/python/tests/data/svg/ts.svg @@ -1,7 +1,7 @@ - + diff --git a/python/tests/data/svg/ts_max_trees.svg b/python/tests/data/svg/ts_max_trees.svg index c50331492d..a012e7e778 100644 --- a/python/tests/data/svg/ts_max_trees.svg +++ b/python/tests/data/svg/ts_max_trees.svg @@ -1,7 +1,7 @@ - + diff --git a/python/tests/data/svg/ts_max_trees_treewise.svg b/python/tests/data/svg/ts_max_trees_treewise.svg index 055c1f3eaf..f088c89e68 100644 --- a/python/tests/data/svg/ts_max_trees_treewise.svg +++ b/python/tests/data/svg/ts_max_trees_treewise.svg @@ -1,7 +1,7 @@ - + diff --git a/python/tests/data/svg/ts_multiroot.svg b/python/tests/data/svg/ts_multiroot.svg index 0d8dd1f35a..9d2637d45a 100644 --- a/python/tests/data/svg/ts_multiroot.svg +++ b/python/tests/data/svg/ts_multiroot.svg @@ -1,7 +1,7 @@ - + diff --git a/python/tests/data/svg/ts_mut_highlight.svg b/python/tests/data/svg/ts_mut_highlight.svg index 54935aa939..72c9fda2d8 100644 --- a/python/tests/data/svg/ts_mut_highlight.svg +++ b/python/tests/data/svg/ts_mut_highlight.svg @@ -1,7 +1,7 @@ - + diff --git a/python/tests/data/svg/ts_mut_times.svg b/python/tests/data/svg/ts_mut_times.svg index a72e9c7a9c..6f4955ec03 100644 --- a/python/tests/data/svg/ts_mut_times.svg +++ b/python/tests/data/svg/ts_mut_times.svg @@ -1,7 +1,7 @@ - + diff --git a/python/tests/data/svg/ts_mut_times_logscale.svg b/python/tests/data/svg/ts_mut_times_logscale.svg index a269eab0c2..4ba7e1a005 100644 --- a/python/tests/data/svg/ts_mut_times_logscale.svg +++ b/python/tests/data/svg/ts_mut_times_logscale.svg @@ -1,7 +1,7 @@ - + diff --git a/python/tests/data/svg/ts_mut_times_titles.svg b/python/tests/data/svg/ts_mut_times_titles.svg index 78aea0d3e1..5a76bff164 100644 --- a/python/tests/data/svg/ts_mut_times_titles.svg +++ b/python/tests/data/svg/ts_mut_times_titles.svg @@ -1,7 +1,7 @@ - + diff --git a/python/tests/data/svg/ts_mutations_no_edges.svg b/python/tests/data/svg/ts_mutations_no_edges.svg index 4b9b1026d0..7d0f4171f5 100644 --- a/python/tests/data/svg/ts_mutations_no_edges.svg +++ b/python/tests/data/svg/ts_mutations_no_edges.svg @@ -1,7 +1,7 @@ - + diff --git a/python/tests/data/svg/ts_mutations_timed_no_edges.svg b/python/tests/data/svg/ts_mutations_timed_no_edges.svg index eb6cb6cdd6..c4f5dfa42b 100644 --- a/python/tests/data/svg/ts_mutations_timed_no_edges.svg +++ b/python/tests/data/svg/ts_mutations_timed_no_edges.svg @@ -1,7 +1,7 @@ - + diff --git a/python/tests/data/svg/ts_no_axes.svg b/python/tests/data/svg/ts_no_axes.svg index 5ce4804a63..d1674bb7cc 100644 --- a/python/tests/data/svg/ts_no_axes.svg +++ b/python/tests/data/svg/ts_no_axes.svg @@ -1,7 +1,7 @@ - + diff --git a/python/tests/data/svg/ts_plain.svg b/python/tests/data/svg/ts_plain.svg index a5c01e3e0e..eab1a91be1 100644 --- a/python/tests/data/svg/ts_plain.svg +++ b/python/tests/data/svg/ts_plain.svg @@ -1,7 +1,7 @@ - + diff --git a/python/tests/data/svg/ts_plain_no_xlab.svg b/python/tests/data/svg/ts_plain_no_xlab.svg index ab4cc34cdf..c17c1fb425 100644 --- a/python/tests/data/svg/ts_plain_no_xlab.svg +++ b/python/tests/data/svg/ts_plain_no_xlab.svg @@ -1,7 +1,7 @@ - + diff --git a/python/tests/data/svg/ts_plain_y.svg b/python/tests/data/svg/ts_plain_y.svg index 72f704f22b..95c591307f 100644 --- a/python/tests/data/svg/ts_plain_y.svg +++ b/python/tests/data/svg/ts_plain_y.svg @@ -1,7 +1,7 @@ - + diff --git a/python/tests/data/svg/ts_rank.svg b/python/tests/data/svg/ts_rank.svg index de9f9390e7..b507842dbc 100644 --- a/python/tests/data/svg/ts_rank.svg +++ b/python/tests/data/svg/ts_rank.svg @@ -1,7 +1,7 @@ - + diff --git a/python/tests/data/svg/ts_x_lim.svg b/python/tests/data/svg/ts_x_lim.svg index dcc6c63032..5a8256ec64 100644 --- a/python/tests/data/svg/ts_x_lim.svg +++ b/python/tests/data/svg/ts_x_lim.svg @@ -1,7 +1,7 @@ - + diff --git a/python/tests/data/svg/ts_xlabel.svg b/python/tests/data/svg/ts_xlabel.svg index 3bf4c97858..f786ee3682 100644 --- a/python/tests/data/svg/ts_xlabel.svg +++ b/python/tests/data/svg/ts_xlabel.svg @@ -1,7 +1,7 @@ - + diff --git a/python/tests/data/svg/ts_y_axis.svg b/python/tests/data/svg/ts_y_axis.svg index 6e4cefac15..d590fef37a 100644 --- a/python/tests/data/svg/ts_y_axis.svg +++ b/python/tests/data/svg/ts_y_axis.svg @@ -1,7 +1,7 @@ - + diff --git a/python/tests/data/svg/ts_y_axis_log.svg b/python/tests/data/svg/ts_y_axis_log.svg index 9b60d4e6ae..e1e5814fcf 100644 --- a/python/tests/data/svg/ts_y_axis_log.svg +++ b/python/tests/data/svg/ts_y_axis_log.svg @@ -1,7 +1,7 @@ - + diff --git a/python/tests/data/svg/ts_y_axis_regular.svg b/python/tests/data/svg/ts_y_axis_regular.svg index d90bd53a7b..1ee704a096 100644 --- a/python/tests/data/svg/ts_y_axis_regular.svg +++ b/python/tests/data/svg/ts_y_axis_regular.svg @@ -1,7 +1,7 @@ - + diff --git a/python/tests/test_drawing.py b/python/tests/test_drawing.py index 0684679ee0..5504c73981 100644 --- a/python/tests/test_drawing.py +++ b/python/tests/test_drawing.py @@ -3051,6 +3051,53 @@ def test_known_svg_tree_subtrees_with_collapsed(self, overwrite_viz, draw_plotbo svg, "tree_subtrees_with_collapsed.svg", overwrite_viz, has_root=False ) + def test_known_svg_tree_polytomy(self, overwrite_viz, draw_plotbox): + tracked_nodes = [20, 24, 25, 27, 28, 29] + tree = tskit.Tree.generate_balanced(30, arity=4) + svg = tree.draw_svg( + time_scale="rank", + debug_box=draw_plotbox, + size=(600, 200), + style="".join(f".n{u} > .sym {{fill: cyan}}" for u in tracked_nodes + [39]), + ) + self.verify_known_svg( + svg, "tree_poly.svg", overwrite_viz, width=600, height=200 + ) + + def test_known_svg_tree_polytomy_tracked(self, overwrite_viz, draw_plotbox): + tracked_nodes = [20, 24, 25, 27, 28, 29] + tree = tskit.Tree.generate_balanced(30, arity=4, tracked_samples=tracked_nodes) + svg = tree.draw_svg( + time_scale="rank", + order=drawing._postorder_tracked_minlex_traversal(tree), + debug_box=draw_plotbox, + pack_untracked_polytomies=True, + size=(600, 200), + style="".join(f".n{u} > .sym {{fill: cyan}}" for u in tracked_nodes + [39]), + ) + self.verify_known_svg( + svg, "tree_poly_tracked.svg", overwrite_viz, width=600, height=200 + ) + + def test_known_svg_tree_polytomy_tracked_collapse( + self, overwrite_viz, draw_plotbox + ): + tracked_nodes = [20, 24, 25, 27, 28, 29] + tree = tskit.Tree.generate_balanced(30, arity=4, tracked_samples=tracked_nodes) + svg = tree.draw_svg( + time_scale="rank", + order=drawing._postorder_tracked_minlex_traversal( + tree, collapse_tracked=True + ), + debug_box=draw_plotbox, + size=(600, 200), + pack_untracked_polytomies=True, + style="".join(f".n{u} > .sym {{fill: cyan}}" for u in tracked_nodes + [39]), + ) + self.verify_known_svg( + svg, "tree_poly_tracked_collapse.svg", overwrite_viz, width=600, height=200 + ) + class TestRounding: def test_rnd(self): @@ -3062,3 +3109,9 @@ def test_rnd(self): assert 1111110 == drawing.rnd(1111111) assert 123.457 == drawing.rnd(123.4567) assert 123.456 == drawing.rnd(123.4564) + + +class TestDrawingTraversals: + # TODO: test drawing._postorder_tracked_minlex_traversal and + # drawing._postorder_tracked_node_traversal + pass diff --git a/python/tskit/drawing.py b/python/tskit/drawing.py index 6191ada42d..dc5ff716b7 100644 --- a/python/tskit/drawing.py +++ b/python/tskit/drawing.py @@ -360,6 +360,12 @@ def rnd(x): return x +def bold_integer(number): + # For simple integers, it's easier to use bold unicode characters + # than to try to get the SVG to render a bold font for part of a string + return "".join("𝟎𝟏𝟐𝟑𝟒𝟓𝟔𝟕𝟖𝟗"[int(digit)] for digit in str(number)) + + def edge_and_sample_nodes(ts, omit_regions=None): """ Return ids of nodes which are mentioned in an edge in this tree sequence or which @@ -383,6 +389,82 @@ def edge_and_sample_nodes(ts, omit_regions=None): ) +def _postorder_tracked_node_traversal(tree, root, collapse_tracked, key_dict=None): + # Postorder traversal that only descends into subtrees if they contain + # a tracked node. Additionally, if collapse_tracked is not None, it is + # interpreted as a proportion, so that we do not descend into a subtree if + # that proportion or greater of the samples in the subtree are tracked. + # If key_dict is provided, use this to sort the children. This allows + # us to put e.g. the subtrees containing the most tracked nodes first. + # Private function, for use only in drawing.postorder_tracked_minlex_traversal() + + # If we deliberately specify the virtual root, it should also be returned + is_virtual_root = root == tree.virtual_root + if root == tskit.NULL: + root = tree.virtual_root + stack = [(root, False)] + while stack: + u, visited = stack.pop() + if visited: + if u != tree.virtual_root or is_virtual_root: + yield u + else: + if tree.num_children(u) == 0: + yield u + elif tree.num_tracked_samples(u) == 0: + yield u + elif ( + collapse_tracked is not None + and tree.num_children(u) != 1 + and tree.num_tracked_samples(u) + >= collapse_tracked * tree.num_samples(u) + ): + yield u + else: + stack.append((u, True)) + if key_dict is None: + stack.extend((c, False) for c in tree.children(u)) + else: + stack.extend( + sorted( + ((c, False) for c in tree.children(u)), + key=lambda v: key_dict[v[0]], + reverse=True, + ) + ) + + +def _postorder_tracked_minlex_traversal(tree, root=None, *, collapse_tracked=None): + """ + Postorder traversal for drawing purposes that places child nodes with the + most tracked sample descendants first (then sorts ties by minlex on leaf node ids). + Additionally, this traversal only descends into subtrees if they contain a tracked + node, and may not descend into other subtree, if the ``collapse_tracked`` + parameter is set to a numeric value. More specifically, if the proportion of + tracked samples in the subtree is greater than or equal to ``collapse_tracked``, + the subtree is not descended into. + """ + + key_dict = {} + parent_array = tree.parent_array + prev = tree.virtual_root + if root is None: + root = tskit.NULL + for u in _postorder_tracked_node_traversal(tree, root, collapse_tracked): + is_tip = parent_array[prev] != u + prev = u + if is_tip: + # Sort by number of tracked samples (desc), then by minlex + key_dict[u] = (-tree.num_tracked_samples(u), u) + else: + min_tip_id = min(key_dict[v][1] for v in tree.children(u) if v in key_dict) + key_dict[u] = (-tree.num_tracked_samples(u), min_tip_id) + + return _postorder_tracked_node_traversal( + tree, root, collapse_tracked, key_dict=key_dict + ) + + def draw_tree( tree, width=None, @@ -688,6 +770,8 @@ class SvgAxisPlot(SvgPlot): ".tree text, .tree-sequence text {dominant-baseline: central}" ".plotbox .lab.lft {text-anchor: end}" ".plotbox .lab.rgt {text-anchor: start}" + ".polytomy line {stroke: black; stroke-dasharray: 1px, 1px}" + ".polytomy text {paint-order:stroke;stroke-width:0.3em;stroke:white}" ) # TODO: we may want to make some of the constants below into parameters @@ -1315,6 +1399,10 @@ class SvgTree(SvgAxisPlot): See :meth:`Tree.draw_svg` for a description of usage and frequently used parameters. """ + PolytomyLine = collections.namedtuple( + "PolytomyLine", "num_branches, num_samples, line_pos" + ) + def __init__( self, tree, @@ -1348,6 +1436,7 @@ def __init__( mutation_label_attrs=None, offsets=None, omit_sites=None, + pack_untracked_polytomies=None, **kwargs, ): if max_time is None and max_tree_height is not None: @@ -1371,6 +1460,7 @@ def __init__( if symbol_size is None: symbol_size = 6 self.symbol_size = symbol_size + self.pack_untracked_polytomies = pack_untracked_polytomies ts = tree.tree_sequence tree_index = tree.index if offsets is not None: @@ -1701,16 +1791,43 @@ def assign_x_coordinates(self): ) # Set up x positions for nodes node_xpos = {} + untracked_children = collections.defaultdict(list) + self.extra_line = {} # To store a dotted line to represent polytomies leaf_x = 0 # First leaf starts at x=1, to give some space between Y axis & leaf - prev = self.tree.virtual_root tree = self.tree + prev = tree.virtual_root for u in self.postorder_nodes: + parent = tree.parent(u) + if parent == prev: + raise ValueError("Nodes must be passed in postorder to Tree.draw_svg()") is_tip = tree.parent(prev) != u if is_tip: - leaf_x += 1 - node_xpos[u] = leaf_x + if self.pack_untracked_polytomies and tree.num_tracked_samples(u) == 0: + untracked_children[parent].append(u) + else: + leaf_x += 1 + node_xpos[u] = leaf_x else: + # Concatenate all the untracked children + num_untracked_children = len(untracked_children[u]) child_x = [node_xpos[c] for c in tree.children(u) if c in node_xpos] + if num_untracked_children > 0: + if num_untracked_children <= 1: + # If only a single non-focal lineage, we might as well show it + for child in untracked_children[u]: + leaf_x += 1 + node_xpos[child] = leaf_x + child_x.append(leaf_x) + else: + # Otherwise show a horizontal line with the number of lineages + # Extra length of line is equal to log of the polytomy size + self.extra_line[u] = self.PolytomyLine( + num_untracked_children, + sum(tree.num_samples(v) for v in untracked_children[u]), + [leaf_x, leaf_x + 1 + np.log(num_untracked_children)], + ) + child_x.append(leaf_x + 1) + leaf_x = self.extra_line[u].line_pos[1] assert len(child_x) != 0 # Must have prev hit somethng defined as a tip if len(child_x) == 1: node_xpos[u] = child_x[0] @@ -1718,15 +1835,15 @@ def assign_x_coordinates(self): a = min(child_x) b = max(child_x) node_xpos[u] = a + (b - a) / 2 - if tree.parent(u) == prev: - raise ValueError("Nodes must be passed in postorder to Tree.draw_svg()") prev = u - # Now rescale to the plot width: leaf_x is the maximum value of the last leaf if len(node_xpos) > 0: scale = self.plotbox.width / leaf_x lft = self.plotbox.left - scale / 2 self.node_x_coord = {k: lft + v * scale for k, v in node_xpos.items()} + for v in self.extra_line.values(): + for i in range(len(v.line_pos)): + v.line_pos[i] = lft + v.line_pos[i] * scale def info_classes(self, focal_node_id): """ @@ -1830,6 +1947,28 @@ def draw_tree(self): o = (0, 0) v = parent_array[u] + # Add polytomy line if necessary + if u in self.extra_line: + info = self.extra_line[u] + x2 = info.line_pos[1] - pu[0] + poly = dwg.g(class_="polytomy") + poly.add( + dwg.line( + start=(0, 0), + end=(x2, 0), + ) + ) + poly.add( + dwg.text( + f"+{info.num_samples}/{bold_integer(info.num_branches)}", + font_style="italic", + x=[rnd(x2)], + dy=[rnd(-self.text_height / 10)], # make the plus sign line up + text_anchor="end", + ) + ) + curr_svg_group.add(poly) + # Add edge above node first => on layer underneath anything else draw_edge_above_node = False try: