def admits(T, mutation, train):
N = coalescent.num_leaves(T)
ones_below = 0
ones_above = 0
zeros_below = 0
zeros_above = 0
for i in range(N):
if i in T.below[mutation[0]]:
if train[i] == 0:
zeros_below += 1
elif train[i] == 1:
ones_below += 1
else:
if train[i] == 0:
zeros_above += 1
elif train[i] == 1:
ones_above += 1
if ((zeros_below == 0) and (ones_above == 0)) or ((zeros_above == 0) and (ones_below == 0)):
return True
else:
return False
def impute(T, edge, train):
N = coalescent.num_leaves(T)
imputed = zeros((N), dtype=int)-nan
for i in where(isnan(train))[0]:
fix = 1
found = False
for j in T.below[edge[0]]:
if not(isnan(train[j])):
if found:
assert train[j] == fix
if train[j] == 0:
fix = 0
elif train[j] == 1:
fix = 1
else:
assert False
found = True
if i in T.below[edge[0]]:
imputed[i] = fix
else:
imputed[i] = 1-fix
return imputed
def observe(T, mutation, p):
N = coalescent.num_leaves(T)
data = zeros((N), dtype=int)
for i in T.below[mutation[0]]:
data[i] = 1
train = zeros((N), dtype=int)+nan
test = zeros((N), dtype=int)+nan
for i in range(N):
if (1-p) > rand():
train[i] = data[i]
else:
test[i] = data[i]
for i in range(N):
if train[i] == 1:
assert i in T.below[mutation[0]]
return (test, train)
