# Draws a genealogical tree (generates a SVG file) based on parent-child relationship information. import json import random import math import argparse TIME = "GENERATIONAL" BALANCE = "MIN" # ------SVG--------- svg_file = 0 svg_line_style = 'stroke="rgb(90%,10%,16%)" stroke-width="1" stroke-opacity="0.8"' svg_dot_style = 'r="2" stroke="black" stroke-width="0.2" fill="red"' svg_spine_line_style = 'stroke="rgb(0%,0%,80%)" stroke-width="2" stroke-opacity="1"' #svg_spine_dot_style = 'r="3" stroke="black" stroke-width="0.4" fill="rgb(50%,50%,100%)"' svg_spine_dot_style = 'r="1" stroke="black" stroke-width="0.2" fill="rgb(50%,50%,100%)"' def svg_add_line(from_pos, to_pos, style=svg_line_style): svg_file.write('') 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] # 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(percent): from_col = [100, 90, 0] to_col = [60, 0, 0] # neon # from_col = [30, 200, 255] # to_col = [240, 0, 220] from_opa = 0.2 to_opa = 1.0 opa = from_opa*(1-percent) + to_opa*percent percent = 1 - ((1-percent)**20) return 'fill="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) + '%)" r="1.5" stroke="black" stroke-width="0.2" fill-opacity="' + str(opa) + '" ' \ 'stroke-opacity="' + str(opa) + '"' # ------------------- def load_data(dir): global firstnode, nodes, inv_nodes f = open(dir) for line in f: sline = line.split(' ', 2) if len(sline) == 3: if sline[1] == "[OFFSPRING]": creature = json.loads(sline[2]) #print("B" +str(creature)) if "FromIDs" in creature: assert(len(creature["FromIDs"]) == 1) nodes[creature["ID"]] = creature["FromIDs"][0] if not creature["FromIDs"][0] in nodes: firstnode = creature["FromIDs"][0] if "Time" in creature: time[creature["ID"]] = creature["Time"] for k, v in sorted(nodes.items()): inv_nodes[v] = inv_nodes.get(v, []) inv_nodes[v].append(k) def load_simple_data(dir): global firstnode, nodes, inv_nodes f = open(dir) 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] 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 for pos in positions: pos = positions[pos] if pos[1] > y-10 or pos[1] < y+10: dy = pos[1]-y dx1 = pos[0]-x1 dx2 = pos[0]-x2 x1_dist += math.sqrt(dy**2 + dx1**2) x2_dist += math.sqrt(dy**2 + dx2**2) return (x1 if x1_dist > x2_dist else x2) # ------------------------------------ def prepos_children_reccurent(node): for c in inv_nodes[node]: #print(node + "->" + c) dissimilarity = 0.5 #random.gauss(0,0.3) if TIME == "BIRTHS": id = "" if c[0] == "c": id = int(c[1:]) else: id = int(c) positions[c] = [xmin_crowd(positions[node][0]-dissimilarity, positions[node][0]+dissimilarity, id), id] elif TIME == "GENERATIONAL": positions[c] = [xmin_crowd(positions[node][0]-dissimilarity, positions[node][0]+dissimilarity, positions[node][1]+1), positions[node][1]+1] elif TIME == "REAL": positions[c] = [xmin_crowd(positions[node][0]-dissimilarity, positions[node][0]+dissimilarity, time[c]), time[c]] for c in inv_nodes[node]: if c in inv_nodes: prepos_children_reccurent(c) def prepos_children(): global max_height, max_width, min_width if not bool(time): print("REAL time requested, but no real time data provided. Assuming BIRTHS time instead.") TIME = "BIRTHS" positions[firstnode] = [0, 0] 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 draw_children_recurrent(node, max_depth): global max_height, max_width, min_width for c in inv_nodes[node]: if c in inv_nodes: draw_children_recurrent(c, max_depth) line_style = (svg_line_style if args.mono_tree else svg_generate_line_style(depth[c]/max_depth)) dot_style = (svg_dot_style if args.mono_tree else svg_generate_dot_style(depth[c]/max_depth)) svg_add_line( (w_margin+w_no_margs*(positions[node][0]-min_width)/(max_width-min_width), h_margin+h_no_margs*positions[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) svg_add_dot( (w_margin+w_no_margs*(positions[c][0]-min_width)/(max_width-min_width), h_margin+h_no_margs*positions[c][1]/max_height), dot_style) def draw_children(): max_depth = 0 for k, v in depth.items(): max_depth = max(max_depth, v) draw_children_recurrent(firstnode, max_depth) dot_style = (svg_dot_style if args.mono_tree else svg_generate_dot_style(depth[firstnode]/max_depth)) svg_add_dot( (w_margin+w_no_margs*(positions[firstnode][0]-min_width)/(max_width-min_width), h_margin+h_no_margs*positions[firstnode][1]/max_height), dot_style) def draw_spine_recurrent(node): global max_height, max_width, min_width for c in inv_nodes[node]: if depth[c] == depth[node] - 1: if c in inv_nodes: draw_spine_recurrent(c) line_style = svg_spine_line_style svg_add_line( (w_margin+w_no_margs*(positions[node][0]-min_width)/(max_width-min_width), h_margin+h_no_margs*positions[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) #svg_add_dot( (w_margin+w_no_margs*(positions[c][0]-min_width)/(max_width-min_width), h_margin+h_no_margs*positions[c][1]/max_height), svg_spine_dot_style) def draw_spine(): draw_spine_recurrent(firstnode) #svg_add_dot( (w_margin+w_no_margs*(positions[firstnode][0]-min_width)/(max_width-min_width), h_margin+h_no_margs*positions[firstnode][1]/max_height), svg_spine_dot_style) def draw_skeleton_reccurent(node, max_depth): global max_height, max_width, min_width for c in inv_nodes[node]: if depth[c] >= min_skeleton_depth or depth[c] == max([depth[q] for q in inv_nodes[node]]): if c in inv_nodes: draw_skeleton_reccurent(c, max_depth) line_style = svg_spine_line_style svg_add_line( (w_margin+w_no_margs*(positions[node][0]-min_width)/(max_width-min_width), h_margin+h_no_margs*positions[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) #svg_add_dot( (w_margin+w_no_margs*(positions[c][0]-min_width)/(max_width-min_width), h_margin+h_no_margs*positions[c][1]/max_height), # svg_spine_dot_style) def draw_skeleton(): max_depth = 0 for k, v in depth.items(): max_depth = max(max_depth, v) draw_skeleton_reccurent(firstnode, max_depth) #svg_add_dot( (w_margin+w_no_margs*(positions[firstnode][0]-min_width)/(max_width-min_width), h_margin+h_no_margs*positions[firstnode][1]/max_height), # svg_spine_dot_style) ##################################################### main ##################################################### args = 0 h = 800 w = 600 h_margin = 10 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 firstnode = "" nodes = {} inv_nodes = {} positions = {} depth = {} time = {} def main(): global svg_file, min_skeleton_depth, args, TIME, BALANCE parser = argparse.ArgumentParser(description='Process some integers.') parser.add_argument('--in', dest='input', required=True, help='input file with stuctured evolutionary data') parser.add_argument('--out', dest='output', required=True, help='output file for the evolutionary tree') 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('--time', default='BIRTHS', dest='time', help='values on vertical axis (REAL/GENERATIONAL)') parser.add_argument('--balance', default='MIN',dest='balance', help='method of placing node in the tree (RANDOM/MIN/DENSITY)') mono_tree_parser = parser.add_mutually_exclusive_group(required=False) mono_tree_parser.add_argument('--mono-tree', dest='mono_tree', action='store_true', help='whether the tree should be drawn with a single color') mono_tree_parser.add_argument('--no-mono-tree', dest='mono_tree', action='store_false') 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.set_defaults(mono_tree=False) 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 BALANCE = args.balance dir = args.input min_skeleton_depth = args.min_skeleton_depth 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() svg_file.write("") svg_file.close() main()