implement an `allele_counter` about tskit HOT 9 CLOSED

import msprime
from tests.tsutil import jukes_cantor

from collections import Counter


def count_alleles(tree, site):
    n = tree.num_tracked_samples()
    U = {site.ancestral_state: n}
    for mut in site.mutations:
        if mut.derived_state not in U:
            U[mut.derived_state] = 0
        U[mut.derived_state] += tree.num_tracked_samples(mut.node)
        parent_state = get_derived_state(site, mut.parent)
        if parent_state not in U:
            U[parent_state] = 0
        U[parent_state] -= tree.num_tracked_samples(mut.node)
    zeros = [a for a in U if U[a] == 0]
    for a in zeros:
        del U[a]
    return U


def get_derived_state(site, mut_id):
    """
    Find the derived state of the mutation with id `mut_id` at site `site`.
    """
    if mut_id == msprime.NULL_NODE:
        state = site.ancestral_state
    else:
        for m in site.mutations:
            if m.id == mut_id:
                state = m.derived_state
    return state

def count_alleles_no_parent(tree, site):
    n = tree.num_tracked_samples()
    U = {site.ancestral_state: n}
    mut_nodes = [mut.node for mut in site.mutations]
    # this version can't handle more than one mut per node
    assert len(mut_nodes) == len(set(mut_nodes))
    for mut in site.mutations:
        if mut.derived_state not in U:
            U[mut.derived_state] = 0
        U[mut.derived_state] += tree.num_tracked_samples(mut.node)
        p = tree.parent(mut.node)
        while ((p != msprime.NULL_NODE) and (p not in mut_nodes)):
            p = tree.parent(p)
        if p == msprime.NULL_NODE:
            parent_state = site.ancestral_state
        else:
            parent_state = site.mutations[mut_nodes.index(p)].derived_state
        if parent_state not in U:
            U[parent_state] = 0
        U[parent_state] -= tree.num_tracked_samples(mut.node)
    zeros = [a for a in U if U[a] == 0]
    for a in zeros:
        del U[a]
    return U


class AlleleCounter(object):

    def __init__(self, ts, samples, with_parent=True):
        self.trees = ts.trees(tracked_samples=samples, sample_counts=True)
        if with_parent:
            self.ac_fun = count_alleles
        else:
            self.ac_fun = count_alleles_no_parent

    def __iter__(self):
        """
        Returns a generator of dictionaries of the form
           allele : number of tracked samples in tree with that allele
        ... one for each site.  Should not contain zeros.
        """
        for t in self.trees:
            for s in t.sites():
                if len(s.mutations) > 0:
                    yield self.ac_fun(t, s)


def hap_counter(ts, samples):
    """
    Quick version for testing.
    """
    haps = list(ts.haplotypes())
    for k in range(ts.num_sites):
        yield Counter(haps[j][k] for j in samples)


ts0 = msprime.simulate(10, recombination_rate=2)
ts = jukes_cantor(ts0, num_sites=10, multiple_per_node=False, mu=1.0)

samples = list(range(4))
ac = AlleleCounter(ts, samples)
ac_np = AlleleCounter(ts, samples)
hc = hap_counter(ts, samples)

for x, xp, y in zip(ac, ac_np, hc):
    print("---------")
    print("AC", x)
    # print("HC", y)
    assert sum(x.values()) == len(samples)
    assert set(x) == set(xp)
    assert set(x) == set(y)
    for a in x:
        assert x[a] == y[a]
        assert x[a] == xp[a]