From b07bee5f8eb88f91170ff34f4860bf2a87062a90 Mon Sep 17 00:00:00 2001 From: Yan Wong Date: Wed, 6 Nov 2024 16:17:48 +0000 Subject: [PATCH] Fix SVG pack_untracked_polytomies bug Fixes #3049 --- python/CHANGELOG.rst | 6 ++++++ python/tests/test_drawing.py | 13 +++++++++++++ python/tskit/drawing.py | 29 ++++++++++++++++------------- 3 files changed, 35 insertions(+), 13 deletions(-) diff --git a/python/CHANGELOG.rst b/python/CHANGELOG.rst index af0a1de8fd..d8073590a4 100644 --- a/python/CHANGELOG.rst +++ b/python/CHANGELOG.rst @@ -2,6 +2,12 @@ [0.6.1] - 2024-XX-XX -------------------- +**Bugfixes** + +- Fix to ``Tree.draw_svg()`` with the (undocumented) experimental + ``pack_untracked_polytomies=True`` option. (:user:`hyanwong`, + :issue:`3049`, :pr:`3050`) + -------------------- [0.6.0] - 2024-10-16 diff --git a/python/tests/test_drawing.py b/python/tests/test_drawing.py index 8da436dab3..0333141837 100644 --- a/python/tests/test_drawing.py +++ b/python/tests/test_drawing.py @@ -2602,6 +2602,19 @@ def test_only_subset_nodes_in_rank(self, caplog): ) assert "Ticks {10: '10'} lie outside the plotted axis" in caplog.text + def test_polytomy_collapsing(self): + tree = tskit.Tree.generate_balanced( + 20, arity=4, tracked_samples=np.arange(2, 8) + ) + svg = tree.draw_svg(pack_untracked_polytomies=True) + # Should have one collapsed node (untracked samples 8 and 9) + # and two "polytomy lines" (from nodes 21 and 28 (the root)) + assert svg.count('class="polytomy"') == 2 # poolytomy lines + collapsed_symbol = re.search("]*>", svg) + assert collapsed_symbol is not None + assert collapsed_symbol.group(0).count("sym") == 1 + assert collapsed_symbol.group(0).count("multi") == 1 + class TestDrawKnownSvg(TestDrawSvgBase): """ diff --git a/python/tskit/drawing.py b/python/tskit/drawing.py index a1c0e5cfee..7c6c591bfa 100644 --- a/python/tskit/drawing.py +++ b/python/tskit/drawing.py @@ -1827,23 +1827,26 @@ def assign_x_coordinates(self): prev = tree.virtual_root for u in self.postorder_nodes: parent = tree.parent(u) + omit = self.pack_untracked_polytomies and tree.num_tracked_samples(u) == 0 if parent == prev: raise ValueError("Nodes must be passed in postorder to Tree.draw_svg()") is_tip = tree.parent(prev) != u if is_tip: - if self.pack_untracked_polytomies and tree.num_tracked_samples(u) == 0: - untracked_children[parent].append(u) - else: + if not omit: leaf_x += 1 node_xpos[u] = leaf_x - else: - # Concatenate all the untracked children - num_untracked_children = len(untracked_children[u]) + elif not omit: + # Untracked children are available for packing into a polytomy summary + untracked_children = [] + if self.pack_untracked_polytomies: + untracked_children += [ + c for c in tree.children(u) if tree.num_tracked_samples(c) == 0 + ] 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]: + if len(untracked_children) > 0: + if len(untracked_children) <= 1: + # If only a single non-focal lineage, treat it as a condensed tip + for child in untracked_children: leaf_x += 1 node_xpos[child] = leaf_x child_x.append(leaf_x) @@ -1851,9 +1854,9 @@ def assign_x_coordinates(self): # 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)], + len(untracked_children), + sum(tree.num_samples(v) for v in untracked_children), + [leaf_x, leaf_x + 1 + np.log(len(untracked_children))], ) child_x.append(leaf_x + 1) leaf_x = self.extra_line[u].line_pos[1]