def generation_counts(tree):
'''
returns a list of tuples of the form (no_admixtures, no_waiting_coalescences, no_sudden_coalescences, no_awaited_coalescences).
'''
leaves, coalescences_nodes, admixture_nodes=get_categories(tree)
ready_lineages=[(key,0) for key in leaves]
coalescences_on_hold=[]
res=[]
while True:
sames, single_coalescences, admixtures=get_destination_of_lineages(tree, ready_lineages)
waiting_coalescences, awaited_coalescences, still_on_hold = matchmake(single_coalescences, coalescences_on_hold)
#print 'sames', sames
#print 'single_coalescences', single_coalescences
#print 'admixtures', admixtures
#print 'waiting_coalescences', waiting_coalescences
#print 'awaited_coalescences', awaited_coalescences
#print 'still_on_hold', still_on_hold
res.append((len(sames), len(waiting_coalescences), len(awaited_coalescences), len(admixtures)))
#updating lineages
coalescences_on_hold=still_on_hold+waiting_coalescences.values()
ready_lineages=propagate_married(tree, awaited_coalescences)
ready_lineages.extend(propagate_married(tree, sames))
ready_lineages.extend(propagate_admixtures(tree, admixtures))
#stop criteria
if len(ready_lineages)==1 and ready_lineages[0][0]=='r':
break
return res
def node_count(tree):
'''
Returns number of leaves, coalescence nodes, and admixture nodes in a tuple of size 3.
'''
leaves, coalescence_nodes, admixture_nodes=get_categories(tree)
return len(leaves), len(coalescence_nodes), len(admixture_nodes)
def tree_to_random_ntree(tree):
leaves,_,admixture_keys=get_categories(tree)
pruned_tree = deepcopy(tree)
for adm_key in admixture_keys:
if adm_key in pruned_tree:
#print '------------------------------------------'
#print 'removing', (adm_key, int_bin) , 'from tree:'
#pretty_print(pruned_tree)
pruned_tree=remove_admixture(pruned_tree, adm_key, int(random()<0.5))
return non_admixture_to_newick(pruned_tree)
def get_possible_permutation_strees(tree):
leaves, _, admixture_keys = get_categories(tree)
k = len(admixture_keys)
format_code = '{0:0' + str(k) + 'b}'
n_trees = []
for i in range(2**k):
pruned_tree = deepcopy(tree)
bina = format_code.format(i)
prop = 1.0
for adm_key, str_bin in zip(admixture_keys, list(bina)):
int_bin = int(str_bin)
if int_bin == 1:
pruned_tree[adm_key] = change_admixture(pruned_tree[adm_key])
n_tree = unique_identifier_and_branch_lengths(pruned_tree)
n_trees.append(n_tree)
return n_trees
def tree_to_mode_ntree(tree):
'''
Every admixture node is collapsed such that the branch with smaller weight is removed. This is not the same majority tree because
'''
leaves,_,admixture_keys=get_categories(tree)
pruned_tree = deepcopy(tree)
for adm_key in admixture_keys:
if adm_key in pruned_tree:
if get_admixture_proportion_from_key(tree, adm_key)>0.5:
remove=1
else:
remove=0
#print '------------------------------------------'
#print 'removing', (adm_key, int_bin) , 'from tree:'
#pretty_print(pruned_tree)
pruned_tree=remove_admixture(pruned_tree, adm_key, remove)
return non_admixture_to_newick(pruned_tree)
def topological_prior(tree):
total_prob = 0
no_admix = get_number_of_admixes(tree)
leaves, coalescences_nodes, admixture_nodes = get_categories(tree)
ready_lineages = [(key, 0) for key in leaves]
totally_free_coalescences = coalescences_nodes + ['r']
free_admixtures = admixture_nodes
coalescences_on_hold = []
res = 0
while True:
sames, single_coalescences, admixtures = get_destination_of_lineages(
tree, ready_lineages)
waiting_coalescences, awaited_coalescences, still_on_hold = matchmake(
single_coalescences, coalescences_on_hold)
#print sames, waiting_coalescences, admixtures
res += _get_selection_probabilities(
no_sames=len(sames),
no_waiting_coalescences=len(waiting_coalescences),
no_awaited_coalescences=len(awaited_coalescences),
no_admixtures=len(admixtures),
no_totally_free_coalescences=len(totally_free_coalescences),
no_free_admixtures=len(free_admixtures),
no_ready_lineages=len(ready_lineages),
no_coalescences_on_hold=len(coalescences_on_hold),
no_admixture_pairs=_get_number_of_admixture_branches(
ready_lineages))
#updating lineages
coalescences_on_hold = still_on_hold + waiting_coalescences.values()
totally_free_coalescences, free_admixtures = _thin_out_frees(
tree, totally_free_coalescences, free_admixtures, ready_lineages)
#print 'free_admixture', free_admixtures
#print 'totally_free_coalescences', totally_free_coalescences
ready_lineages = propagate_married(tree, awaited_coalescences)
ready_lineages.extend(propagate_married(tree, sames))
ready_lineages.extend(propagate_admixtures(tree, admixtures))
#print 'ready_lineages', ready_lineages
#stop criteria
if len(ready_lineages) == 1 and ready_lineages[0][0] == 'r':
break
return res
def non_admixture_to_newick(tree):
leaves,_,_=get_categories(tree)
keys_to_pops={l:l for l in leaves}
while len(keys_to_pops)>1:
next_gen={} #dictionary of mapping from parent to a list of children
dup_children=[]
dup_parents=[]
for key in keys_to_pops:
parent_key=tree[key][0]
if parent_key in next_gen:
next_gen[parent_key]=[next_gen[parent_key][0], key]
dup_parents.append(parent_key)
dup_children.append(next_gen[parent_key])
else:
next_gen[parent_key]=[key]
# print next_gen
# print dup_parents
# print dup_children
for (c1,c2),p in zip(dup_children, dup_parents):
keys_to_pops[p]='('+','.join(sorted([keys_to_pops[c1],keys_to_pops[c2]]))+')'
del keys_to_pops[c1]
del keys_to_pops[c2]
return keys_to_pops.values()[0]
def get_timing(tree):
'''
returns a list of tuples of the form (no_admixtures, no_waiting_coalescences, no_sudden_coalescences, no_awaited_coalescences).
'''
leaves, coalescences_nodes, admixture_nodes=get_categories(tree)
ready_lineages=[(key,0) for key in leaves]
coalescences_on_hold=[]
res={key:0.0 for key in leaves}
count=1
while True:
sames, single_coalescences, admixtures=get_destination_of_lineages(tree, ready_lineages)
waiting_coalescences, awaited_coalescences, still_on_hold = matchmake(single_coalescences, coalescences_on_hold)
#print 'sames', sames
#print 'single_coalescences', single_coalescences
#print 'admixtures', admixtures
#print 'waiting_coalescences', waiting_coalescences
#print 'awaited_coalescences', awaited_coalescences
#print 'still_on_hold', still_on_hold
#updating lineages
coalescences_on_hold=still_on_hold+waiting_coalescences.values()
ready_lineages=propagate_married(tree, awaited_coalescences)
ready_lineages.extend(propagate_married(tree, sames))
ready_lineages.extend(propagate_admixtures(tree, admixtures))
res.update({key:count for key,_ in ready_lineages})
#stop criteria
if len(ready_lineages)==1 and ready_lineages[0][0]=='r':
res['r']=count
break
count+=1
return res
def tree_to_prop_newicks(tree, leaves=None):
'''
Transforms a tree in the admixture identifier format into a dictionary of newick trees where the keys are the trees and the values are the proportion of trees of that form.
'''
leaves,_,admixture_keys=get_categories(tree)
k=len(admixture_keys)
format_code='{0:0'+str(k)+'b}'
n_trees={}
max_count=65
for i in range(2**k):
pruned_tree = deepcopy(tree)
bina= format_code.format(i)
prop=1.0
for adm_key,str_bin in zip(admixture_keys, list(bina)):
int_bin=int(str_bin)
if int_bin==0:
prop*=1.0-get_admixture_proportion_from_key(tree, adm_key)
else:
prop*=get_admixture_proportion_from_key(tree, adm_key)
if adm_key in pruned_tree:
#print '------------------------------------------'
#print 'removing', (adm_key, int_bin) , 'from tree:'
#pretty_print(pruned_tree)
pruned_tree=remove_admixture(pruned_tree, adm_key, int_bin)
n_tree= non_admixture_to_newick(pruned_tree)
if n_tree in n_trees:
n_trees[n_tree]+=prop
else:
n_trees[n_tree]=prop
if i>max_count:
break
return n_trees
def unique_identifier_and_branch_lengths(tree, leaf_order=None):
leaves, coalescences_nodes, admixture_nodes=get_categories(tree)
if leaf_order is not None:
assert set(leaves)==set(leaf_order), 'the specified leaf order did not match the leaves of the tree.'
leaves_ordered=leaf_order
else:
leaves_ordered=sorted(leaves)
ready_lineages=[(key,0) for key in leaves_ordered]
lineages=deepcopy(ready_lineages)
gen_to_column=range(len(ready_lineages))
list_of_gens=[]
coalescences_on_hold=[]
gone=[]
res=[]
branch_lengths=[]
admixture_proportions=[]
while True:
sames, single_coalescences, admixtures=get_destination_of_lineages(tree, ready_lineages)
waiting_coalescences, awaited_coalescences, still_on_hold = matchmake(single_coalescences, coalescences_on_hold)
res.append((len(sames), len(waiting_coalescences), len(awaited_coalescences), len(admixtures)))
sames_dic, first_sames, second_sames = make_dics_first_and_second(sames)
awaited_dic, first_awaited, second_awaited = make_dics_first_and_second(awaited_coalescences)
waiting=waiting_coalescences.keys()
gen=[]
#print 'lineages',lineages
#print 'sames', sames, sames_dic, first_sames, second_sames
#print 'awaited', awaited_coalescences, awaited_dic, first_awaited, second_awaited
for n,element in enumerate(lineages):
if element in gone:
print 'entered gone!'
gen.append('_')
elif element in sames_dic:
partner_index=lineages.index(sames_dic[element])
if n<partner_index:
gen.append('c')
branch_lengths.append(get_branch_length(tree, element[0],element[1]))
else:
gen.append(partner_index)
branch_lengths.append(get_branch_length(tree, element[0],element[1]))
elif element in awaited_dic:
partner_index=lineages.index(awaited_dic[element])
if n<partner_index:
gen.append('c')
branch_lengths.append(get_branch_length(tree, element[0],element[1]))
else:
gen.append(partner_index)
branch_lengths.append(get_branch_length(tree, element[0],element[1]))
elif element in admixtures:
gen.append('a')
branch_lengths.append(get_branch_length(tree, element[0],element[1]))
admixture_proportions.append(get_admixture_proportion(tree, child_key=element[0],child_branch=element[1]))
else:
gen.append('w')
#print 'gen',gen
#print 'gone', gone
list_of_gens,gone, lineages =update_lineages(list_of_gens,gen,gone, lineages, tree)
for gon in gone:
lineages.remove(gon)
gone=[]
#updating lineages
coalescences_on_hold=still_on_hold+waiting_coalescences.values()
ready_lineages=propagate_married(tree, awaited_coalescences)
ready_lineages.extend(propagate_married(tree, sames))
ready_lineages.extend(propagate_admixtures(tree, admixtures))
#stop criteria
if len(ready_lineages)==1 and ready_lineages[0][0]=='r':
break
return ';'.join([_list_identifier_to_string(list_of_gens),
_list_double_to_string(branch_lengths, 9),
_list_double_to_string(admixture_proportions, 3)])