Skip to content

Instantly share code, notes, and snippets.

@ivan-krukov
Created August 29, 2019 20:11
Show Gist options
  • Save ivan-krukov/238d3f5c2b52846c1e6157da1305e0bb to your computer and use it in GitHub Desktop.
Save ivan-krukov/238d3f5c2b52846c1e6157da1305e0bb to your computer and use it in GitHub Desktop.
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])))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment