Displaying a tanglegram (i.e. "cophylogeny plot")

[hyanwong]

I have put together a simple draw_tanglegram function for comparing trees, which might be useful to some people:

import msprime
ts = msprime.sim_ancestry(7, sequence_length=1000, recombination_rate=0.001, random_seed=123)
display(draw_tanglegram(ts.at_index(0), ts.at_index(1)))

import tsconvert
ts1 = tsconvert.from_newick("((B:1,C:1):5,(A:5,((D:2,E:2):2,(F:3,G:3):1):1):1);", min_edge_length=1, span=1000)
ts2 = tsconvert.from_newick("((A:1,B:1):5,(C:5,((D:2,E:2):2,(G:3,F:3):1):1):1);", min_edge_length=1, span=1000)
display(draw_tanglegram(ts1.first(), ts2.first(), metadata_label='name', height=200))

Here's the code:

def draw_tanglegram(tree1, tree2, metadata_label=None, width=400, height=500):
    """
    Draw a "tanglegram" that compares the topology leading to leaf nodes on
    two trees, by drawing them facing each other, and plotting lines
    between the same leaves in each tree. Leaves are considered the same
    if they have the same value for `node.metadata[metadata_label]`, or if
    `use_metadata_label` is None, they are compared using leaf node ID.
    """
    w = width / 2  # width (tree height) of one of the plotted trees, after 90° rotation
    md_lab = metadata_label  # shorthand
    ts1 = tree1.tree_sequence
    ts2 = tree2.tree_sequence
    node_labels = {
        u: u if md_lab is None else ts2.node(u).metadata.get(md_lab, '')
        for u in tree2.nodes()
    }
    svg = tree2.draw_svg(
        node_labels=node_labels,
        root_svg_attributes={'class': 'rgt_tree', 'x': w},
        size=(height, w),  # will be rotated, so switch height & width
        canvas_size=(w, height),
    )
    # Relabel the leaf nodes of the tree sequence associated with tree 1,
    # This means that the minlex order ends up reversing the plotted node order
    node_map = np.arange(tree1.tree_sequence.num_nodes)
    leaves = np.array([u for u in tree1.nodes(order="minlex_postorder") if tree1.is_leaf(u)])[::-1]
    node_map[np.sort(leaves)] = leaves
    tmp_ts = tree1.tree_sequence.subset(node_map, remove_unreferenced=False)
    tree1 = tmp_ts.at_index(tree1.index)

    # Assume samples are tips
    node_order1 = [
        u if md_lab is None else ts1.node(u).metadata[md_lab] for u in leaves
    ]
    i = 0
    node_order2 = {}
    for u in tree2.nodes(order="minlex_postorder"):
        if tree2.is_leaf(u):
            key = u if md_lab is None else tree2.tree_sequence.node(u).metadata[md_lab]
            node_order2[key] = i
            i += 1
    
    margin_lft, margin_rgt = 20, 20  # default margins in tskit draw_svg()
    tip_spacing = (height - margin_lft - margin_rgt) / (len(node_order1))
    tip_pos = np.array([margin_lft + tip_spacing/2 + i*tip_spacing for i in range(len(node_order1))])
    
    lines = [
        f'<line x1="{tip_pos[i]}" y1="{w-10}" x2="{height-tip_pos[node_order2[tip]]}" y2="{w+10}" stroke="blue" />'
        for i, tip in enumerate(node_order1)
        if tip in node_order2
    ]

    node_labels = {
        u: node_map[u] if md_lab is None else tree1.tree_sequence.node(u).metadata.get(md_lab, '')
        for u in tree1.nodes()
    }
    return tree1.draw_svg(
        preamble=svg + '<g class="tangle_lines">' + ''.join(lines) + '</g>',
        size=(height, w),
        canvas_size=(w*2, height),
        node_labels=node_labels,
        root_svg_attributes={'class': 'lft_tree'},
        style=(
            ".lft_tree > g.tree, .lft_tree > g.tangle_lines {transform: translate(0, " + str(height) + "px) rotate(-90deg)}" +
            ".lft_tree > g.tree .leaf > .lab {text-anchor: start; transform: rotate(90deg) translate(6px)}" +
            ".rgt_tree > g.tree {transform: translate(" + str(w) + "px, 0) rotate(90deg)}" +
            ".rgt_tree > g.tree .leaf > .lab {text-anchor: end; transform: rotate(-90deg) translate(-6px)}"
        )
    )

[jeromekelleher]

Cool! So there's no attempt to reorder the tips here to optimise, just take both in minlex and draw the lines?

[hyanwong]

No, but I'm experimenting with using tanglegram.untangle from https://github.com/schlegelp/tanglegram to do this. Slightly annoyingly it requires the trees in script linkage matrix format, but I have found some conversion code to do this.

[hyanwong]

I think this code might be nice to add to the examples at the end of the viz tutorial. I'll get round to that at some point, but don't have time at the moment.

[petrelharp]

Wow, this is very cool.

FWIW the node matching capability in tscompare might be helpful in untangling? (This gives a node-to-node "best match" for internal nodes - tree sequence-wide, which you don't need here? - which seems useful.

[hyanwong]

Here's the code for running the untangling algorithm(s) described in tanglegram.untangle (paper at https://academic.oup.com/bioinformaticsadvances/article/2/1/vbac014/6539778). It would be good to find an implementation of ShUnTan

from Bio import Phylo
import numpy as np
import io
from scipy.cluster.hierarchy import is_valid_linkage
import tanglegram

def newick_string_to_linkage_matrix(newick_str):
    """
    Convert a newick to a scipy.cluster.hierarchy linakes matrix. Returns the
    matrix and an array of labels
    """
    # definition of linkage matrix: 
    # https://docs.scipy.org/doc/scipy/reference/generated/scipy.cluster.hierarchy.linkage.html
    
    # create tree object via Biopython (Bio.Phylo)
    tree = next(Phylo.parse(io.StringIO(newick_str), format="newick"))
    # calculate tree height
    tree_height = max(tree.distance(c) for c in tree.find_clades(terminal=True))
    # add comment with id and span of terminal nodes:
    id_map = {}
    for i, c in enumerate(tree.find_clades(terminal=True)):
        c.comment = (i, 1) 
        id_map[i] = c.name
    # ancestor list orderred by distance from the root
    anc_lst = []
    for c in tree.find_clades(terminal=False):
        d = tree.distance(c)
        anc_lst.append((c, list(c), d))
    anc_lst.sort(key=lambda x:x[2], reverse=True)
    # running number of node
    nodes = len(list(tree.find_clades(terminal=True)))
    lnk_lst = []
    for anc,children, anc_d in anc_lst:
        n_children = len(children)
        assert n_children>=2
        child1 = children[0]
        for child2 in children[1:]:
            id1, n_leaves1 = child1.comment
            id2, n_leaves2 = child2.comment
            total_leaves = n_leaves1 + n_leaves2
            anc.comment = (nodes, total_leaves)
            distance = tree_height - anc_d
            nodes += 1
            row = [id1, id2, distance, total_leaves]
            lnk_lst.append(row)
            child1 = anc
    return np.array(lnk_lst), [id_map[i] for i in range(len(id_map))]

def to_newick_tree(Z, leaf_names):
    """
    Convert a linkage matrix into a 
    string in `Newick <https://en.wikipedia.org/wiki/Newick_format>`_ format.
    Newick is a commonly-used format to represent trees. It is used by packages 
    such as `ete3 <http://etetoolkit.org/>`_  for tree manipulation and visualisation.
    Parameters
    ----------
    Z : ndarray
        The linkage matrix in proper form (see the `linkage`
        function documentation).
    leaf_names : list
        List of strings giving the names of the leaves of the tree. The order
        of strings is assumed to be the the same as the order of the data points 
        given to `linkage`.
    Returns
    -------
    tree : string
    See Also
    --------
    linkage, to_tree
    Examples
    --------
    >>> import numpy as np
    >>> from scipy.cluster import hierarchy
    >>> Z = np.asarray([[0, 1, 3.0, 2], [3, 2, 4.0, 3]], dtype=np.double)
    >>> leaf_names = ["l0", "l1", "l2"]
    >>> hierarchy.to_newick_tree(Z, leaf_names)
    ((l0:3.0,l1:3.0):1.0,l2:4.0);
    """
    Z = np.asarray(Z, order='c')

    # Number of original objects is equal to the number of rows plus 1.
    n = Z.shape[0] + 1
    n_leaves = len(leaf_names)
    if n != n_leaves:
        raise ValueError(f"Expected {n} leaf names, got {n_leaves}")

    is_valid_linkage(Z, throw=True, name='Z')

    newick_intermediates = leaf_names + [None] * Z.shape[0]
    cluster_dists = [0] * (n + Z.shape[0])
    for i, row in enumerate(Z):
        dist = row[2]
        fi = int(row[0])
        fj = int(row[1])
        cdi = dist - cluster_dists[fi]
        cdj = dist - cluster_dists[fj]
        newick_subtree = (
                f"({newick_intermediates[fi]}:{cdi},"
                f"{newick_intermediates[fj]}:{cdj})"
                )
        newick_intermediates[i + n] = newick_subtree
        cluster_dists[i + n] = dist
        newick_intermediates[fi] = None
        newick_intermediates[fj] = None
    return newick_intermediates[i + n] + ";"

def untangle_newicks(newick1, newick2, method, **kwargs):
    Z1, labels1 = newick_string_to_linkage_matrix(newick1)
    Z2, labels2 = newick_string_to_linkage_matrix(newick2)
    untangled1, untangled2 = tanglegram.untangle(Z1, Z2, labels1, labels2, method, **kwargs)
    return to_newick_tree(untangled1, labels1), to_newick_tree(untangled2, labels2)

The untangle_newicks function can be combined with tskit like this:

import string
import msprime
import tskit

# make a tree sequence. We'll compare the k'th tree with the first tree
ts = msprime.simulate(10, recombination_rate=10)
tables = ts.dump_tables()
tables.nodes.metadata_schema = tskit.MetadataSchema.permissive_json()
tables.nodes.packset_metadata([
tables.nodes.metadata_schema.validate_and_encode_row(
            {'name': string.ascii_uppercase[nd.id]} if nd.is_sample() else {}
) for nd in ts.nodes()])
ts = tables.tree_sequence()

k=20
nwk1, nwk2 = untangle_newicks(
    ts.at_index(0).as_newick(node_labels={nd.id: nd.metadata.get('name', '') for nd in ts.nodes()}),
    ts.at_index(k).as_newick(node_labels={nd.id: nd.metadata.get('name', '') for nd in ts.nodes()}),
    method="step2side"
)

display(draw_tanglegram(ts.at_index(0), ts.at_index(k), 'name', width=800, height=300))
display(draw_tanglegram(tsconvert.from_newick(nwk1).first(), tsconvert.from_newick(nwk2).first(), 'name', width=800, height=300))

I'm unclear if this produces as good a result as the java Dendroscope untangling algorithm, but it is linked to a good paper, and is certainly easier to use. Probably we would want to combine the step2side method with a couple of different random start orders.

[jeromekelleher]

Note there's some code here for converting linkage matrix to tskit-like format. Perhaps this would be something to move over to tsconvert?

[hyanwong]

Oh, very neat. I didn't know about that. tsconvert seems like a good home. I assume at some point (but not yet) we will put that repo on PyPI.

[jeromekelleher]

Maybe, I don't know. Might as well move the linkage matrix code in there anyway so it's findable.