From 979da99e5b02275d445ca4327d1dff6d19bed2bb Mon Sep 17 00:00:00 2001 From: Yan Wong Date: Wed, 25 Sep 2024 17:23:13 +0100 Subject: [PATCH] Add ability to plot annotations (e.g. genes) along the x-axis --- python/CHANGELOG.rst | 3 + python/tests/data/svg/internal_sample_ts.svg | 2 +- python/tests/data/svg/tree.svg | 2 +- python/tests/data/svg/tree_both_axes.svg | 2 +- python/tests/data/svg/tree_muts.svg | 2 +- python/tests/data/svg/tree_muts_all_edge.svg | 2 +- python/tests/data/svg/tree_timed_muts.svg | 2 +- python/tests/data/svg/tree_x_axis.svg | 2 +- python/tests/data/svg/tree_y_axis_rank.svg | 2 +- python/tests/data/svg/ts.svg | 2 +- python/tests/data/svg/ts_max_trees.svg | 2 +- .../tests/data/svg/ts_max_trees_treewise.svg | 2 +- python/tests/data/svg/ts_multiroot.svg | 2 +- python/tests/data/svg/ts_mut_highlight.svg | 2 +- python/tests/data/svg/ts_mut_times.svg | 2 +- .../tests/data/svg/ts_mut_times_logscale.svg | 2 +- .../tests/data/svg/ts_mutations_no_edges.svg | 2 +- .../data/svg/ts_mutations_timed_no_edges.svg | 2 +- python/tests/data/svg/ts_no_axes.svg | 2 +- python/tests/data/svg/ts_plain.svg | 2 +- python/tests/data/svg/ts_plain_no_xlab.svg | 2 +- python/tests/data/svg/ts_plain_y.svg | 2 +- python/tests/data/svg/ts_rank.svg | 2 +- python/tests/data/svg/ts_x_lim.svg | 2 +- python/tests/data/svg/ts_xlabel.svg | 2 +- python/tests/data/svg/ts_y_axis.svg | 2 +- python/tests/data/svg/ts_y_axis_log.svg | 2 +- python/tests/data/svg/ts_y_axis_regular.svg | 2 +- python/tests/test_drawing.py | 65 ++++++- python/tskit/drawing.py | 172 ++++++++++-------- python/tskit/trees.py | 12 ++ 31 files changed, 202 insertions(+), 104 deletions(-) diff --git a/python/CHANGELOG.rst b/python/CHANGELOG.rst index badd7ecb0f..4500539db4 100644 --- a/python/CHANGELOG.rst +++ b/python/CHANGELOG.rst @@ -52,6 +52,9 @@ - Add provenance information to the HTML notebook representation of a tree sequence. (:user:`benjeffery`, :pr:`3001`) +- The ``.draw_svg()`` methods can add annotated genomic regions (e.g. genes) to the + x-axis. (:user:`hyanwong`, :pr:`3002`) + -------------------- [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 42d55392a0..67553fa8dd 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 0ab913202c..f578dcb89c 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 f90ce1c539..c0f14edcdf 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 3d2c017317..f02ef51002 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 adf4eb1a09..5e69e262d2 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_timed_muts.svg b/python/tests/data/svg/tree_timed_muts.svg index 0b79065c61..d5b356fff7 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 e6c7af5cd6..66492b4d84 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 9373285104..fc72850c02 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 d413b0a08b..71ba72fd55 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 7720e50ece..ad7ab7f392 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 0d1164dc65..519d7f9ee4 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 fea6053303..2ee9dc66df 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 0f7276d245..836a8bd617 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 3bd6fb5ef3..5dce0fff66 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 86382d3cf8..f2e576324e 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_mutations_no_edges.svg b/python/tests/data/svg/ts_mutations_no_edges.svg index 4feb1e2a7a..4b9b1026d0 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 de064ffc36..eb6cb6cdd6 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 051cbb1fb0..9897f1bcb8 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 6bb71f35a8..bedd75adf6 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 648a306187..4a4cc4e348 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 9aa6f26c88..9a63f2d30a 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 b8527b6638..ac3c558d13 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 e0aef7f41d..0bcdfc48df 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 8af7c9dd36..f83f3aeb0b 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 b1c5450b45..ed0b57c2af 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 70afefd41f..57a48517cd 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 c73dc6f40c..95b1c9b160 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 95dbc4face..31b9d36314 100644 --- a/python/tests/test_drawing.py +++ b/python/tests/test_drawing.py @@ -2256,12 +2256,12 @@ def test_no_edges_show_empty(self): # SVG should just be a row of 10 sample nodes svg = ts.draw_svg(time_scale=time_scale, x_lim=[0, ts.sequence_length]) self.verify_basic_svg(svg) - assert svg.count("rect") == 10 # Sample nodes are rectangles - assert svg.count('path class="edge"') == 0 + assert svg.count(" 0 # Should have some singletons svg = ts.draw_svg() self.verify_basic_svg(svg) - assert svg.count("rect") == 10 + assert svg.count(" mutation to plot as ticks on the x axis alternate_dash_positions=None, # Where to alternate the axis from solid to dash + x_regions=None, # A dict of (left, right):label items to place in boxes ): - if not self.x_axis and not self.x_label: + if not self.x_axis: return if alternate_dash_positions is None: alternate_dash_positions = np.array([]) + if x_regions is None: + x_regions = {} dwg = self.drawing axes = self.get_axes() x_axis = axes.add(dwg.g(class_="x-axis")) @@ -791,82 +795,98 @@ def draw_x_axis( transform="translate(0 -11)", text_anchor="middle", ) - if self.x_axis: - if tick_length_upper is None: - tick_length_upper = tick_length_lower - y = rnd(self.plotbox.max_y - self.x_axis_offset) - dash_locs = np.concatenate( - ( - [self.plotbox.left], - self.x_transform(alternate_dash_positions), - [self.plotbox.right], + if len(x_regions) > 0: + regions_group = x_axis.add(dwg.g(class_="x-regions")) + for i, ((left, right), label) in enumerate(x_regions.items()): + if not (0 <= left < right <= self.ts.sequence_length): + raise ValueError( + f"Invalid coordinates ({left} to {right}) for x-axis region" + ) + x1 = self.x_transform(left) + x2 = self.x_transform(right) + y = self.plotbox.max_y - self.x_axis_offset + region = regions_group.add(dwg.g(class_=f"r{i}")) + region.add( + dwg.rect((x1, y), (x2 - x1, self.line_height), class_="r{i}") + ) + self.add_text_in_group( + label, + region, + pos=((x2 + x1) / 2, y + self.line_height / 2), + class_="lab", + text_anchor="middle", ) + if tick_length_upper is None: + tick_length_upper = tick_length_lower + y = rnd(self.plotbox.max_y - self.x_axis_offset) + dash_locs = np.concatenate( + ( + [self.plotbox.left], + self.x_transform(alternate_dash_positions), + [self.plotbox.right], ) - for i, (x1, x2) in enumerate(zip(dash_locs[:-1], dash_locs[1:])): - x_axis.add( - dwg.line( - (rnd(x1), y), - (rnd(x2), y), - class_="ax-skip" if i % 2 else "ax-line", - ) + ) + for i, (x1, x2) in enumerate(zip(dash_locs[:-1], dash_locs[1:])): + x_axis.add( + dwg.line( + (rnd(x1), y), + (rnd(x2), y), + class_="ax-skip" if i % 2 else "ax-line", ) - if tick_positions is not None: - if tick_labels is None or isinstance(tick_labels, np.ndarray): - if tick_labels is None: - tick_labels = tick_positions - tick_labels = create_tick_labels(tick_labels) # format integers - - upper_length = -tick_length_upper if site_muts is None else 0 - ticks_group = x_axis.add(dwg.g(class_="ticks")) - for pos, lab in itertools.zip_longest(tick_positions, tick_labels): - tick = ticks_group.add( - dwg.g( - class_="tick", - transform=f"translate({rnd(self.x_transform(pos))} {y})", - ) - ) - tick.add( - dwg.line((0, rnd(upper_length)), (0, rnd(tick_length_lower))) + ) + if tick_positions is not None: + if tick_labels is None or isinstance(tick_labels, np.ndarray): + if tick_labels is None: + tick_labels = tick_positions + tick_labels = create_tick_labels(tick_labels) # format integers + + upper_length = -tick_length_upper if site_muts is None else 0 + ticks_group = x_axis.add(dwg.g(class_="ticks")) + for pos, lab in itertools.zip_longest(tick_positions, tick_labels): + tick = ticks_group.add( + dwg.g( + class_="tick", + transform=f"translate({rnd(self.x_transform(pos))} {y})", ) - self.add_text_in_group( - lab, - tick, - class_="lab", - # place origin at the bottom of the tick plus a single px space - pos=(0, tick_length_lower + 1), + ) + tick.add(dwg.line((0, rnd(upper_length)), (0, rnd(tick_length_lower)))) + self.add_text_in_group( + lab, + tick, + class_="lab", + # place origin at the bottom of the tick plus a single px space + pos=(0, tick_length_lower + 1), + ) + if not self.omit_sites and site_muts is not None: + # Add sites as vertical lines with overlaid mutations as upper chevrons + for s_id, mutations in site_muts.items(): + s = self.ts.site(s_id) + x = self.x_transform(s.position) + site = x_axis.add( + dwg.g( + class_=f"site s{s.id + self.offsets.site}", + transform=f"translate({rnd(x)} {y})", ) - if not self.omit_sites and site_muts is not None: - # Add sites as vertical lines with overlaid mutations as upper chevrons - for s_id, mutations in site_muts.items(): - s = self.ts.site(s_id) - x = self.x_transform(s.position) - site = x_axis.add( - dwg.g( - class_=f"site s{s.id + self.offsets.site}", - transform=f"translate({rnd(x)} {y})", + ) + site.add(dwg.line((0, 0), (0, rnd(-tick_length_upper)), class_="sym")) + for i, m in enumerate(reversed(mutations)): + mutation_class = f"mut m{m.id + self.offsets.mutation}" + if m.id in self.mutations_outside_tree: + mutation_class += " extra" + mut = dwg.g(class_=mutation_class) + h = -i * 4 - 1.5 + w = tick_length_upper / 4 + mut.add( + dwg.polyline( + [ + (rnd(w), rnd(h - 2 * w)), + (0, rnd(h)), + (rnd(-w), rnd(h - 2 * w)), + ], + class_="sym", ) ) - site.add( - dwg.line((0, 0), (0, rnd(-tick_length_upper)), class_="sym") - ) - for i, m in enumerate(reversed(mutations)): - mutation_class = f"mut m{m.id + self.offsets.mutation}" - if m.id in self.mutations_outside_tree: - mutation_class += " extra" - mut = dwg.g(class_=mutation_class) - h = -i * 4 - 1.5 - w = tick_length_upper / 4 - mut.add( - dwg.polyline( - [ - (rnd(w), rnd(h - 2 * w)), - (0, rnd(h)), - (rnd(-w), rnd(h - 2 * w)), - ], - class_="sym", - ) - ) - site.add(mut) + site.add(mut) def draw_y_axis( self, @@ -1002,7 +1022,8 @@ def __init__( x_label, y_label, y_ticks, - y_gridlines, + x_regions=None, + y_gridlines=None, x_lim=None, max_time=None, min_time=None, @@ -1116,6 +1137,7 @@ def __init__( skipbreaks, tick_length_lower=self.default_tick_length, # TODO - parameterize tick_length_upper=self.default_tick_length_site, # TODO - parameterize + x_regions=x_regions, ) y_low = self.tree_plotbox.bottom if y_axis is not None: @@ -1181,6 +1203,7 @@ def draw_x_axis( tree_is_used, breaks, skipbreaks, + x_regions, tick_length_lower=SvgAxisPlot.default_tick_length, tick_length_upper=SvgAxisPlot.default_tick_length_site, ): @@ -1243,6 +1266,8 @@ def draw_x_axis( site_muts = None # It doesn't make sense to plot sites for "treewise" plots tick_length_upper = None # No sites plotted, so use the default upper tick + if x_regions is not None and len(x_regions) > 0: + raise ValueError("x_regions are not supported for treewise plots") # NB: no background shading needed if x_scale is "treewise" @@ -1258,6 +1283,7 @@ def draw_x_axis( tick_length_upper=tick_length_upper, site_muts=site_muts, alternate_dash_positions=skipregion_pos, + x_regions=x_regions, ) @@ -1286,6 +1312,7 @@ def __init__( y_axis=None, x_label=None, y_label=None, + x_regions=None, y_ticks=None, y_gridlines=None, all_edge_mutations=None, @@ -1478,6 +1505,7 @@ def __init__( tick_length_lower=tick_length_lower, tick_length_upper=tick_length_upper, site_muts=site_muts, + x_regions=x_regions, ) if y_ticks is None: y_ticks = {h: ts.node(u).time for u, h in self.node_height.items()} diff --git a/python/tskit/trees.py b/python/tskit/trees.py index eea322e4dc..e617c7ea8d 100644 --- a/python/tskit/trees.py +++ b/python/tskit/trees.py @@ -1811,6 +1811,7 @@ def draw_svg( symbol_size=None, x_axis=None, x_label=None, + x_regions=None, y_axis=None, y_label=None, y_ticks=None, @@ -1890,6 +1891,10 @@ def draw_svg( plot an X axis. :param str x_label: Place a label under the plot. If ``None`` (default) and there is an X axis, create and place an appropriate label. + :param dict x_regions: A dictionary mapping (left, right) tuples to names. This + draws a box, labelled with the name, on the X axis between the left and + right positions, and can be used for annotating genomic regions (e.g. + genes) on the X axis. If ``None`` (default) do not plot any regions. :param bool y_axis: Should the plot have an Y axis line, showing time (or ranked node time if ``time_scale="rank"``). If ``None`` (default) do not plot a Y axis. @@ -1935,6 +1940,7 @@ def draw_svg( symbol_size=symbol_size, x_axis=x_axis, x_label=x_label, + x_regions=x_regions, y_axis=y_axis, y_label=y_label, y_ticks=y_ticks, @@ -7178,6 +7184,7 @@ def draw_svg( x_axis=None, x_label=None, x_lim=None, + x_regions=None, y_axis=None, y_label=None, y_ticks=None, @@ -7254,6 +7261,10 @@ def draw_svg( information discarded: this means that mutations outside the interval will not be shown. To force display of the entire tree sequence, including empty flanking regions, specify ``x_lim=[0, ts.sequence_length]``. + :param dict x_regions: A dictionary mapping (left, right) tuples to names. This + draws a box, labelled with the name, on the X axis between the left and + right positions, and can be used for annotating genomic regions (e.g. + genes) on the X axis. If ``None`` (default) do not plot any regions. :param bool y_axis: Should the plot have an Y axis line, showing time (or ranked node time if ``time_scale="rank"``. If ``None`` (default) do not plot a Y axis. @@ -7307,6 +7318,7 @@ def draw_svg( x_axis=x_axis, x_label=x_label, x_lim=x_lim, + x_regions=x_regions, y_axis=y_axis, y_label=y_label, y_ticks=y_ticks,