pair_picks.py

import numpy as np
from collections import Counter, defaultdict


def __f(x, i, j):
    '''Consume a state to generate a sequence of haplotypes picked'''
    pairs = []
    y = x.copy()
    while True:

        while not y[i]:
            i += 1
            i %= 4
        y[i] -= 1

        while not y[j]:
            j += 1
            j %= 4
        y[j] -= 1

        pairs.append((i, j))
        if (y == 0).all():
            return sorted(pairs)


def diploid_dist(state, verbose=True):

    # iterate over the starting state and get configurations
    outcomes = []
    for i in range(4):
        for j in range(4):
            outcomes.append(__f(state, i, j))
    if verbose:
        print(outcomes)

    size = len(outcomes) * (sum(state) / 2)
    if verbose:
        print(size)

    # Count up the pairs, turn counts into probabilities
    prob = defaultdict(float)
    for k, c in Counter([x for s in outcomes for x in s]).items():
        # make sure (x,y) and (y,x) are counted together
        ht = tuple(sorted(k))
        prob[ht] += c / size

    assert (np.allclose(sum(prob.values()), 1))

    return prob


print(diploid_dist(np.array([0, 2, 0, 8])))
print(diploid_dist(np.array([2, 2, 3, 3])))