def construct_ej_es_string(tree, times, leaf_keys, final_pop_size=1.0):
s_times=sorted([(v,k) for k,v in times.items()])
dic_of_lineages={(key,0):(n+1) for n,key in enumerate(leaf_keys)}
#print dic_of_lineages
population_count=len(dic_of_lineages)
res_string=''
for time,key in s_times:
if key=='r':
i,j=get_affected_populations(dic_of_lineages, get_real_children_root(tree, key))
res_string+='-ej '+str(time)+' '+str(i)+' '+str(j)+' '
dic_of_lineages[(key,0)]=j
res_string+='-en '+str(time)+' '+str(dic_of_lineages[(key,0)])+' '+str(final_pop_size)+' '
break
node=tree[key]
if node_is_coalescence(node):
i,j=get_affected_populations(dic_of_lineages, get_real_children_root(tree, key))
res_string+='-ej '+str(time)+' '+str(i)+' '+str(j)+' '
dic_of_lineages[(key,0)]=j
if node_is_admixture(node):
population_count+=1
i=get_affected_populations(dic_of_lineages, get_real_children_root(tree, key))[0]
res_string+='-es '+str(time)+' '+str(i)+' '+str(1.0-node[2])+' '
dic_of_lineages[(key,0)]=i
dic_of_lineages[(key,1)]=population_count
return res_string
def adjust_admixture_proportions(rTree_structure, vertice_dictionary,
admixtures):
for key, node in rTree_structure.items():
if node_is_admixture(node):
## We are now in the situation that
## source
## / \
## \ / \ /
## node b_child_1
## | \
## | |
## a_child_1 b_child_2
## / \ / \
## \ / |
## a_child_2 ...
## | |
## ... b_child_m
## | / \
## a_child_k |
## / \ d_child_1
## |
## c_child_1
## / \
##And we want to find the c_child_1 and d_child_1
#because the admixtures-dictionary contains the admixture proportion with the item
#(c_child_1_N,d_child_1):adm_proportion.
admixture_parent = node[1]
c_child_1 = find_first_non_admix(key, rTree_structure)
d_child_1 = find_first_non_admix(admixture_parent, rTree_structure)
c_child_1_N = vertice_dictionary[c_child_1].N_name
d_child_1_N = vertice_dictionary[d_child_1].N_name
rTree_structure[key][2] = 1.0 - admixtures[(d_child_1_N,
c_child_1_N)]
return rTree_structure
def _get_removable_admixture_branches(tree):
res = []
for key, node in tree.items():
if node_is_admixture(node):
if _check_node(tree, key, 0):
res.append((key, 0))
if _check_node(tree, key, 1):
res.append((key, 1))
return res
def admixes_are_sadmixes_popz(tree):
pops=get_populations(tree)
#print pops
for key,node in tree.items():
if node_is_admixture(node):
sadmix_bool=admix_is_sadmix(tree, (key,1), pops)
if not sadmix_bool:
return False
return True
def pretty_print(tree):
keys, vals = tree.keys(), tree.values()
chosen_indexes = []
print '{ '
for key, val in zip(keys, vals):
if node_is_leaf_node(val):
print ' ' + key + ': ' + str(val)
print ' ,'
for key, val in zip(keys, vals):
if node_is_admixture(val):
print ' ' + key + ': ' + str(val)
print ' ,'
for key, val in zip(keys, vals):
if node_is_coalescence(val):
print ' ' + key + ': ' + str(val)
print '}'
def pretty_string(tree):
keys, vals = tree.keys(), tree.values()
res = ''
res += '{ ' + '\n'
for key, val in zip(keys, vals):
if node_is_leaf_node(val):
res += ' ' + key + ': ' + str(val) + '\n'
res += ' ,' + '\n'
for key, val in zip(keys, vals):
if node_is_admixture(val):
res += ' ' + key + ': ' + str(val) + '\n'
res += ' ,' + '\n'
for key, val in zip(keys, vals):
if node_is_coalescence(val):
res += ' ' + key + ': ' + str(val) + '\n'
res += '}'
return res
def first_admixture_proportion_custom_summary(Rtree, **kwargs):
'''
This is an example of a custom function. It returns the admixture proportion of the immediate ancestors of population s1 and s3.
The function is being read by downstream_analysis_parser because it ends with '_custom_summary'.
:param Rtree: A tree
:param kwargs: all the parameters passed to all summary functions that are not used by this function
:return: a float. the admixture proportion.
'''
#print 'entering first admixture proportion'
for key, node in Rtree.items():
if node_is_admixture(
node): #checking for the pattern s1/ nX \ a1 \s2 /a1 / nY \ s3
p1, p2 = get_parents(node)
p1_other_children = get_sibling_on_parent_side(Rtree, p1, key)
p2_other_children = get_sibling_on_parent_side(Rtree, p2, key)
if 's1' in p1_other_children and 's3' in p2_other_children:
return node[2]
if 's3' in p1_other_children and 's1' in p2_other_children:
return 1.0 - node[2]
else:
return '.'
def construct_ej_en_es_string(tree, times, leaf_keys, final_pop_size=1.0):
s_times=sorted([(v,k) for k,v in times.items()])
dic_of_lineages={(key,0):(n+1) for n,key in enumerate(leaf_keys)}
#print dic_of_lineages
population_count=len(dic_of_lineages)
res_string=''
add_on=1e-6
for time,key in s_times:
if key=='r':
i,j=get_affected_populations(dic_of_lineages, get_real_children_root(tree, key))
res_string+='-ej '+str(time)+' '+str(i)+' '+str(j)+' '
dic_of_lineages[(key,0)]=j
pop_size=final_pop_size
res_string+='-en '+str(time+add_on)+' '+str(dic_of_lineages[(key,0)])+' '+str(pop_size)+' '
break
node=tree[key]
if node_is_leaf_node(node):
pop_size=calculate_pop_size(node[3])
res_string+='-en '+str(time)+' '+str(dic_of_lineages[(key,0)])+' '+str(pop_size)+' '
if node_is_coalescence(node):
i,j=get_affected_populations(dic_of_lineages, get_real_children_root(tree, key))
res_string+='-ej '+str(time)+' '+str(i)+' '+str(j)+' '
dic_of_lineages[(key,0)]=j
pop_size=calculate_pop_size(node[3])
res_string+='-en '+str(time+add_on)+' '+str(dic_of_lineages[(key,0)])+' '+str(pop_size)+' '
if node_is_admixture(node):
population_count+=1
i=get_affected_populations(dic_of_lineages, get_real_children_root(tree, key))[0]
res_string+='-es '+str(time)+' '+str(i)+' '+str(node[2])+' '# WRONG EARLIER VERSION: str(1.0-node[2])+' '
dic_of_lineages[(key,0)]=i
dic_of_lineages[(key,1)]=population_count
pop_size1=calculate_pop_size(node[3])
pop_size2=calculate_pop_size(node[4])
res_string+='-en '+str(time+add_on)+' '+str(dic_of_lineages[(key,0)])+' '+str(pop_size1)+' '
res_string+='-en '+str(time+add_on)+' '+str(dic_of_lineages[(key,1)])+' '+str(pop_size2)+' '
return res_string
def produce_p_matrix(tree,
nSNP,
clip=True,
middle_start=False,
allele_dependent=True,
fixed_in_outgroup='out'):
ps = uniform.rvs(size=nSNP)
if middle_start:
ps = ps * 0.2 + 0.4
else:
ps = ps * 0.9998 + 0.0001
if clip:
if allele_dependent:
add_n = add_noise
else:
add_n = add_noise3
else:
add_n = add_noise2
ps2 = deepcopy(ps)
p_org = deepcopy(ps)
res = {}
(child_key1, child_branch1, l1), (child_key2, child_branch2,
l2) = find_rooted_nodes(tree)
print find_rooted_nodes(tree)
if fixed_in_outgroup:
if child_key1 == fixed_in_outgroup:
tree[child_key2][child_branch2 + 3] = l1 + l2
tree[child_key1][child_branch1 + 3] = 0
else:
tree[child_key2][child_branch2 + 3] = 0
tree[child_key1][child_branch1 + 3] = l1 + l2
ready_lineages = [(child_key1, child_branch1, ps),
(child_key2, child_branch2, ps2)]
print ready_lineages
started_admixtures = {}
while ready_lineages:
k, b, p = ready_lineages.pop()
branch_length = get_branch_length(tree, k, b)
branch = (k, b)
p = add_n(p, branch_length)
node = tree[k]
children = get_children(node)
if node_is_leaf_node(node):
res[k] = p
elif node_is_admixture(node):
mate = (k, other_branch(b))
if mate in started_admixtures:
w = get_admixture_proportion_from_key(tree, k)
p = merge_ps(w * (1 - b) + b * (1 - w), p,
started_admixtures[mate])
del started_admixtures[mate]
k_new = children[0]
b_new = mother_or_father(tree, k_new, k)
ready_lineages.append((k_new, b_new, p))
else:
started_admixtures[branch] = p
else:
for child in children:
k_new = child
b_new = mother_or_father(tree, k_new, k)
ready_lineages.append((k_new, b_new, p))
return res
