Efficiently calculating number of mutations per site

[hyanwong]

What's an efficient way to calculate an array of the number of mutations per site? Obviously I can loop through ts.sites() counting up the size of the mutations list, or do np.unique(ts.mutations_site, return_counts=True), but since all the mutations are sorted by site ID, there is presumably a more efficient way which doesn't sort?

I quite often want to get a boolean array selecting only those sites that have exactly one mutation, so a quick shorthand would be useful. AI suggests this:

def single_mutation_sites(ts):
    diff_mask = np.diff(ts.mutations_site) != 0
    return ts.mutations_site[np.concatenate(([True], diff_mask)) & np.concatenate((diff_mask, [True]))]

Which I think is right?

[nspope]

I use numpy.bincount(ts.mutations_site, minlength=ts.num_sites)

[jeromekelleher]

Yep, that's definitely the way to do it.

[hyanwong]

Maybe I'll add that to the docs somewhere then. It's a useful one-liner.