Get the edge ids for every mutation in a TS

[hyanwong]

I want to get a list of all mutations above each edge in a tree sequence. I guess the easiest way is to create a column of length ts.num_mutations with the edge ID for each corresponding mutation in it. Here's how we currently count the number of mutations on each edge in tsdate in an array of length ts.num_edges, but it doesn't return the mutation IDs. Presumably there should be a nice numpy based way of getting the edge IDs instead?

    def get_mut_edges(ts):
        """
        Get the number of mutations on each edge in the tree sequence.
        """
        edge_diff_iter = ts.edge_diffs()
        right = 0
        edges_by_child = {}  # contains {child_node:edge_id}
        mut_edges = np.zeros(ts.num_edges, dtype=np.int64)
        for site in ts.sites():
            while right <= site.position:
                (left, right), edges_out, edges_in = next(edge_diff_iter)
                for e in edges_out:
                    del edges_by_child[e.child]
                for e in edges_in:
                    assert e.child not in edges_by_child
                    edges_by_child[e.child] = e.id
            for m in site.mutations:
                # In some cases, mutations occur above the root
                # These don't provide any information for the inside step
                if m.node in edges_by_child:
                    edge_id = edges_by_child[m.node]
                    mut_edges[edge_id] += 1
        return mut_edges

[hyanwong]

Now that the mutations have an edge ID associated with them (see #2279) we should be able to do this with a simple loop. For example

def new_get_num_muts_for_edges(ts):
    mut_edges = np.zeros(ts.num_edges, dtype=int)
    for m in ts.mutations():
        mut_edges[m.edge] += 1
    return mut_edges

def new_get_muts_for_edges(ts):
    return np.array([m.edge for m in ts.mutations()])