Merge pull request #369 from hyanwong/viz-subtree

Plot RE nodes as open circles

sc2ts/info.py

@@ -1614,6 +1614,7 @@ def draw_subtree(
         mutation_labels=None,
         size=None,
         style="",
+        symbol_size=4,
         **kwargs,
     ):
         """
@@ -1796,7 +1797,12 @@ def draw_subtree(
             lab_css = ", ".join(f".mut.m{m} .lab" for m in reverted_mutations)
             sym_css = ", ".join(f".mut.m{m} .sym" for m in reverted_mutations)
             styles.append(lab_css + "{fill: magenta}" + sym_css + "{stroke: magenta}")
-
+        # Recombination nodes as larger open circles
+        re_nodes = np.where(ts.nodes_flags & core.NODE_IS_RECOMBINANT)[0]
+        styles.append(
+            ",".join([f".node.n{u} > .sym" for u in re_nodes]) +
+            f"{{r:{symbol_size/2*1.5:.2f}px; stroke:black; fill:white}}"
+        )
         return tree.draw_svg(
             time_scale=time_scale,
             y_axis=True,
@@ -1806,7 +1812,7 @@ def draw_subtree(
             order=order,
             mutation_labels=mutation_labels,
             all_edge_mutations=True,
-            symbol_size=4,
+            symbol_size=symbol_size,
             pack_untracked_polytomies=pack_untracked_polytomies,
             style="".join(styles) + style,
             **kwargs,
@@ -1893,6 +1899,7 @@ def draw_svg(
         highlight_universal_mutations=None,
         x_regions=None,
         node_labels=None,
+        symbol_size=3,
         **kwargs,
     ):
         """
@@ -2011,6 +2018,7 @@ def draw_svg(
             lab_css = ", ".join(f".mut.m{m} .lab" for m in reverted_mutations)
             sym_css = ", ".join(f".mut.m{m} .sym" for m in reverted_mutations)
             styles.append(lab_css + "{fill: magenta}" + sym_css + "{stroke: magenta}")
+        # mutations shared by all the samples (above their mrca)
         if len(universal_mutations) > 0:
             lab_css = ", ".join(f".mut.m{m} .lab" for m in universal_mutations)
             sym_css = ", ".join(f".mut.m{m} .sym" for m in universal_mutations)
@@ -2021,6 +2029,12 @@ def draw_svg(
                 lab_css + "{font-weight: bold}" + sym_css + "{stroke-width: 3}"
             )
             styles.append(sym_ax_css + "{stroke-width: 8}")
+        # recombination nodes in larger open white circles
+        re_nodes = np.where(ts.nodes_flags & core.NODE_IS_RECOMBINANT)[0]
+        styles.append(
+            ",".join([f".node.n{u} > .sym" for u in re_nodes]) +
+            f"{{r: {symbol_size/2*1.5:.2f}px; stroke: black; fill: white}}"
+        )
         svg = self.ts.draw_svg(
             size=size,
             time_scale=time_scale,
@@ -2029,6 +2043,7 @@ def draw_svg(
             y_ticks=y_ticks,
             node_labels=node_labels,
             style="".join(styles) + style,
+            symbol_size=symbol_size,
             **kwargs,
         )
 

tests/test_info.py

@@ -186,3 +186,29 @@ def test_plots(self, fx_ti_2020_02_13, method):
         assert isinstance(fig, matplotlib.figure.Figure)
         for ax in axes:
             assert isinstance(ax, matplotlib.axes.Axes)
+
+    def test_draw_pango_lineage_subtree(self, fx_ti_2020_02_13):
+        ti = fx_ti_2020_02_13
+        svg = ti.draw_pango_lineage_subtree("A")
+        svg2 = ti.draw_subtree(tracked_pango=["A"])
+        assert svg == svg2
+        assert svg.startswith("<svg")
+        for u in ti.pango_lineage_samples['A']:
+            assert f"node n{u}" in svg
+
+    def test_draw_subtree(self, fx_ti_2020_02_13):
+        ti = fx_ti_2020_02_13
+        samples = ti.ts.samples()[0:5]
+        svg = ti.draw_subtree(tracked_samples=samples)
+        assert svg.startswith("<svg")
+        for u in samples:
+            assert f"node n{u}" in svg
+
+
+class TestSampleGroupInfo:
+    def test_draw_svg(self, fx_ti_2020_02_13):
+        ti = fx_ti_2020_02_13
+        sg = list(ti.nodes_sample_group.keys())[0]
+        sg_info = ti.get_sample_group_info(sg)
+        svg = sg_info.draw_svg()
+        assert svg.startswith("<svg")