How to check if a tree sequence has any locally unary nodes

[hyanwong]

We want to do this in tsdate (see tskit-dev/tsdate#320). Here's the fastest way I have found (takes about 90 secs on my laptop for a 138MB / 258547-tree ts (chr 2q from the unified genealogy):

def has_locally_unary_nodes(ts):
    for tree, ediff in zip(ts.trees(), ts.edge_diffs()):
        changed = {e.parent for edges in (ediff.edges_out, ediff.edges_in) for e in edges}
        if (tree.num_children_array[list(changed)] == 1).any():
            return True
    return False

@jeromekelleher mentioned in that tskit issue that there might be a speedier way using numba, though.

[jeromekelleher]

I'm surprised messing with the edge diffs makes any difference - out of interest what's the time difference with:

for tree in ts.trees():
    if np.any(tree.num_children_array == 1):
           return True
return False

But yeah, this can definitely be done more quickly using the numba incremental algorithm approach.

[barneyhill]

Yeah double-checked this on the quebec ts (1050794 ts, 564.8 MiB) I average 22 seconds on the has_locally_unary_nodes and 3 minutes with Jerome's method.