Returning node x,y locations in a plotted tree

[hyanwong]

For creating extra annotations on plots, it would be helpful to be able to access the x and y locations of nodes in a plotted tree. I think these are stored in node_x_coord_map and node_y_coord_map, so should be trivial to return.

One discrete way to to this would be to add an attribute (or method) to the SVGstring class that returned a dict of this data, e.g.

svg_string = ts.first().draw_svg()
positions = svg_string.positions
print(positions[node_id])  # prints tuple of (x, y)

I wonder if this is an abuse of the SVGstring class though, which is simply a subclass of the python String?

[jeromekelleher]

A return_node_positions argument might be more obvious I think, in which case we return a tuple. I guess if we changed the return type to SVGTree then it would be quite natural to add these types of attributes to it , though.

[benjeffery]

SVGTree sounds like a nice idea. Still can be a subclass of string, but lets us add stuff in non-intrusive way.

[hyanwong]

Yes, I prefer not to change the return type depending on parameters in the function, if possible.

[jeromekelleher]

Doesn't even need to be a subclass of string I guess - once it has a __str__ then it quacks like a string?

[hyanwong]

True, although it can be helpful to do e.g. find, search and replace, encode(), etc. I guess the point is, does it seem like a sort of string, or something else?

[jeromekelleher]

Well it seems like a something else if we're going to be adding a bunch of attributes about a plotted tree to it.

Does it break our tests much if we return something that just has a __str__ method?

[hyanwong]

Well it seems like a something else if we're going to be adding a bunch of attributes about a plotted tree to it.

I tend to agree.

Does it break our tests much if we return something that just has a __str__ method?

Yes, quite a lot. We have many things in the tests that fail, like:

svg_string = ts.draw_svg()
assert "Genome position" not in svg_string
assert svg_string.count('class="site ') == 0
svg_string[svg_string.find("</style>") :]
assert re.search(rf"t{tree.index}[^\d]", svg_string) is not None  # re.search expects svg_string to be a string-like object
xml.dom.minidom.parseString(svg_string) # also expects svg_string to be a string-like object

[hyanwong]

I figured out a nicer way to do this (although it uses internal undocumented classes): just make a SvgTree and access
svgtree.node_x_coord and svgtree.timescaling.transform. This is much more flexible, as you can convert arbitrary times (e.g. for putting horizontal lines across the plot, even in log-time plots) and with #3197 you can then just call svgtree.draw() to get the final output. Given it's such a specialised case, I think this is a perfectly acceptable workaround, and I'll put a (suitably warning-containing) example in the tutorials.

def node_positions(svgtree):
    # the internal `node_x_coord` dict maps node IDs to horizontal pos
    x = svgtree.node_x_coord
    # the internal `timescaling.transform` function returns vertical positions
    node_times = svgtree.ts.nodes_time  # grab node times from the associated ts
    if svgtree.time_scale == "rank":  # NB: the "rank" scale doesn't use raw ages
        within_tree_node_times = np.unique(node_times[svgtree.tree.preorder()])
        node_times = np.searchsorted(within_tree_node_times, node_times)
    y = svgtree.timescaling.transform(node_times)
    return x, y

# Create an SvgTree object using the same parameters as for draw_svg()
internal_obj = tskit.drawing.SvgTree(
    ts_small.first(),
    canvas_size=(300, 200),
    time_scale="log_time",
    y_axis=True,
    y_ticks=[0, 10, 100, 1000],
    node_labels={},
)

horiz_node_pos, vert_node_pos = node_positions(internal_obj)

x10, y10 = horiz_node_pos[10], vert_node_pos[10]  # Node 10
x8, y8 = horiz_node_pos[8], vert_node_pos[8]  # Node 8

internal_obj.preamble = (
   f'<circle cx="{x10}" cy="{y10}" r="10" stroke="blue" fill="white" />'
   f'<text x="{x10+8}" y="{y10-8}" fill="blue">This is node 10!</text>'
   f'<line x1="{x8}" y1="{y8}" x2="{x8 + 80}" y2="{y8 - 20}" stroke="red" />'
   f'<text x="{x8 + 82}" y="{y8 - 22}" fill="red">This is node 8!</text>'
)
internal_obj.draw()