def WriteResults(species_tree_fn_or_text, roots, S, clades, clusters_counter, output_dir):
# for c in clusters_counter:
# print((clusters_counter[c], c))
print("\nResults written to:\n" + os.path.realpath(output_dir))
# Label species tree nodes
species_tree = tree.Tree(species_tree_fn_or_text)
thisRoot = roots[0]
species_tree = RootAtClade(species_tree, thisRoot)
iNode = 0
for n in species_tree.traverse():
if not n.is_leaf():
n.name = "N%d" % iNode
iNode+=1
species_tree.write(outfile=output_dir + "Species_tree_labelled.tre", format=1)
# print(species_tree)
# species_tree = tree.Tree(output_dir + "Species_tree_labelled.tre", format=1)
# Calculate probabilities
qBinary = True
for n in species_tree.traverse():
if len(n.get_children()) > 2:
qBinary = False
if qBinary:
p_final = probroot.GetProbabilities(species_tree, S, clades, clusters_counter)
else:
print("Probability distribution for root location is not supported for non-binary trees")
print("To get a probability distribution for the root, please supply a fully resolved input species tree")
# Write numbers of duplications
table = dict()
new_tree = tree.Tree(output_dir + "Species_tree_labelled.tre", format=1)
for clade in clades + [frozenset([s]) for s in S]:
qAnti = False
anticlade = S.difference(clade)
if len(clade) == 1:
node = new_tree & list(clade)[0]
else:
node = new_tree.get_common_ancestor(clade)
if node == new_tree:
node = new_tree.get_common_ancestor(anticlade)
qAnti = True
x = anticlade if qAnti else clade
y = clade if qAnti else anticlade
X = ("(%d)" % clusters_counter[x]) if len(clade) == 1 else clusters_counter[x]
if qBinary:
p = p_final[clade] if clade in p_final else p_final[anticlade]
else:
p = 0.
table[node.name] = [node.name, "X" if (clade in roots or anticlade in roots) else "", "%0.1f%%" % (100.*p) , X, clusters_counter[y]]
with open(output_dir + "Duplication_counts.csv", 'wb') as outfile:
writer = csv.writer(outfile, delimiter="\t")
writer.writerow(["Branch", "MP Root", "Probability", "Duplications supporting clade", "Duplications supporting opposite clade"])
qSingle = len(thisRoot) == 1
root_branches = [n.name for n in new_tree.get_children()]
writer.writerow([root_branches[0] + " (& " + root_branches[1] + ")"] + table[root_branches[0]][1:])
for i in range(2 if qSingle else 3, iNode):
name = "N%d" % i
if name in table:
writer.writerow(table[name])
else:
print("Skipping %s" % name)
for sp in S:
if sp in table:
if qSingle and sp in thisRoot: continue
writer.writerow(table[sp])
def WriteResults(species_tree_fn_or_text, roots, S, clades, clusters_counter, output_dir):
# for c in clusters_counter:
# print((clusters_counter[c], c))
print("\nResults written to:\n" + os.path.realpath(output_dir))
# Label species tree nodes
species_tree = tree.Tree(species_tree_fn_or_text)
thisRoot = roots[0]
species_tree = RootAtClade(species_tree, thisRoot)
iNode = 0
for n in species_tree.traverse():
if not n.is_leaf():
n.name = "N%d" % iNode
iNode+=1
species_tree.write(outfile=output_dir + "Species_tree_labelled.tre", format=1)
# print(species_tree)
# species_tree = tree.Tree(output_dir + "Species_tree_labelled.tre", format=1)
# Calculate probabilities
p_final = probroot.GetProbabilities(species_tree, S, clades, clusters_counter)
# Write numbers of duplications
table = dict()
new_tree = tree.Tree(output_dir + "Species_tree_labelled.tre", format=1)
for clade in clades + [frozenset([s]) for s in S]:
# print(clade)
qAnti = False
anticlade = S.difference(clade)
if len(clade) == 1:
node = new_tree & list(clade)[0]
# if len(anticlade) == 1: qAnti = True:
# elif len(anticlade) == 1:
else:
node = new_tree.get_common_ancestor(clade)
if node == new_tree:
node = new_tree.get_common_ancestor(anticlade)
qAnti = True
# print("")
# print(clade)
# print(anticlade)
# print(node.name)
# print(node)
p = p_final[clade] if clade in p_final else p_final[anticlade]
# qThisRooting = (clade == thisRoot) or (anticlade == thisRoot)
x = anticlade if qAnti else clade
y = clade if qAnti else anticlade
X = ("(%d)" % clusters_counter[x]) if len(clade) == 1 else clusters_counter[x]
table[node.name] = [node.name, "X" if (clade in roots or anticlade in roots) else "", "%0.1f%%" % (100.*p) , X, clusters_counter[y]]
# print(len(table))
# for t in table.items():
# print(t)
with open(output_dir + "Duplication_counts.csv", 'wb') as outfile:
writer = csv.writer(outfile)
writer.writerow(["Branch", "MP Root", "Probability", "Duplications supporting clade", "Duplications supporting opposite clade"])
qSingle = len(thisRoot) == 1
root_branches = [n.name for n in new_tree.get_children()]
writer.writerow([root_branches[0] + " (& " + root_branches[1] + ")"] + table[root_branches[0]][1:])
# if qSingle:
# else:
# root_data = table['N1']
# writer.writerow(['N1 (& N2)'] + root_data[1:])
for i in range(2 if qSingle else 3, iNode):
name = "N%d" % i
if name in table:
writer.writerow(table[name])
else:
print("Skipping %s" % name)
# writer.writerow([name, "X" if (clade in roots or anticlade in roots) else "", "%0.1f%%" % (100.*p) , clusters_counter[clade], clusters_counter[anticlade]])
for sp in S:
if sp in table:
if qSingle and sp in thisRoot: continue
writer.writerow(table[sp])
