Relabel leaf nodes in "canonical" order

[hyanwong]

I have a large number of degree-1 tree sequences (with a single tree), all of which have the same (unlabelled) topology, but where the labelled topologies might differ. E.g.

8.56┊    6    ┊             1.96┊    6    ┊  
    ┊  ┏━┻━┓  ┊                 ┊  ┏━┻━┓  ┊  
1.81┊  5   ┃  ┊             0.67┊  ┃   5  ┊  
    ┊ ┏┻┓  ┃  ┊                 ┊  ┃  ┏┻┓ ┊  
0.50┊ ┃ ┃  4  ┊             0.56┊  4  ┃ ┃ ┊  
    ┊ ┃ ┃ ┏┻┓ ┊                 ┊ ┏┻┓ ┃ ┃ ┊  
0.00┊ 0 2 1 3 ┊             0.00┊ 0 3 1 2 ┊  
  0.00      1.00              0.00      1.00 

I want to renumber the sample nodes so that they all conform to a "canonical" labelling. I can get a canonical ordering using @daniel-goldstein 's nice "rank" method. Here's an example:

tree = msprime.simulate(6, random_seed=13).first()

orig_ranktree = tskit.combinatorics.RankTree.from_tsk_tree(tree)
canonical_ranktree = tskit.combinatorics.RankTree.unrank(tree.num_samples(), (tree.rank()[0], 0))

print(tree.tree_sequence.draw_text())
print(canonical_ranktree.to_tsk_tree().tree_sequence.draw_text())

# formatted side-by-side
4.54┊       10    ┊          
    ┊     ┏━━┻━━┓ ┊          
2.67┊     9     ┃ ┊          4.00┊   10        ┊  
    ┊  ┏━━┻━┓   ┃ ┊              ┊ ┏━━┻━━┓     ┊  
1.72┊  ┃    8   ┃ ┊          3.00┊ ┃     9     ┊  
    ┊  ┃   ┏┻━┓ ┃ ┊              ┊ ┃  ┏━━┻━┓   ┊  
0.59┊  7   ┃  ┃ ┃ ┊          2.00┊ ┃  ┃    8   ┊  
    ┊ ┏┻┓  ┃  ┃ ┃ ┊              ┊ ┃  ┃  ┏━┻┓  ┊  
0.20┊ ┃ ┃  6  ┃ ┃ ┊          1.00┊ ┃  6  ┃  7  ┊  
    ┊ ┃ ┃ ┏┻┓ ┃ ┃ ┊              ┊ ┃ ┏┻┓ ┃ ┏┻┓ ┊  
0.00┊ 0 1 2 4 3 5 ┊          0.00┊ 0 1 2 3 4 5 ┊  
  0.00          1.00           0.00          1.00 

So I need to find out that e.g. sample node 5 in the first tree should be changed to a 0, and samples 0 and 1 in the first tree should be mapped to 1 and 2 (or 2 and 1, it doesn't matter which). How can I get this mapping?

[hyanwong]

I think I can do this by using the order of labels returned by standard traversal though a RankTree. But there's no function that does this. Here's one of my own that I rolled (name updated to leaves as an equivalent to Tree.leaves() after discussion with @daniel-goldstein ) (edit: this method now PRed in #1292)

def leaves(self):
    if self.is_leaf():
        return self.label,
    return [l for c in self.children for l in c.ordered_labels()]
tskit.combinatorics.RankTree.leaves = leaves

Used as follows (note that the internal node numbers are not the same between tree 1 and tree 3, which is to be expected)

tree = msprime.simulate(6, random_seed=13).first()
orig_ranktree = tskit.combinatorics.RankTree.from_tsk_tree(tree)
# I think I don't need this because in the canonical tree all the labels are in order, so I can just use np.arange(num_samples)
canonical_ranktree = tskit.combinatorics.RankTree.unrank(tree.num_samples(), (tree.rank()[0], 0))
# Check that assumption
assert sorted(canonical_ranktree.ordered_labels()) == canonical_ranktree.leaves()

mapping = np.arange(tree.tree_sequence.num_nodes, dtype=tree.tree_sequence.tables.edges.child.dtype)
mapping[orig_ranktree.leaves()] = mapping[canonical_ranktree.leaves()]

# swap the node IDs
tables = tree.tree_sequence.dump_tables()
nodes = tables.nodes
edges = tables.edges
tables.nodes.set_columns(time=nodes.time[mapping], flags=nodes.flags[mapping])
tables.edges.set_columns(left=edges.left, right=edges.right, parent=mapping[edges.parent], child=mapping[edges.child])
tables.sort()
new_ts = tables.tree_sequence()

print("Original", tree.tree_sequence.draw_text(), sep="\n")
print("Canonical", canonical_ranktree.to_tsk_tree().tree_sequence.draw_text(), sep="\n")
print("Original renumbered", new_ts.draw_text(), sep="\n")

Which gives the correct node renumbering in this case:

Original
4.54┊       10    ┊  
    ┊     ┏━━┻━━┓ ┊  
2.67┊     9     ┃ ┊  
    ┊  ┏━━┻━┓   ┃ ┊  
1.72┊  ┃    8   ┃ ┊  
    ┊  ┃   ┏┻━┓ ┃ ┊  
0.59┊  7   ┃  ┃ ┃ ┊  
    ┊ ┏┻┓  ┃  ┃ ┃ ┊  
0.20┊ ┃ ┃  6  ┃ ┃ ┊  
    ┊ ┃ ┃ ┏┻┓ ┃ ┃ ┊  
0.00┊ 0 1 2 4 3 5 ┊  
  0.00          1.00 

Canonical
4.00┊   10        ┊  
    ┊ ┏━━┻━━┓     ┊  
3.00┊ ┃     9     ┊  
    ┊ ┃  ┏━━┻━┓   ┊  
2.00┊ ┃  ┃    8   ┊  
    ┊ ┃  ┃  ┏━┻┓  ┊  
1.00┊ ┃  6  ┃  7  ┊  
    ┊ ┃ ┏┻┓ ┃ ┏┻┓ ┊  
0.00┊ 0 1 2 3 4 5 ┊  
  0.00          1.00 

Original renumbered
4.54┊   10        ┊  
    ┊ ┏━━┻━━┓     ┊  
2.67┊ ┃     9     ┊  
    ┊ ┃  ┏━━┻━┓   ┊  
1.72┊ ┃  ┃    8   ┊  
    ┊ ┃  ┃  ┏━┻┓  ┊  
0.59┊ ┃  7  ┃  ┃  ┊  
    ┊ ┃ ┏┻┓ ┃  ┃  ┊  
0.20┊ ┃ ┃ ┃ ┃  6  ┊  
    ┊ ┃ ┃ ┃ ┃ ┏┻┓ ┊  
0.00┊ 0 1 2 3 4 5 ┊  
  0.00          1.00 

But I'm not 100% sure the logic is correct. I wondered if @daniel-goldstein had any thoughts?

[daniel-goldstein]

This looks right to me! And yes the label ordering for a given topology is a lexicographic ordering, so the 0th labelling will be the input list of labels, left to right, or range(n) I believe if no labels are provided.