ARG traversal: iterate over edges by child node

[hyanwong]

Traversing from bottom to top of a tree sequence, visiting the edges that belong to each parent, is easy, because edges are ordered by parent time (itertools.groupby(ts.edges(), operator.attrgetter("parent")) should do it). But It is often useful to traverse from top to bottom in a tree sequence, visiting the all the edges that nodes in age order. For example, we do this in tsdate, here.

Here's a nice compact function to do this:

import itertools
import operator
def edges_by_child(ts, order="timeasc"):
    """
    Return an itertools.groupby object of edges grouped by child in order
    of the time of the child.
    """
    if order not in {"timeasc", "timedesc"}:
        raise ValueError('order must be "timeasc" or "timedesc"')
    flip = 1 if order == "timeasc" else -1
    it = (ts.edge(u) for u in np.lexsort((ts.edges_child, ts.nodes_time[ts.edges_child] * flip)))
    return itertools.groupby(it, operator.attrgetter("child"))

An example of using it:

simple_ts = msprime.simulate(4, recombination_rate=1, random_seed = 1)
for u, edges in edges_by_child(simple_ts):
    print("child", u)
    print(list(edges))