Number of nodes (or equivalent) under a node

[hyanwong]

As well as having a postorder array It is often useful in phylogenetic algorithms to also have the total number of nodes under each element of the postorder array (or an equivalent number like the index into the postorder array to skip to in order to bypass the subtree). For example, this can be used to create a "nested set" representation of a tree, or could be used to iterate over the postorder array but skip certain subtrees. Is there a way of getting this info somehow, without traversing through each subtree in turn?

[jeromekelleher]

You can make the array in a single pass through the postorder array, right?

total_nodes = np.zeros_like(tree.tree_sequence.num_nodes)
for u in tree.postorder():
     for v in tree.children(u):
           total_nodes[u] += total_nodes[v]
     total_nodes[u] += 1

Or such like?

[hyanwong]

Here's an alternative version of the code above. I'm not sure if it's more efficient to store the counts as a Counter() or a (potentially large) array of all the nodes in the tree sequence (including those not in the tree):

from collections import Counter
nodes_in_subtree = Counter()  # keyed by the ID of the node that defines the subtree
for u in tree.nodes(order="postorder"):
     for v in tree.children(u):
          nodes_in_subtree[u] += nodes_in_subtree[v]
     nodes_in_subtree[u] += 1  # Also count the node itself

[hyanwong]

FWIW the second version with Counter is a fraction quicker on a 1M tip tree (but both take about 1 second). It makes no difference in that case whether we use the tree.postorder() array or the tree.nodes(order="postorder") iterator.

[benjeffery]

If this is a commonly used quantity we could consider adding it to the tree arrays - num_descendants similar to num_children?

[hyanwong]

I think it is worth considering if people who do lots of phylogenetic algorithms would use it. So let's leave this discussion open and point to it e.g. from the phylogenetics / tree analysis tut (https://tskit.dev/tutorials/analysing_trees.html), with a comment that if this would be useful, to add to the discussion. We already do this e.g. in the getting started tutorial:

[jeromekelleher]

It's not worth it I think, certainly not worth storing for every tree. There's a whole bunch of related quantities you might want to sum post-order wise, so this would be better off as tutorial fodder.

[hyanwong]

Here's a version which is twice as fast as the original answer, and returns correct values for the virtual root (hence has an array of size num_nodes + 1:

total_nodes = np.ones(tree.tree_sequence.num_nodes + 1, dtype=ts.edges_child.dtype)
total_nodes[tskit.NULL] = 0  # this is the virtual root, and should not count itself
for u in tree.postorder():
    p = tree.parent(u)  # will be -1 == tskit.NULL for a root, so virtual root correctly updated
    total_nodes[p] += total_nodes[u]