Add num_children to Tree arrays

[jeromekelleher]

There are lots of times when we want to ask questions about the arity (number of children) of the nodes in a tree. In particular, many of the operations discussed in #2245 required this as input. While it's easy enough to compute in a linear pass through the nodes in the tree, it's also something we can maintain during the tree transition process very cheaply. The basic algorithm just needs to increment/decrement a count for the parent node:

ts = msprime.sim_ancestry(10, random_seed=2, sequence_length=100, recombination_rate=0.1)

num_children = np.zeros(ts.num_nodes, dtype=int)

for tree, (_, edges_out, edges_in) in zip(ts.trees(), ts.edge_diffs()):
    for e in edges_out:
        num_children[e.parent] -= 1
    for e in edges_in:
        num_children[e.parent] += 1

    tree_num_children = [tree.num_children(u) for u in range(ts.num_nodes)]
    assert tree_num_children == list(num_children)

So, the computational cost would essentially be free, and it's just a question about whether we want to add the extra O(num_nodes) space cost to the Tree struct by default or not.

My feeling is that this is quite a basic thing and it'll be annoying to have to specify a Tree option if you want to do things like compute stats based on node arity, so it's probably worth the extra bit of memory.

If we make it a standard part of the tree building process, then we can export it as a numpy array to Python (like the other tree arrays) and then things like computing the arity histogram, is_binary etc can be done efficiently in Python.

[jeromekelleher]

Another advantage of this is that we'd be able to efficiently count the number of roots, since num_children[tree.virtual_root] should be tree.num_roots.

[hyanwong]

I definitely think this is worth it for the storage cost.

[andrewkern]

+1

i came across a use case for this just earlier today.

[benjeffery]

I've added this to the release after the imminent one, there are quite a few things on it already but we can prioritise after 0.5.0.

[jeromekelleher]

Closed in #2316