def resimulate_moved_branch(tree, key, branch, delta_L, alpha):
old_length= get_branch_length(tree, key, branch)
new_length= gamma_restricted.rvs(old_length, -delta_L, alpha)
logpdf=gamma_restricted.logpdf(new_length, old_length, -delta_L, alpha)
logpdf_back=gamma_restricted.logpdf(old_length, new_length, delta_L, alpha)
backward=exp(-logpdf+logpdf_back) #for stability, both forward and backward are put in backward.
ratio= exp(logpdf-logpdf_back)
update_branch_length(tree, key, branch, new_length)
#print 'resimulated', old_length, 'to', new_length, 'on branch', str((key,branch)), 'and change in lattitude at', delta_L, 'backward', backward
return tree, 1.0, backward
def follow(self, tree, visited_keys=[]):
node = tree[self.key]
new_lineages = []
pieces = []
new_keys_visited = []
for key in get_real_children(tree[self.key]):
if key not in visited_keys:
branch = mother_or_father(tree,
child_key=key,
parent_key=self.key)
l = get_branch_length(tree, key, branch)
pieces.append(
piece(self.lattitude, self.lattitude - l, self.distance,
self.distance + l, key, branch, self.key))
new_lineages.append(
lineage(key,
self.distance + l,
self.lattitude - l,
topological_distance=self.topological_distance))
for key in get_real_parents(tree[self.key]):
if key not in visited_keys:
branch = mother_or_father(tree,
child_key=self.key,
parent_key=key)
l = get_branch_length(tree, self.key, branch)
pieces.append(
piece(self.lattitude, self.lattitude + l, self.distance,
self.distance + l, self.key, branch, key))
new_lineages.append(
lineage(key,
self.distance + l,
self.lattitude + l,
topological_distance=self.topological_distance))
if self.key == 'r': #add the very long piece
pieces.append(
piece(self.lattitude, None, self.distance, None, 'r', 0, None))
return new_lineages, pieces
def regraft(tree,
remove_key,
remove_branch,
add_to_branch,
new_node=None,
which_branch=0):
if add_to_branch == 'r':
insertion_spot, q = simulate_and_forward_density('r')
else:
branch_length = get_branch_length(tree, add_to_branch, which_branch)
#print branch_length
insertion_spot, q = simulate_and_forward_density('u', branch_length)
if new_node is None:
new_node = str(getrandbits(68)).strip()
tree = graft(tree,
remove_key,
add_to_branch,
insertion_spot,
new_node,
which_branch,
remove_branch=remove_branch)
return tree, q
def insert_admix(tree,
source_key,
source_branch,
sink_key,
sink_branch,
source_name=None,
sink_name=None,
pks={},
new_branch_length=None,
new_to_root_length=None,
preserve_root_distance=False):
#print "new branch", new_branch_length
if source_key == 'r':
u1, q1 = get_root_branch_length()
if new_to_root_length is not None:
u1, q1 = new_to_root_length, get_root_branch_length(
new_to_root_length)
else:
u1, q1 = get_insertion_spot(
length=get_branch_length(tree, source_key, source_branch))
t1 = get_branch_length(tree, sink_key, sink_branch)
u2, q2 = get_insertion_spot(
length=get_branch_length(tree, sink_key, sink_branch))
if new_branch_length is not None:
t4, q4 = new_branch_length, get_admixture_branch_length(
new_branch_length)
else:
t4, q4 = get_admixture_branch_length()
u3, q3 = get_admixture_proportion()
if sink_name is None:
sink_name = str(getrandbits(128)).strip()
if source_name is None:
source_name = str(getrandbits(128)).strip()
tree = insert_admixture_node_halfly(tree,
sink_key,
sink_branch,
u2,
admix_b_length=t4,
new_node_name=sink_name,
admixture_proportion=u3)
#print 'tree after inserting admixture', tree
tree = graft(tree,
sink_name,
source_key,
u1,
source_name,
source_branch,
remove_branch=1)
#print 'old_t1', tree[sink_name][3]
if preserve_root_distance:
tree[sink_name], multip = readjust_length(tree[sink_name])
else:
multip = 1.0
#print 'new_t1', tree[sink_name][3]
new_branch = 1
if random() < 0.5:
new_branch = 0
tree[sink_name] = change_admixture(tree[sink_name])
pks['t5'] = t4
pks['t1'] = u1
pks['sink_new_name'] = sink_name
pks['sink_new_branch'] = new_branch
#print 'tree after grafting', tree
return tree, q1 * q2 * q3 * q4, 1, multip
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)])
def insert_admix(tree,
source_key,
source_branch,
sink_key,
sink_branch,
source_name=None,
sink_name=None,
pks={},
new_branch_length=None,
new_to_root_length=None,
preserve_root_distance=False):
#print "new branch", new_branch_length
branches = get_all_branches(tree)
new_branches = deepcopy(branches)
branch_indices = {branch: n for n, branch in enumerate(branches)}
no_branches = len(branches)
U_matrix = identity(no_branches)
branch_lengths = get_specific_branch_lengths(tree, branches)
total_length = sum(branch_lengths)
if sink_name is None:
sink_name = str(getrandbits(128)).strip()
if source_name is None:
source_name = str(getrandbits(128)).strip()
if source_key == 'r':
u1, q1 = get_root_branch_length()
if new_to_root_length is not None:
u1, q1 = new_to_root_length, get_root_branch_length(
new_to_root_length)
U_matrix = insert(U_matrix,
U_matrix.shape[0],
[u1 / total_length] * U_matrix.shape[1],
axis=0)
else:
u1, q1 = get_insertion_spot(
length=get_branch_length(tree, source_key, source_branch))
index = branch_indices[(source_key, source_branch)]
insertion = [0] * U_matrix.shape[1]
insertion[index] = 1.0 - u1
U_matrix[index, index] = u1
U_matrix = insert(U_matrix, U_matrix.shape[0], insertion, axis=0)
branches.append((source_name, 0))
u2, q2 = get_insertion_spot(
length=get_branch_length(tree, sink_key, sink_branch))
index = branch_indices[(sink_key, sink_branch)]
insertion = [0] * U_matrix.shape[1]
insertion[index] = 1.0 - u2
U_matrix[index, index] = u2
#print U_matrix
U_matrix = insert(U_matrix, U_matrix.shape[0], insertion, axis=0)
branches.append((sink_name, 0))
if new_branch_length is not None:
t4, q4 = new_branch_length, get_admixture_branch_length(
new_branch_length)
else:
t4, q4 = get_admixture_branch_length()
u3, q3 = get_admixture_proportion()
U_matrix = insert(U_matrix,
U_matrix.shape[0],
[t4 / total_length] * U_matrix.shape[1],
axis=0)
branches.append((sink_name, 1))
tree = insert_admixture_node_halfly(tree,
sink_key,
sink_branch,
u2,
admix_b_length=t4,
new_node_name=sink_name,
admixture_proportion=u3)
#print 'tree after inserting admixture', tree
tree = graft(tree,
sink_name,
source_key,
u1,
source_name,
source_branch,
remove_branch=1)
#print 'old_t1', tree[sink_name][3]
if preserve_root_distance:
tree[sink_name], multip = readjust_length(tree[sink_name])
else:
multip = 1.0
#print 'new_t1', tree[sink_name][3]
new_branch = 1
if random() < 0.5:
new_branch = 0
tree[sink_name] = change_admixture(tree[sink_name])
branches[-2:] = list(reversed(branches[-2:]))
pks['t5'] = t4
pks['t1'] = u1
pks['sink_new_name'] = sink_name
pks['sink_new_branch'] = new_branch
#print 'tree after grafting', tree
return tree, q1 * q2 * q3 * q4, 1, multip, U_matrix, branches
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