parser.add_option("-e", "--eventfile", dest="eventfile", help="events file") parser.add_option("-o", "--output", dest="output", help="output file") parser.add_option("-p", "--prune", dest="prune", type="float",default=0,help="Do not draw subs with probability less than this value") parser.add_option("-s", "--scale", dest="scale", type="int", help="8000 for little ksu, 3 for krya") parser.add_option("-c", "--circle_size", dest="size", type="float", help="20 for little ksu, ? for krya") parser.add_option("-w", "--width", dest="width", type="int", help="870 for little ksu, ? for krya") parser.add_option("-l", "--large", dest="large",action="store_true", default=False, help="do not set vertical margin between branches( branch margins used for scheme)") (options, args) = parser.parse_args() t = Tree(options.treefile, format=1) #flexible with internal node names farthest, dist = t.get_farthest_node() print ("The farthest node from root is", farthest.name, "with dist=", dist) if not options.scale: options.scale = int(1700/dist) #900/dist #root = t.get_tree_root() #farthest = root.get_farthest_leaf() #print "The height is ", t.get_distance(root.name, farthest.name) # Basic tree style ts = TreeStyle() ts.show_leaf_name = False if options.large: ts.show_scale = False #ts.optimal_scale_level = "full" ts.scale = options.scale #3 for krya # pixels per branch length unit if not options.large:
class um_tree: def __init__(self, tree, PATH): self.tree = Tree(tree, format=1) self.tree2 = open(tree) self.tree.resolve_polytomy(default_dist=0.000001, recursive=True) self.tree.dist = 0 self.tree.add_feature("age", 0) self.nodes = self.tree.get_descendants() self.PATH = PATH internal_node = [] cnt = 0 for n in self.nodes: node_age = n.get_distance(self.tree) n.add_feature("age", node_age) if not n.is_leaf(): n.add_feature("id", cnt) cnt = cnt + 1 internal_node.append(n) self.nodes = internal_node one_leaf = self.tree.get_farthest_node()[0] one_leaf.add_feature("id", cnt + 1) if one_leaf.is_leaf(): self.nodes.append(one_leaf) self.nodes.sort(key=self.__compare_node) self.species_list = [] self.coa_roots = None def __compare_node(self, node): return node.age def get_waiting_times(self, threshold_node=None, threshold_node_idx=0): wt_list = [] reach_t = False curr_age = 0.0 curr_spe = 2 curr_num_coa = 0 coa_roots = [] min_brl = 1000 num_spe = -1 if threshold_node == None: threshold_node = self.nodes[threshold_node_idx] last_coa_num = 0 tcnt = 0 for node in self.nodes: num_children = len(node.get_children()) wt = None times = node.age - curr_age if times >= 0: if times < min_brl and times > 0: min_brl = times curr_age = node.age assert curr_spe >= 0 if reach_t: if tcnt == 0: last_coa_num = 2 fnode = node.up coa_root = None idx = 0 while not fnode.is_root(): idx = 0 for coa_r in coa_roots: if coa_r.id == fnode.id: coa_root = coa_r break idx = idx + 1 if coa_root != None: break else: fnode = fnode.up wt = waiting_time(length=times, num_coas=curr_num_coa, num_lines=curr_spe) for coa_r in coa_roots: coa = coalescent(num_individual=coa_r.curr_n) wt.coas.add_coalescent(coa) wt.coas.coas_idx = last_coa_num wt.num_curr_coa = last_coa_num if (coa_root == None ): # here can be modified to use multiple T curr_spe = curr_spe - 1 curr_num_coa = curr_num_coa + 1 node.add_feature("curr_n", 2) coa_roots.append(node) last_coa_num = 2 else: curr_n = coa_root.curr_n coa_root.add_feature("curr_n", curr_n + 1) last_coa_num = curr_n + 1 tcnt = tcnt + 1 else: if node.id == threshold_node.id: reach_t = True tcnt = 0 wt = waiting_time(length=times, num_coas=0, num_lines=curr_spe) num_spe = curr_spe curr_spe = curr_spe - 1 curr_num_coa = 2 node.add_feature("curr_n", 2) coa_roots.append(node) else: wt = waiting_time(length=times, num_coas=0, num_lines=curr_spe) curr_spe = curr_spe + 1 if times > 0.00000001: wt_list.append(wt) for wt in wt_list: wt.count_num_lines() self.species_list = [] all_coa_leaves = [] self.coa_roots = coa_roots for coa_r in coa_roots: leaves = coa_r.get_leaves() all_coa_leaves.extend(leaves) self.species_list.append(leaves) all_leaves = self.tree.get_leaves() for leaf in all_leaves: if leaf not in all_coa_leaves: self.species_list.append([leaf]) return wt_list, num_spe def show(self, wt_list): cnt = 1 for wt in wt_list: print(("Waitting interval " + repr(cnt))) print(wt) cnt = cnt + 1 def get_species(self): sp_list = [] for sp in self.species_list: spe = [] for taxa in sp: spe.append(taxa.name) sp_list.append(spe) all_taxa_name = [] # self.tree.convert_to_ultrametric(tree_length = 1.0, strategy='balanced') for leaf in self.tree.get_leaves(): all_taxa_name.append(leaf.name) style0 = NodeStyle() style0["fgcolor"] = "#000000" # style2["shape"] = "circle" style0["vt_line_color"] = "#0000aa" style0["hz_line_color"] = "#0000aa" style0["vt_line_width"] = 2 style0["hz_line_width"] = 2 style0["vt_line_type"] = 0 # 0 solid, 1 dashed, 2 dotted style0["hz_line_type"] = 0 style0["size"] = 0 for node in self.tree.get_descendants(): node.set_style(style0) node.img_style["size"] = 0 self.tree.set_style(style0) self.tree.img_style["size"] = 0 style1 = NodeStyle() style1["fgcolor"] = "#000000" # style2["shape"] = "circle" style1["vt_line_color"] = "#ff0000" style1["hz_line_color"] = "#0000aa" style1["vt_line_width"] = 2 style1["hz_line_width"] = 2 style1["vt_line_type"] = 0 # 0 solid, 1 dashed, 2 dotted style1["hz_line_type"] = 0 style1["size"] = 0 style2 = NodeStyle() style2["fgcolor"] = "#0f0f0f" # style2["shape"] = "circle" style2["vt_line_color"] = "#ff0000" style2["hz_line_color"] = "#ff0000" style2["vt_line_width"] = 2 style2["hz_line_width"] = 2 style2["vt_line_type"] = 0 # 0 solid, 1 dashed, 2 dotted style2["hz_line_type"] = 0 style2["size"] = 0 for node in self.coa_roots: node.set_style(style1) node.img_style["size"] = 0 for des in node.get_descendants(): des.set_style(style2) des.img_style["size"] = 0 return [all_taxa_name], sp_list def print_species(self): # tree_path = os.path.dirname(self.tree2.name) sp_out = open(os.path.join(self.PATH, "GMYC/GMYC_MOTU.txt"), "w+") cnt = 1 for sp in self.species_list: # print("Species " + repr(cnt) + ":") sp_out.write("Species " + repr(cnt) + "\n") cnt = cnt + 1 taxas = "" for taxa in sp: taxas = taxas + taxa.name + ", " # print(" " + taxas[:-1]) sp_out.write(" " + taxas[:-1] + "\n") def output_species(self, taxa_order=[]): """taxa_order is a list of taxa names, the paritions will be output as the same order""" if len(taxa_order) == 0: taxa_order = self.tree.get_leaf_names() num_taxa = 0 for sp in self.species_list: for taxa in sp: num_taxa = num_taxa + 1 if not len(taxa_order) == num_taxa: print("error error, taxa_order != num_taxa!") return None, None else: partion = [-1] * num_taxa cnt = 1 for sp in self.species_list: for taxa in sp: idx = taxa_order.index(taxa.name) partion[idx] = cnt cnt = cnt + 1 return taxa_order, partion def num_lineages(self, wt_list): nl_list = [] times = [] last_time = 0.0 for wt in wt_list: nl_list.append(wt.get_num_branches()) times.append(last_time) last_time = wt.length + last_time plt.plot(times, nl_list) plt.ylabel("Number of lineages") plt.xlabel("Time") plt.savefig("Time_Lines") plt.show()
class um_tree: def __init__(self, tree): self.tree = Tree(tree, format=1) self.tree.resolve_polytomy(default_dist=0.000001, recursive=True) self.tree.dist = 0 self.tree.add_feature("age", 0) self.nodes = self.tree.get_descendants() internal_node = [] cnt = 0 for n in self.nodes: node_age = n.get_distance(self.tree) n.add_feature("age", node_age) if not n.is_leaf(): n.add_feature("id", cnt) cnt = cnt + 1 internal_node.append(n) self.nodes = internal_node one_leaf = self.tree.get_farthest_node()[0] one_leaf.add_feature("id", cnt + 1) if one_leaf.is_leaf(): self.nodes.append(one_leaf) self.nodes.sort(key=self.__compare_node) self.species_list = [] self.coa_roots = None def __compare_node(self, node): return node.age def get_waiting_times(self, threshold_node=None, threshold_node_idx=0): wt_list = [] reach_t = False curr_age = 0.0 curr_spe = 2 curr_num_coa = 0 coa_roots = [] min_brl = 1000 num_spe = -1 if threshold_node == None: threshold_node = self.nodes[threshold_node_idx] last_coa_num = 0 tcnt = 0 for node in self.nodes: num_children = len(node.get_children()) wt = None times = node.age - curr_age if times >= 0: if times < min_brl and times > 0: min_brl = times curr_age = node.age assert curr_spe >= 0 if reach_t: if tcnt == 0: last_coa_num = 2 fnode = node.up coa_root = None idx = 0 while not fnode.is_root(): idx = 0 for coa_r in coa_roots: if coa_r.id == fnode.id: coa_root = coa_r break idx = idx + 1 if coa_root != None: break else: fnode = fnode.up wt = waiting_time(length=times, num_coas=curr_num_coa, num_lines=curr_spe) for coa_r in coa_roots: coa = coalescent(num_individual=coa_r.curr_n) wt.coas.add_coalescent(coa) wt.coas.coas_idx = last_coa_num wt.num_curr_coa = last_coa_num if coa_root == None: #here can be modified to use multiple T curr_spe = curr_spe - 1 curr_num_coa = curr_num_coa + 1 node.add_feature("curr_n", 2) coa_roots.append(node) last_coa_num = 2 else: curr_n = coa_root.curr_n coa_root.add_feature("curr_n", curr_n + 1) last_coa_num = curr_n + 1 tcnt = tcnt + 1 else: if node.id == threshold_node.id: reach_t = True tcnt = 0 wt = waiting_time(length=times, num_coas=0, num_lines=curr_spe) num_spe = curr_spe curr_spe = curr_spe - 1 curr_num_coa = 2 node.add_feature("curr_n", 2) coa_roots.append(node) else: wt = waiting_time(length=times, num_coas=0, num_lines=curr_spe) curr_spe = curr_spe + 1 if times > 0.00000001: wt_list.append(wt) for wt in wt_list: wt.count_num_lines() self.species_list = [] all_coa_leaves = [] self.coa_roots = coa_roots for coa_r in coa_roots: leaves = coa_r.get_leaves() all_coa_leaves.extend(leaves) self.species_list.append(leaves) all_leaves = self.tree.get_leaves() for leaf in all_leaves: if leaf not in all_coa_leaves: self.species_list.append([leaf]) return wt_list, num_spe def show(self, wt_list): cnt = 1 for wt in wt_list: print("Waitting interval " + repr(cnt)) print(wt) cnt = cnt + 1 def get_species(self): sp_list = [] for sp in self.species_list: spe = [] for taxa in sp: spe.append(taxa.name) sp_list.append(spe) all_taxa_name = [] for leaf in self.tree.get_leaves(): all_taxa_name.append(leaf.name) style0 = NodeStyle() style0["fgcolor"] = "#000000" style0["vt_line_color"] = "#0000aa" style0["hz_line_color"] = "#0000aa" style0["vt_line_width"] = 2 style0["hz_line_width"] = 2 style0["vt_line_type"] = 0 # 0 solid, 1 dashed, 2 dotted style0["hz_line_type"] = 0 style0["size"] = 0 for node in self.tree.get_descendants(): node.set_style(style0) node.img_style["size"] = 0 self.tree.set_style(style0) self.tree.img_style["size"] = 0 style1 = NodeStyle() style1["fgcolor"] = "#000000" style1["vt_line_color"] = "#ff0000" style1["hz_line_color"] = "#0000aa" style1["vt_line_width"] = 2 style1["hz_line_width"] = 2 style1["vt_line_type"] = 0 # 0 solid, 1 dashed, 2 dotted style1["hz_line_type"] = 0 style1["size"] = 0 style2 = NodeStyle() style2["fgcolor"] = "#0f0f0f" style2["vt_line_color"] = "#ff0000" style2["hz_line_color"] = "#ff0000" style2["vt_line_width"] = 2 style2["hz_line_width"] = 2 style2["vt_line_type"] = 0 # 0 solid, 1 dashed, 2 dotted style2["hz_line_type"] = 0 style2["size"] = 0 for node in self.coa_roots: node.set_style(style1) node.img_style["size"] = 0 for des in node.get_descendants(): des.set_style(style2) des.img_style["size"] = 0 return [all_taxa_name], sp_list def print_species(self, save_file): cnt = 1 file3 = open(os.path.join(save_file, "partition.txt"), "w+") for sp in self.species_list: print("Species " + repr(cnt) + ":", file=file3) cnt = cnt + 1 taxas = "" for taxa in sp: taxas = taxas + taxa.name + ", " print("" + taxas[:-1], file=file3) def print_species_spart(self, save_file): cnt = 1 file3 = open(os.path.join(save_file, "partition.spart"), "w+") # for sp in self.species_list: # print("Species " + repr(cnt) + ":", file= file3) # cnt = cnt + 1 # taxas = "" # for taxa in sp: # print(taxa) # taxas = taxas + taxa.name + ", " # print("" + taxas[:-1], file= file3) file3.write("Filename=GMYC delimitation\n") file3.write(f'{datetime.datetime.now().astimezone().isoformat()}\n\n') file3.write(f"Npartition={1};GMYC\n") file3.write(f'Nsamples={sum(len(sp) for sp in self.species_list)}\n') file3.write( f'Nsubsets={len(self.species_list)};{",".join(["?" for i in range(len(self.species_list))])}\n\n' ) file3.write("#this is my first comment\n") file3.write("#this is my second comment\n\n") file3.write("Assignment\n") cnt = 1 for sp in self.species_list: print(repr(sp)) for taxa in sp: print(repr(taxa)) xx = taxa.name + "\t" + repr(cnt) + ";" + "?" file3.write(f"{xx}\n") cnt += 1 file3.write("\nPartition_score=\n") file3.close() def output_species(self, taxa_order=[]): if len(taxa_order) == 0: taxa_order = self.tree.get_leaf_names() num_taxa = 0 for sp in self.species_list: for taxa in sp: num_taxa = num_taxa + 1 if not len(taxa_order) == num_taxa: print("error error, taxa_order != num_taxa!") return None, None else: partion = [-1] * num_taxa cnt = 1 for sp in self.species_list: for taxa in sp: idx = taxa_order.index(taxa.name) partion[idx] = cnt cnt = cnt + 1 return taxa_order, partion def num_lineages(self, wt_list, save_file): nl_list = [] times = [] last_time = 0.0 for wt in wt_list: nl_list.append(wt.get_num_branches()) times.append(last_time) last_time = wt.length + last_time plt.plot(times, nl_list) plt.ylabel('Number of lineages') plt.xlabel('Time') plt.savefig(os.path.join(save_file, "Time_Lines.png"))
else: sys.exit("Invalid parent-daughter: %s, %s" % (parentloc, thisloc)) return 5 * base + add def FindLocationOfParent(node): parent = node.up try: return infodic[parent.name]["location"] except AttributeError: sys.exit("Error finding parent of %s" % node.name) tree = Tree(sys.argv[1], format=1) rootheight = tree.get_farthest_node()[1] print("Tree height: %s" % rootheight) infodic = {} with open(sys.argv[2], "rU") as infofile: csvreader = csv.reader(infofile, delimiter=",") header = csvreader.next() for row in csvreader: infodic[row[0]] = { "height": row[1], "length": row[2], "location": row[3] } calendar = [[0 for year in xrange(int(math.ceil(rootheight)) + 1)]