# Draws a genealogical tree (generates a SVG file) based on parent-child relationship information. # Supports files generated by Framsticks experiments. import json import random import math import argparse import time as ttime TIME = "" # BIRTHS / GENERATIONAL / REAL BALANCE = "" # MIN / DENSITY DOT_STYLE = "" # NONE / NORMAL / CLEAR JITTER = "" # # ------SVG--------- svg_file = 0 svg_line_style = 'stroke="rgb(90%,10%,16%)" stroke-width="1" stroke-opacity="0.7"' svg_mutation_line_style = 'stroke-width="1"' svg_crossover_line_style = 'stroke-width="1"' svg_spine_line_style = 'stroke="rgb(0%,90%,40%)" stroke-width="2" stroke-opacity="1"' svg_scale_line_style = 'stroke="black" stroke-width="0.5" stroke-opacity="1" stroke-dasharray="5, 5"' svg_dot_style = 'r="2" stroke="black" stroke-width="0.2" fill="red"' svg_clear_dot_style = 'r="2" stroke="black" stroke-width="0.4" fill="none"' svg_spine_dot_style = 'r="1" stroke="black" stroke-width="0.2" fill="rgb(50%,50%,100%)"' svg_scale_text_style = 'style="font-family: Arial; font-size: 12; fill: #000000;"' def hex_to_style(hex): default_style = ' stroke="black" stroke-opacity="0.5" ' if hex[0] == "#": hex = hex[1:] if len(hex) == 6 or len(hex) == 8: try: int(hex, 16) except: print("Invalid characters in the color's hex #" + hex + "! Assuming black.") return default_style red = 100*int(hex[0:2], 16)/255 green = 100*int(hex[2:4], 16)/255 blue = 100*int(hex[4:6], 16)/255 opacity = 0.5 if len(hex) == 8: opacity = int(hex[6:8], 16)/255 return ' stroke="rgb(' +str(red)+ '%,' +str(green)+ '%,' +str(blue)+ '%)" stroke-opacity="' +str(opacity)+ '" ' else: print("Invalid number of digits in the color's hex #" + hex + "! Assuming black.") return default_style def svg_add_line(from_pos, to_pos, style=svg_line_style): svg_file.write('') def svg_add_text(text, pos, anchor, style=svg_scale_text_style): svg_file.write('' + text + '') def svg_add_dot(pos, style=svg_dot_style): svg_file.write('') def svg_generate_line_style(percent): # hotdog from_col = [100, 70, 0] to_col = [60, 0, 0] # lava # from_col = [100, 80, 0] # to_col = [100, 0, 0] # neon # from_col = [30, 200, 255] # to_col = [240, 0, 220] from_opa = 0.2 to_opa = 1.0 from_stroke = 1 to_stroke = 3 opa = from_opa*(1-percent) + to_opa*percent stroke = from_stroke*(1-percent) + to_stroke*percent percent = 1 - ((1-percent)**20) return 'stroke="rgb(' + str(from_col[0]*(1-percent) + to_col[0]*percent) + '%,' \ + str(from_col[1]*(1-percent) + to_col[1]*percent) + '%,' \ + str(from_col[2]*(1-percent) + to_col[2]*percent) + '%)" stroke-width="' + str(stroke) + '" stroke-opacity="' + str(opa) + '"' def svg_generate_dot_style(kind): kinds = ["red", "lawngreen", "royalblue", "magenta", "yellow", "cyan", "white", "black"] r = min(2500/len(nodes), 10) return 'fill="' + kinds[kind] + '" r="' + str(r) + '" stroke="black" stroke-width="' + str(r/10) + '" fill-opacity="1.0" ' \ 'stroke-opacity="1.0"' # ------------------- def load_data(dir): global firstnode, nodes, inv_nodes, time f = open(dir) loaded = 0 for line in f: sline = line.split(' ', 1) if len(sline) == 2: if sline[0] == "[OFFSPRING]": creature = json.loads(sline[1]) #print("B" +str(creature)) if "FromIDs" in creature: if not creature["ID"] in nodes: nodes[creature["ID"]] = {} # we assign to each parent its contribution to the genotype of the child for i in range(0, len(creature["FromIDs"])): inherited = 1 #(creature["Inherited"][i] if 'Inherited' in creature else 1) #ONLY FOR NOW nodes[creature["ID"]][creature["FromIDs"][i]] = inherited else: print("Duplicated entry for " + creature["ID"]) quit() if not creature["FromIDs"][0] in nodes and firstnode == None: firstnode = creature["FromIDs"][0] if "Time" in creature: time[creature["ID"]] = creature["Time"] if "Kind" in creature: kind[creature["ID"]] = creature["Kind"] loaded += 1 if loaded == max_nodes and max_nodes != 0: break for k, v in sorted(nodes.items()): for val in sorted(v): inv_nodes[val] = inv_nodes.get(val, []) inv_nodes[val].append(k) print(len(nodes)) def load_simple_data(dir): global firstnode, nodes, inv_nodes f = open(dir) loaded = 0 for line in f: sline = line.split() if len(sline) > 1: #if int(sline[0]) > 15000: # break if sline[0] == firstnode: continue nodes[sline[0]] = str(max(int(sline[1]), int(firstnode))) else: firstnode = sline[0] loaded += 1 if loaded == max_nodes and max_nodes != 0: break for k, v in sorted(nodes.items()): inv_nodes[v] = inv_nodes.get(v, []) inv_nodes[v].append(k) #print(str(inv_nodes)) #quit() def compute_depth(node): my_depth = 0 if node in inv_nodes: for c in inv_nodes[node]: my_depth = max(my_depth, compute_depth(c)+1) depth[node] = my_depth return my_depth # ------------------------------------ def xmin_crowd(x1, x2, y): if BALANCE == "RANDOM": return (x1 if random.randrange(2) == 0 else x2) elif BALANCE == "MIN": x1_closest = 999999 x2_closest = 999999 for pos in positions: pos = positions[pos] if pos[1] == y: x1_closest = min(x1_closest, abs(x1-pos[0])) x2_closest = min(x2_closest, abs(x2-pos[0])) return (x1 if x1_closest > x2_closest else x2) elif BALANCE == "DENSITY": x1_dist = 0 x2_dist = 0 ymin = y-10 ymax = y+10 for pos in positions: pos = positions[pos] if pos[1] > ymin or pos[1] < ymax: dysq = (pos[1]-y)**2 dx1 = pos[0]-x1 dx2 = pos[0]-x2 x1_dist += math.sqrt(dysq + dx1**2) x2_dist += math.sqrt(dysq + dx2**2) return (x1 if x1_dist > x2_dist else x2) # ------------------------------------ def prepos_children(): global max_height, max_width, min_width, visited, TIME print("firstnode " + firstnode) if not bool(time): print("REAL time requested, but no real time data provided. Assuming BIRTHS time instead.") TIME = "BIRTHS" positions[firstnode] = [0, 0] #visited = {} #visited[firstnode] = True nodes_to_visit = [firstnode] ccc = 0 timet = ttime.time() while True: ccc += 1 if ccc%1000 == 0 : print(str(ccc) + " " + str(ttime.time()-timet)) timet = ttime.time() current_node = nodes_to_visit[0] if current_node in inv_nodes: for c in inv_nodes[current_node]: # we want to visit the node just once, after all of its parents if c not in nodes_to_visit: nodes_to_visit.append(c) cy = 0 if TIME == "BIRTHS": if c[0] == "c": cy = int(c[1:]) else: cy = int(c) elif TIME == "GENERATIONAL": cy = positions[current_node][1]+1 elif TIME == "REAL": cy = time[c] if len(nodes[c]) == 1: dissimilarity = 0 if JITTER == True: dissimilarity = random.gauss(0,1) else: dissimilarity = 1 positions[c] = [xmin_crowd(positions[current_node][0]-dissimilarity, positions[current_node][0]+dissimilarity, cy), cy] else: vsum = sum([v for k, v in nodes[c].items()]) cx = sum([positions[k][0]*v/vsum for k, v in nodes[c].items()]) if JITTER == True: positions[c] = [cx + random.gauss(0, 0.1), cy] else: positions[c] = [cx, cy] #if c in inv_nodes: # prepos_children_reccurent(c) nodes_to_visit = nodes_to_visit[1:] # if none left, we can stop if len(nodes_to_visit) == 0: break # prepos_children_reccurent(firstnode) for pos in positions: max_height = max(max_height, positions[pos][1]) max_width = max(max_width, positions[pos][0]) min_width = min(min_width, positions[pos][0]) # ------------------------------------ def all_parents_visited(node): apv = True for k, v in sorted(nodes[node].items()): if not k in visited: apv = False break return apv # ------------------------------------ def draw_children(): max_depth = 0 for k, v in depth.items(): max_depth = max(max_depth, v) nodes_to_visit = [firstnode] while True: current_node = nodes_to_visit[0] if current_node in inv_nodes: for c in inv_nodes[current_node]: # inv_node => p->c if not c in nodes_to_visit: nodes_to_visit.append(c) line_style = "" if COLORING == "NONE": line_style = svg_line_style elif COLORING == "TYPE": line_style = (svg_mutation_line_style if len(nodes[c]) == 1 else svg_crossover_line_style) else: # IMPORTANCE, default line_style = svg_generate_line_style(depth[c]/max_depth) svg_add_line( (w_margin+w_no_margs*(positions[current_node][0]-min_width)/(max_width-min_width), h_margin+h_no_margs*positions[current_node][1]/max_height), (w_margin+w_no_margs*(positions[c][0]-min_width)/(max_width-min_width), h_margin+h_no_margs*positions[c][1]/max_height), line_style) # we want to draw the node just once if DOT_STYLE == "NONE": continue elif DOT_STYLE == "TYPE": dot_style = svg_generate_dot_style(kind[current_node] if current_node in kind else 0) #type else: # NORMAL, default dot_style = svg_clear_dot_style #svg_generate_dot_style(depth[c]/max_depth) svg_add_dot( (w_margin+w_no_margs*(positions[current_node][0]-min_width)/(max_width-min_width), h_margin+h_no_margs*positions[current_node][1]/max_height), dot_style) #svg_add_text( str(depth[current_node]), (w_margin+w_no_margs*(positions[current_node][0]-min_width)/(max_width-min_width), # h_margin+h_no_margs*positions[current_node][1]/max_height), "end") # we remove the current node from the list nodes_to_visit = nodes_to_visit[1:] # if none left, we can stop if len(nodes_to_visit) == 0: break def draw_spine(): nodes_to_visit = [firstnode] while True: current_node = nodes_to_visit[0] if current_node in inv_nodes: for c in inv_nodes[current_node]: # inv_node => p->c if depth[c] == depth[current_node] - 1: if not c in nodes_to_visit: nodes_to_visit.append(c) line_style = svg_spine_line_style svg_add_line( (w_margin+w_no_margs*(positions[current_node][0]-min_width)/(max_width-min_width), h_margin+h_no_margs*positions[current_node][1]/max_height), (w_margin+w_no_margs*(positions[c][0]-min_width)/(max_width-min_width), h_margin+h_no_margs*positions[c][1]/max_height), line_style) # we remove the current node from the list nodes_to_visit = nodes_to_visit[1:] # if none left, we can stop if len(nodes_to_visit) == 0: break def draw_skeleton(): nodes_to_visit = [firstnode] while True: current_node = nodes_to_visit[0] if current_node in inv_nodes: for c in inv_nodes[current_node]: # inv_node => p->c if depth[c] >= min_skeleton_depth: if not c in nodes_to_visit: nodes_to_visit.append(c) line_style = svg_spine_line_style svg_add_line( (w_margin+w_no_margs*(positions[current_node][0]-min_width)/(max_width-min_width), h_margin+h_no_margs*positions[current_node][1]/max_height), (w_margin+w_no_margs*(positions[c][0]-min_width)/(max_width-min_width), h_margin+h_no_margs*positions[c][1]/max_height), line_style) # we remove the current node from the list nodes_to_visit = nodes_to_visit[1:] # if none left, we can stop if len(nodes_to_visit) == 0: break # ------------------------------------ def draw_scale(filename ,type): svg_add_text("Generated from " + filename.split("\\")[-1], (5, 15), "start") svg_add_line( (w*0.7, h_margin), (w, h_margin), svg_scale_line_style) start_text = "" if TIME == "BIRTHS": start_text = "Birth #" + str(min([int(k[1:]) for k, v in nodes.items()])) if TIME == "REAL": start_text = "Time " + str(min([v for k, v in time.items()])) if TIME == "GENERATIONAL": start_text = "Depth " + str(min([v for k, v in depth.items()])) svg_add_text( start_text, (w, h_margin + 15), "end") svg_add_line( (w*0.7, h-h_margin), (w, h-h_margin), svg_scale_line_style) end_text = "" if TIME == "BIRTHS": end_text = "Birth #" + str(max([int(k[1:]) for k, v in nodes.items()])) if TIME == "REAL": end_text = "Time " + str(max([v for k, v in time.items()])) if TIME == "GENERATIONAL": end_text = "Depth " + str(max([v for k, v in depth.items()])) svg_add_text( end_text, (w, h-h_margin + 15), "end") ##################################################### main ##################################################### args = 0 h = 800 w = 600 h_margin = 20 w_margin = 10 h_no_margs = h - 2* h_margin w_no_margs = w - 2* w_margin max_height = 0 max_width = 0 min_width = 9999999999 min_skeleton_depth = 0 max_nodes = 0 firstnode = None nodes = {} inv_nodes = {} positions = {} visited= {} depth = {} time = {} kind = {} def main(): global svg_file, min_skeleton_depth, max_nodes, args, \ TIME, BALANCE, DOT_STYLE, COLORING, JITTER, \ svg_mutation_line_style, svg_crossover_line_style parser = argparse.ArgumentParser(description='Draws a genealogical tree (generates a SVG file) based on parent-child relationship information from a text file. Supports files generated by Framsticks experiments.') parser.add_argument('-i', '--in', dest='input', required=True, help='input file name with stuctured evolutionary data') parser.add_argument('-o', '--out', dest='output', required=True, help='output file name for the evolutionary tree (SVG format)') draw_tree_parser = parser.add_mutually_exclusive_group(required=False) draw_tree_parser.add_argument('--draw-tree', dest='draw_tree', action='store_true', help='whether drawing the full tree should be skipped') draw_tree_parser.add_argument('--no-draw-tree', dest='draw_tree', action='store_false') draw_skeleton_parser = parser.add_mutually_exclusive_group(required=False) draw_skeleton_parser.add_argument('--draw-skeleton', dest='draw_skeleton', action='store_true', help='whether the skeleton of the tree should be drawn') draw_skeleton_parser.add_argument('--no-draw-skeleton', dest='draw_skeleton', action='store_false') draw_spine_parser = parser.add_mutually_exclusive_group(required=False) draw_spine_parser.add_argument('--draw-spine', dest='draw_spine', action='store_true', help='whether the spine of the tree should be drawn') draw_spine_parser.add_argument('--no-draw-spine', dest='draw_spine', action='store_false') #TODO: better names for those parameters parser.add_argument('-t', '--time', default='GENERATIONAL', dest='time', help='values on vertical axis (BIRTHS/GENERATIONAL(d)/REAL); ' 'BIRTHS: time measured as the number of births since the beginning; ' 'GENERATIONAL: time measured as number of ancestors; ' 'REAL: real time of the simulation') parser.add_argument('-b', '--balance', default='DENSITY', dest='balance', help='method of placing nodes in the tree (RANDOM/MIN/DENSITY(d))') parser.add_argument('-s', '--scale', default='NONE', dest='scale', help='type of timescale added to the tree (NONE(d)/SIMPLE)') parser.add_argument('-c', '--coloring', default='IMPORTANCE', dest="coloring", help='method of coloring the tree (NONE/IMPORTANCE(d)/TYPE)') parser.add_argument('-d', '--dots', default='TYPE', dest='dots', help='method of drawing dots (individuals) (NONE/NORMAL/TYPE(d))') parser.add_argument('-j', '--jitter', dest="jitter", action='store_true', help='draw horizontal positions of children from the normal distribution') parser.add_argument('--color-mut', default="#000000", dest="color_mut", help='color of clone/mutation lines in rgba (e.g. #FF60B240) for TYPE coloring') parser.add_argument('--color-cross', default="#660198", dest="color_cross", help='color of crossover lines in rgba (e.g. #FF60B240) for TYPE coloring') parser.add_argument('--min-skeleton-depth', type=int, default=2, dest='min_skeleton_depth', help='minimal distance from the leafs for the nodes in the skeleton') parser.add_argument('--seed', type=int, dest='seed', help='seed for the random number generator (-1 for random)') parser.add_argument('--simple-data', type=bool, dest='simple_data', help='input data are given in a simple format (#child #parent)') parser.add_argument('-x', '--max-nodes', type=int, default=0, dest='max_nodes', help='maximum number of nodes drawn (starting from the first one)') parser.set_defaults(draw_tree=True) parser.set_defaults(draw_skeleton=False) parser.set_defaults(draw_spine=False) parser.set_defaults(seed=-1) args = parser.parse_args() TIME = args.time.upper() BALANCE = args.balance.upper() DOT_STYLE = args.dots.upper() COLORING = args.coloring.upper() SCALE = args.scale.upper() JITTER = args.jitter if not TIME in ['BIRTHS', 'GENERATIONAL', 'REAL']\ or not BALANCE in ['RANDOM', 'MIN', 'DENSITY']\ or not DOT_STYLE in ['NONE', 'NORMAL', 'TYPE']\ or not COLORING in ['NONE', 'IMPORTANCE', 'TYPE']\ or not SCALE in ['NONE', 'SIMPLE']: print("Incorrect value of one of the parameters! Closing the program.") #TODO don't be lazy, figure out which parameter is wrong... return svg_mutation_line_style += hex_to_style(args.color_mut) svg_crossover_line_style += hex_to_style(args.color_cross) dir = args.input min_skeleton_depth = args.min_skeleton_depth max_nodes = args.max_nodes seed = args.seed if seed == -1: seed = random.randint(0, 10000) random.seed(seed) print("seed:", seed) if args.simple_data: load_simple_data(dir) else: load_data(dir) compute_depth(firstnode) svg_file = open(args.output, "w") svg_file.write('') prepos_children() if args.draw_tree: draw_children() if args.draw_skeleton: draw_skeleton() if args.draw_spine: draw_spine() draw_scale(dir, SCALE) svg_file.write("") svg_file.close() main()