source: mds-and-trees/tree-genealogy.py @ 709

Last change on this file since 709 was 709, checked in by konrad, 6 years ago

tree-genealogy.py should now work correctly with the order of events in the logs BORN -> OFFSPRING -> DIED

File size: 38.7 KB
Line 
1import json
2import math
3import random
4import argparse
5import bisect
6import copy
7import time as timelib
8from PIL import Image, ImageDraw, ImageFont
9from scipy import stats
10from matplotlib import colors
11import numpy as np
12
13class LoadingError(Exception):
14    pass
15
16class Drawer:
17
18    def __init__(self, design, config_file, w=600, h=800, w_margin=10, h_margin=20):
19        self.design = design
20        self.width = w
21        self.height = h
22        self.w_margin = w_margin
23        self.h_margin = h_margin
24        self.w_no_margs = w - 2* w_margin
25        self.h_no_margs = h - 2* h_margin
26
27        self.color_converter = colors.ColorConverter()
28
29        self.settings = {
30            'colors_of_kinds': ['red', 'green', 'blue', 'magenta', 'yellow', 'cyan', 'orange', 'purple'],
31            'dots': {
32                'color': {
33                    'meaning': 'Lifespan',
34                    'normalize_cmap': False,
35                    'cmap': {},
36                    'start': 'red',
37                    'end': 'green',
38                    'bias': 1
39                    },
40                'size': {
41                    'meaning': 'EnergyEaten',
42                    'start': 1,
43                    'end': 6,
44                    'bias': 0.5
45                    },
46                'opacity': {
47                    'meaning': 'EnergyEaten',
48                    'start': 0.2,
49                    'end': 1,
50                    'bias': 1
51                    }
52            },
53            'lines': {
54                'color': {
55                    'meaning': 'adepth',
56                    'normalize_cmap': False,
57                    'cmap': {},
58                    'start': 'black',
59                    'end': 'red',
60                    'bias': 3
61                    },
62                'width': {
63                    'meaning': 'adepth',
64                    'start': 0.1,
65                    'end': 4,
66                    'bias': 3
67                    },
68                'opacity': {
69                    'meaning': 'adepth',
70                    'start': 0.1,
71                    'end': 0.8,
72                    'bias': 5
73                    }
74            }
75        }
76
77        def merge(source, destination):
78            for key, value in source.items():
79                if isinstance(value, dict):
80                    node = destination.setdefault(key, {})
81                    merge(value, node)
82                else:
83                    destination[key] = value
84            return destination
85
86        if config_file != "":
87            with open(config_file) as config:
88                c = json.load(config)
89            self.settings = merge(c, self.settings)
90            #print(json.dumps(self.settings, indent=4, sort_keys=True))
91
92        self.compile_cmaps()
93
94    def compile_cmaps(self):
95        def normalize_and_compile_cmap(cmap):
96            for key in cmap:
97                for arr in cmap[key]:
98                    arr[0] = (arr[0] - cmap[key][0][0]) / (cmap[key][-1][0] - cmap[key][0][0])
99            return colors.LinearSegmentedColormap('Custom', cmap)
100
101        for part in ['dots', 'lines']:
102            if self.settings[part]['color']['cmap']:
103                if self.settings[part]['color']['normalize_cmap']:
104                    cmap = self.settings[part]['color']['cmap']
105                    min = self.design.props[self.settings[part]['color']['meaning'] + "_min"]
106                    max = self.design.props[self.settings[part]['color']['meaning'] + "_max"]
107
108                    for key in cmap:
109                        if cmap[key][0][0] > min:
110                            cmap[key].insert(0, cmap[key][0][:])
111                            cmap[key][0][0] = min
112                        if cmap[key][-1][0] < max:
113                            cmap[key].append(cmap[key][-1][:])
114                            cmap[key][-1][0] = max
115
116                    og_cmap = normalize_and_compile_cmap(copy.deepcopy(cmap))
117
118                    col2key = {'red':0, 'green':1, 'blue':2}
119                    for key in cmap:
120                        # for color from (r/g/b) #n's should be the same for all keys!
121                        n_min = (min - cmap[key][0][0]) / (cmap[key][-1][0] - cmap[key][0][0])
122                        n_max = (max - cmap[key][0][0]) / (cmap[key][-1][0] - cmap[key][0][0])
123
124                        min_col = og_cmap(n_min)
125                        max_col = og_cmap(n_max)
126
127                        cmap[key][0] = [min, min_col[col2key[key]], min_col[col2key[key]]]
128                        cmap[key][-1] = [max, max_col[col2key[key]], max_col[col2key[key]]]
129                print(self.settings[part]['color']['cmap'])
130                self.settings[part]['color']['cmap'] = normalize_and_compile_cmap(self.settings[part]['color']['cmap'])
131
132    def draw_dots(self, file, min_width, max_width, max_height):
133        for i in range(len(self.design.positions)):
134            node = self.design.positions[i]
135            if 'x' not in node:
136                continue
137            dot_style = self.compute_dot_style(node=i)
138            self.add_dot(file, (self.w_margin+self.w_no_margs*(node['x']-min_width)/(max_width-min_width),
139                               self.h_margin+self.h_no_margs*node['y']/max_height), dot_style)
140
141    def draw_lines(self, file, min_width, max_width, max_height):
142        for parent in range(len(self.design.positions)):
143            par_pos = self.design.positions[parent]
144            if not 'x' in par_pos:
145                continue
146            for child in self.design.tree.children[parent]:
147                chi_pos = self.design.positions[child]
148                if 'x' not in chi_pos:
149                    continue
150                line_style = self.compute_line_style(parent, child)
151                self.add_line(file, (self.w_margin+self.w_no_margs*(par_pos['x']-min_width)/(max_width-min_width),
152                                  self.h_margin+self.h_no_margs*par_pos['y']/max_height),
153                                  (self.w_margin+self.w_no_margs*(chi_pos['x']-min_width)/(max_width-min_width),
154                                  self.h_margin+self.h_no_margs*chi_pos['y']/max_height), line_style)
155
156    def draw_scale(self, file, filename):
157        self.add_text(file, "Generated from " + filename.split("\\")[-1], (5, 5), "start")
158
159        start_text = ""
160        end_text = ""
161        if self.design.TIME == "BIRTHS":
162           start_text = "Birth #0"
163           end_text = "Birth #" + str(len(self.design.positions)-1)
164        if self.design.TIME == "REAL":
165           start_text = "Time " + str(min(self.design.tree.time))
166           end_text = "Time " + str(max(self.design.tree.time))
167        if self.design.TIME == "GENERATIONAL":
168           start_text = "Depth " + str(self.design.props['adepth_min'])
169           end_text = "Depth " + str(self.design.props['adepth_max'])
170
171        self.add_dashed_line(file, (self.width*0.7, self.h_margin), (self.width, self.h_margin))
172        self.add_text(file, start_text, (self.width, self.h_margin), "end")
173        self.add_dashed_line(file, (self.width*0.7, self.height-self.h_margin), (self.width, self.height-self.h_margin))
174        self.add_text(file, end_text, (self.width, self.height-self.h_margin), "end")
175
176    def compute_property(self, part, prop, node):
177        start = self.settings[part][prop]['start']
178        end = self.settings[part][prop]['end']
179        value = (self.design.props[self.settings[part][prop]['meaning']][node]
180                 if self.settings[part][prop]['meaning'] in self.design.props else 0 )
181        bias = self.settings[part][prop]['bias']
182        if prop == "color":
183            if not self.settings[part][prop]['cmap']:
184                return self.compute_color(start, end, value, bias)
185            else:
186                return self.compute_color_from_cmap(self.settings[part][prop]['cmap'], value, bias)
187        else:
188            return self.compute_value(start, end, value, bias)
189
190    def compute_color_from_cmap(self, cmap, value, bias=1):
191        value = 1 - (1-value)**bias
192        rgba = cmap(value)
193        return (100*rgba[0], 100*rgba[1], 100*rgba[2])
194
195
196    def compute_color(self, start, end, value, bias=1):
197        if isinstance(value, str):
198            value = int(value)
199            r, g, b = self.color_converter.to_rgb(self.settings['colors_of_kinds'][value])
200        else:
201            start_color = self.color_converter.to_rgb(start)
202            end_color = self.color_converter.to_rgb(end)
203            value = 1 - (1-value)**bias
204            r = start_color[0]*(1-value)+end_color[0]*value
205            g = start_color[1]*(1-value)+end_color[1]*value
206            b = start_color[2]*(1-value)+end_color[2]*value
207        return (100*r, 100*g, 100*b)
208
209    def compute_value(self, start, end, value, bias=1):
210        value = 1 - (1-value)**bias
211        return start*(1-value) + end*value
212
213class PngDrawer(Drawer):
214
215    def scale_up(self):
216        self.width *= self.multi
217        self.height *= self.multi
218        self.w_margin *= self.multi
219        self.h_margin *= self.multi
220        self.h_no_margs *= self.multi
221        self.w_no_margs *= self.multi
222
223    def scale_down(self):
224        self.width /= self.multi
225        self.height /= self.multi
226        self.w_margin /= self.multi
227        self.h_margin /= self.multi
228        self.h_no_margs /= self.multi
229        self.w_no_margs /= self.multi
230
231    def draw_design(self, filename, input_filename, multi=1, scale="SIMPLE"):
232        print("Drawing...")
233
234        self.multi=multi
235        self.scale_up()
236
237        back = Image.new('RGBA', (self.width, self.height), (255,255,255,0))
238
239        min_width = min([x['x'] for x in self.design.positions if 'x' in x])
240        max_width = max([x['x'] for x in self.design.positions if 'x' in x])
241        max_height = max([x['y'] for x in self.design.positions if 'y' in x])
242
243        self.draw_lines(back, min_width, max_width, max_height)
244        self.draw_dots(back, min_width, max_width, max_height)
245
246        if scale == "SIMPLE":
247            self.draw_scale(back, input_filename)
248
249        #back.show()
250        self.scale_down()
251
252        back.thumbnail((self.width, self.height), Image.ANTIALIAS)
253
254        back.save(filename)
255
256    def add_dot(self, file, pos, style):
257        x, y = int(pos[0]), int(pos[1])
258        r = style['r']*self.multi
259        offset = (int(x - r), int(y - r))
260        size = (2*int(r), 2*int(r))
261
262        c = style['color']
263
264        img = Image.new('RGBA', size)
265        ImageDraw.Draw(img).ellipse((1, 1, size[0]-1, size[1]-1),
266                                    (int(2.55*c[0]), int(2.55*c[1]), int(2.55*c[2]), int(255*style['opacity'])))
267        file.paste(img, offset, mask=img)
268
269    def add_line(self, file, from_pos, to_pos, style):
270        fx, fy, tx, ty = int(from_pos[0]), int(from_pos[1]), int(to_pos[0]), int(to_pos[1])
271        w = int(style['width'])*self.multi
272
273        offset = (min(fx-w, tx-w), min(fy-w, ty-w))
274        size = (abs(fx-tx)+2*w, abs(fy-ty)+2*w)
275        if size[0] == 0 or size[1] == 0:
276            return
277
278        c = style['color']
279
280        img = Image.new('RGBA', size)
281        ImageDraw.Draw(img).line((w, w, size[0]-w, size[1]-w) if (fx-tx)*(fy-ty)>0 else (size[0]-w, w, w, size[1]-w),
282                                  (int(2.55*c[0]), int(2.55*c[1]), int(2.55*c[2]), int(255*style['opacity'])), w)
283        file.paste(img, offset, mask=img)
284
285    def add_dashed_line(self, file, from_pos, to_pos):
286        style = {'color': (0,0,0), 'width': 1, 'opacity': 1}
287        sublines = 50
288        # TODO could be faster: compute delta and only add delta each time (but currently we do not use it often)
289        normdiv = 2*sublines-1
290        for i in range(sublines):
291            from_pos_sub = (self.compute_value(from_pos[0], to_pos[0], 2*i/normdiv, 1),
292                            self.compute_value(from_pos[1], to_pos[1], 2*i/normdiv, 1))
293            to_pos_sub = (self.compute_value(from_pos[0], to_pos[0], (2*i+1)/normdiv, 1),
294                          self.compute_value(from_pos[1], to_pos[1], (2*i+1)/normdiv, 1))
295            self.add_line(file, from_pos_sub, to_pos_sub, style)
296
297    def add_text(self, file, text, pos, anchor, style=''):
298        font = ImageFont.truetype("Vera.ttf", 16*self.multi)
299
300        img = Image.new('RGBA', (self.width, self.height))
301        draw = ImageDraw.Draw(img)
302        txtsize = draw.textsize(text, font=font)
303        pos = pos if anchor == "start" else (pos[0]-txtsize[0], pos[1])
304        draw.text(pos, text, (0,0,0), font=font)
305        file.paste(img, (0,0), mask=img)
306
307    def compute_line_style(self, parent, child):
308        return {'color': self.compute_property('lines', 'color', child),
309                'width': self.compute_property('lines', 'width', child),
310                'opacity': self.compute_property('lines', 'opacity', child)}
311
312    def compute_dot_style(self, node):
313        return {'color': self.compute_property('dots', 'color', node),
314                'r': self.compute_property('dots', 'size', node),
315                'opacity': self.compute_property('dots', 'opacity', node)}
316
317class SvgDrawer(Drawer):
318    def draw_design(self, filename, input_filename, multi=1, scale="SIMPLE"):
319        print("Drawing...")
320        file = open(filename, "w")
321
322        min_width = min([x['x'] for x in self.design.positions if 'x' in x])
323        max_width = max([x['x'] for x in self.design.positions if 'x' in x])
324        max_height = max([x['y'] for x in self.design.positions if 'y' in x])
325
326        file.write('<svg xmlns:svg="http://www.w3.org/2000/svg" xmlns="http://www.w3.org/2000/svg" '
327                   'xmlns:xlink="http://www.w3.org/1999/xlink" version="1.0" '
328                   'width="' + str(self.width) + '" height="' + str(self.height) + '">')
329
330        self.draw_lines(file, min_width, max_width, max_height)
331        self.draw_dots(file, min_width, max_width, max_height)
332
333        if scale == "SIMPLE":
334            self.draw_scale(file, input_filename)
335
336        file.write("</svg>")
337        file.close()
338
339    def add_text(self, file, text, pos, anchor, style=''):
340        style = (style if style != '' else 'style="font-family: Arial; font-size: 12; fill: #000000;"')
341        # assuming font size 12, it should be taken from the style string!
342        file.write('<text ' + style + ' text-anchor="' + anchor + '" x="' + str(pos[0]) + '" y="' + str(pos[1]+12) + '" >' + text + '</text>')
343
344    def add_dot(self, file, pos, style):
345        file.write('<circle ' + style + ' cx="' + str(pos[0]) + '" cy="' + str(pos[1]) + '" />')
346
347    def add_line(self, file, from_pos, to_pos, style):
348        file.write('<line ' + style + ' x1="' + str(from_pos[0]) + '" x2="' + str(to_pos[0]) +
349                       '" y1="' + str(from_pos[1]) + '" y2="' + str(to_pos[1]) + '"  fill="none"/>')
350
351    def add_dashed_line(self, file, from_pos, to_pos):
352        style = 'stroke="black" stroke-width="0.5" stroke-opacity="1" stroke-dasharray="5, 5"'
353        self.add_line(file, from_pos, to_pos, style)
354
355    def compute_line_style(self, parent, child):
356        return self.compute_stroke_color('lines', child) + ' ' \
357               + self.compute_stroke_width('lines', child) + ' ' \
358               + self.compute_stroke_opacity(child)
359
360    def compute_dot_style(self, node):
361        return self.compute_dot_size(node) + ' ' \
362               + self.compute_fill_opacity(node) + ' ' \
363               + self.compute_dot_fill(node)
364
365    def compute_stroke_color(self, part, node):
366        color = self.compute_property(part, 'color', node)
367        return 'stroke="rgb(' + str(color[0]) + '%,' + str(color[1]) + '%,' + str(color[2]) + '%)"'
368
369    def compute_stroke_width(self, part, node):
370        return 'stroke-width="' + str(self.compute_property(part, 'width', node)) + '"'
371
372    def compute_stroke_opacity(self, node):
373        return 'stroke-opacity="' + str(self.compute_property('lines', 'opacity', node)) + '"'
374
375    def compute_fill_opacity(self, node):
376        return 'fill-opacity="' + str(self.compute_property('dots', 'opacity', node)) + '"'
377
378    def compute_dot_size(self, node):
379        return 'r="' + str(self.compute_property('dots', 'size', node)) + '"'
380
381    def compute_dot_fill(self, node):
382        color = self.compute_property('dots', 'color', node)
383        return 'fill="rgb(' + str(color[0]) + '%,' + str(color[1]) + '%,' + str(color[2]) + '%)"'
384
385class Designer:
386
387    def __init__(self, tree, jitter=False, time="GENERATIONAL", balance="DENSITY"):
388        self.props = {}
389
390        self.tree = tree
391
392        self.TIME = time
393        self.JITTER = jitter
394
395        if balance == "RANDOM":
396            self.xmin_crowd = self.xmin_crowd_random
397        elif balance == "MIN":
398            self.xmin_crowd = self.xmin_crowd_min
399        elif balance == "DENSITY":
400            self.xmin_crowd = self.xmin_crowd_density
401        else:
402            raise ValueError("Error, the value of BALANCE does not match any expected value.")
403
404    def calculate_measures(self):
405        print("Calculating measures...")
406        self.compute_depth()
407        self.compute_maxdepth()
408        self.compute_adepth()
409        self.compute_children()
410        self.compute_kind()
411        self.compute_time()
412        self.compute_progress()
413        self.compute_custom()
414
415    def xmin_crowd_random(self, x1, x2, y):
416        return (x1 if random.randrange(2) == 0 else x2)
417
418    def xmin_crowd_min(self, x1, x2, y):
419        x1_closest = 999999
420        x2_closest = 999999
421        miny = y-3
422        maxy = y+3
423        i = bisect.bisect_left(self.y_sorted, miny)
424        while True:
425            if len(self.positions_sorted) <= i or self.positions_sorted[i]['y'] > maxy:
426                break
427            pos = self.positions_sorted[i]
428
429            x1_closest = min(x1_closest, abs(x1-pos['x']))
430            x2_closest = min(x2_closest, abs(x2-pos['x']))
431
432            i += 1
433        return (x1 if x1_closest > x2_closest else x2)
434
435    def xmin_crowd_density(self, x1, x2, y):
436        # TODO experimental - requires further work to make it less 'jumpy' and more predictable
437        CONST_LOCAL_AREA_RADIUS = 5
438        CONST_GLOBAL_AREA_RADIUS = 10
439        CONST_WINDOW_SIZE = 20000 #TODO should depend on the maxY ?
440        x1_dist_loc = 0
441        x2_dist_loc = 0
442        count_loc = 1
443        x1_dist_glob = 0
444        x2_dist_glob = 0
445        count_glob = 1
446        miny = y-CONST_WINDOW_SIZE
447        maxy = y+CONST_WINDOW_SIZE
448        i_left = bisect.bisect_left(self.y_sorted, miny)
449        i_right = bisect.bisect_right(self.y_sorted, maxy)
450        #TODO test: maxy=y should give the same results, right?
451
452        def include_pos(pos):
453            nonlocal x1_dist_loc, x2_dist_loc, x1_dist_glob, x2_dist_glob, count_loc, count_glob
454
455            dysq = (pos['y']-y)**2 + 1 #+1 so 1/dysq is at most 1
456            dx1 = math.fabs(pos['x']-x1)
457            dx2 = math.fabs(pos['x']-x2)
458
459            d = math.fabs(pos['x'] - (x1+x2)/2)
460
461            if d < CONST_LOCAL_AREA_RADIUS:
462                x1_dist_loc += math.sqrt(dx1/dysq + dx1**2)
463                x2_dist_loc += math.sqrt(dx2/dysq + dx2**2)
464                count_loc += 1
465            elif d > CONST_GLOBAL_AREA_RADIUS:
466                x1_dist_glob += math.sqrt(dx1/dysq + dx1**2)
467                x2_dist_glob += math.sqrt(dx2/dysq + dx2**2)
468                count_glob += 1
469
470        # optimized to draw from all the nodes, if less than 10 nodes in the range
471        if len(self.positions_sorted) > i_left:
472            if i_right - i_left < 10:
473                for j in range(i_left, i_right):
474                    include_pos(self.positions_sorted[j])
475            else:
476                for j in range(10):
477                    pos = self.positions_sorted[random.randrange(i_left, i_right)]
478                    include_pos(pos)
479
480        return (x1 if (x1_dist_loc-x2_dist_loc)/count_loc-(x1_dist_glob-x2_dist_glob)/count_glob > 0  else x2)
481        #return (x1 if x1_dist +random.gauss(0, 0.00001) > x2_dist +random.gauss(0, 0.00001)  else x2)
482        #print(x1_dist, x2_dist)
483        #x1_dist = x1_dist**2
484        #x2_dist = x2_dist**2
485        #return x1 if x1_dist+x2_dist==0 else (x1*x1_dist + x2*x2_dist) / (x1_dist+x2_dist) + random.gauss(0, 0.01)
486        #return (x1 if random.randint(0, int(x1_dist+x2_dist)) < x1_dist else x2)
487
488    def calculate_node_positions(self, ignore_last=0):
489        print("Calculating positions...")
490
491        def add_node(node):
492            index = bisect.bisect_left(self.y_sorted, node['y'])
493            self.y_sorted.insert(index, node['y'])
494            self.positions_sorted.insert(index, node)
495            self.positions[node['id']] = node
496
497        self.positions_sorted = [{'x':0, 'y':0, 'id':0}]
498        self.y_sorted = [0]
499        self.positions = [{} for x in range(len(self.tree.parents))]
500        self.positions[0] = {'x':0, 'y':0, 'id':0}
501
502        # order by maximum depth of the parent guarantees that co child is evaluated before its parent
503        visiting_order = [i for i in range(0, len(self.tree.parents))]
504        visiting_order = sorted(visiting_order, key=lambda q:\
505                            0 if q == 0 else self.props["maxdepth"][q])
506
507        start_time = timelib.time()
508
509        # for each child of the current node
510        for node_counter,child in enumerate(visiting_order, start=1):
511            # debug info - elapsed time
512            if node_counter % 100000 == 0:
513               print("%d%%\t%d\t%g" % (node_counter*100/len(self.tree.parents), node_counter, timelib.time()-start_time))
514               start_time = timelib.time()
515
516            # using normalized adepth
517            if self.props['adepth'][child] >= ignore_last/self.props['adepth_max']:
518
519                ypos = 0
520                if self.TIME == "BIRTHS":
521                    ypos = child
522                elif self.TIME == "GENERATIONAL":
523                    # one more than its parent (what if more than one parent?)
524                    ypos = max([self.positions[par]['y'] for par, v in self.tree.parents[child].items()])+1 \
525                        if self.tree.parents[child] else 0
526                elif self.TIME == "REAL":
527                    ypos = self.tree.time[child]
528
529                if len(self.tree.parents[child]) == 1:
530                # if current_node is the only parent
531                    parent, similarity = [(par, v) for par, v in self.tree.parents[child].items()][0]
532
533                    if self.JITTER:
534                        dissimilarity = (1-similarity) + random.gauss(0, 0.01) + 0.001
535                    else:
536                        dissimilarity = (1-similarity) + 0.001
537                    add_node({'id':child, 'y':ypos, 'x':
538                             self.xmin_crowd(self.positions[parent]['x']-dissimilarity,
539                              self.positions[parent]['x']+dissimilarity, ypos)})
540                else:
541                    # position weighted by the degree of inheritence from each parent
542                    total_inheretance = sum([v for k, v in self.tree.parents[child].items()])
543                    xpos = sum([self.positions[k]['x']*v/total_inheretance
544                               for k, v in self.tree.parents[child].items()])
545                    if self.JITTER:
546                        add_node({'id':child, 'y':ypos, 'x':xpos + random.gauss(0, 0.1)})
547                    else:
548                        add_node({'id':child, 'y':ypos, 'x':xpos})
549
550
551    def compute_custom(self):
552        for prop in self.tree.props:
553            self.props[prop] = [None for x in range(len(self.tree.children))]
554
555            for i in range(len(self.props[prop])):
556                self.props[prop][i] = self.tree.props[prop][i]
557
558            self.normalize_prop(prop)
559
560    def compute_time(self):
561        # simple rewrite from the tree
562        self.props["time"] = [0 for x in range(len(self.tree.children))]
563
564        for i in range(len(self.props['time'])):
565            self.props['time'][i] = self.tree.time[i]
566
567        self.normalize_prop('time')
568
569    def compute_kind(self):
570        # simple rewrite from the tree
571        self.props["kind"] = [0 for x in range(len(self.tree.children))]
572
573        for i in range (len(self.props['kind'])):
574            self.props['kind'][i] = str(self.tree.kind[i])
575
576    def compute_depth(self):
577        self.props["depth"] = [999999999 for x in range(len(self.tree.children))]
578        visited = [0 for x in range(len(self.tree.children))]
579
580        nodes_to_visit = [0]
581        visited[0] = 1
582        self.props["depth"][0] = 0
583        while True:
584            current_node = nodes_to_visit[0]
585
586            for child in self.tree.children[current_node]:
587                if visited[child] == 0:
588                    visited[child] = 1
589                    nodes_to_visit.append(child)
590                    self.props["depth"][child] = self.props["depth"][current_node]+1
591            nodes_to_visit = nodes_to_visit[1:]
592            if len(nodes_to_visit) == 0:
593                break
594
595        self.normalize_prop('depth')
596
597    def compute_maxdepth(self):
598        self.props["maxdepth"] = [999999999 for x in range(len(self.tree.children))]
599        visited = [0 for x in range(len(self.tree.children))]
600
601        nodes_to_visit = [0]
602        visited[0] = 1
603        self.props["maxdepth"][0] = 0
604        while True:
605            current_node = nodes_to_visit[0]
606
607            for child in self.tree.children[current_node]:
608                if visited[child] == 0:
609                    visited[child] = 1
610                    nodes_to_visit.append(child)
611                    self.props["maxdepth"][child] = self.props["maxdepth"][current_node]+1
612                elif self.props["maxdepth"][child] < self.props["maxdepth"][current_node]+1:
613                    self.props["maxdepth"][child] = self.props["maxdepth"][current_node]+1
614                    if child not in  nodes_to_visit:
615                        nodes_to_visit.append(child)
616
617            nodes_to_visit = nodes_to_visit[1:]
618            if len(nodes_to_visit) == 0:
619                break
620
621        self.normalize_prop('maxdepth')
622
623    def compute_adepth(self):
624        self.props["adepth"] = [0 for x in range(len(self.tree.children))]
625
626        # order by maximum depth of the parent guarantees that co child is evaluated before its parent
627        visiting_order = [i for i in range(0, len(self.tree.parents))]
628        visiting_order = sorted(visiting_order, key=lambda q: self.props["maxdepth"][q])[::-1]
629
630        for node in visiting_order:
631            children = self.tree.children[node]
632            if len(children) != 0:
633                # 0 by default
634                self.props["adepth"][node] = max([self.props["adepth"][child] for child in children])+1
635        self.normalize_prop('adepth')
636
637    def compute_children(self):
638        self.props["children"] = [0 for x in range(len(self.tree.children))]
639        for i in range (len(self.props['children'])):
640            self.props['children'][i] = len(self.tree.children[i])
641
642        self.normalize_prop('children')
643
644    def compute_progress(self):
645        self.props["progress"] = [0 for x in range(len(self.tree.children))]
646        for i in range(len(self.props['children'])):
647            times = sorted([self.props["time"][self.tree.children[i][j]]*100000 for j in range(len(self.tree.children[i]))])
648            if len(times) > 4:
649                times = [times[i+1] - times[i] for i in range(len(times)-1)]
650                #print(times)
651                slope, intercept, r_value, p_value, std_err = stats.linregress(range(len(times)), times)
652                self.props['progress'][i] = slope if not np.isnan(slope) and not np.isinf(slope) else 0
653
654        for i in range(0, 5):
655            self.props['progress'][self.props['progress'].index(min(self.props['progress']))] = 0
656            self.props['progress'][self.props['progress'].index(max(self.props['progress']))] = 0
657
658        mini = min(self.props['progress'])
659        maxi = max(self.props['progress'])
660        for k in range(len(self.props['progress'])):
661            if self.props['progress'][k] == 0:
662                self.props['progress'][k] = mini
663
664        #for k in range(len(self.props['progress'])):
665        #        self.props['progress'][k] = 1-self.props['progress'][k]
666
667        self.normalize_prop('progress')
668
669    def normalize_prop(self, prop):
670        noneless = [v for v in self.props[prop] if (type(v)!=str and type(v)!=list)]
671        if len(noneless) > 0:
672            max_val = max(noneless)
673            min_val = min(noneless)
674            print("%s: [%g, %g]" % (prop, min_val, max_val))
675            self.props[prop +'_max'] = max_val
676            self.props[prop +'_min'] = min_val
677            for i in range(len(self.props[prop])):
678                if self.props[prop][i] is not None:
679                    qqq = self.props[prop][i]
680                    self.props[prop][i] = 0 if max_val == min_val else (self.props[prop][i] - min_val) / (max_val - min_val)
681
682class TreeData:
683    simple_data = None
684
685    children = []
686    parents = []
687    time = []
688    kind = []
689
690    def __init__(self): #, simple_data=False):
691        #self.simple_data = simple_data
692        pass
693
694    def load(self, filename, max_nodes=0):
695        print("Loading...")
696
697        CLI_PREFIX = "Script.Message:"
698        default_props = ["Time", "FromIDs", "ID", "Operation", "Inherited"]
699
700        merged_with_virtual_parent = [] #this list will contain individuals for which the parent could not be found
701
702        self.ids = {}
703        def get_id(id, createOnError = True):
704            if createOnError:
705                if id not in self.ids:
706                    self.ids[id] = len(self.ids)
707            else:
708                if id not in self.ids:
709                    return None
710
711            return self.ids[id]
712
713        def try_to_load(input):
714            creature = False
715            try:
716                creature = json.loads(input)
717            except ValueError:
718                print("Json format error - the line cannot be read. Breaking the loading loop.")
719                # fixing arrays by removing the last element
720                # ! assuming that only the last line is broken !
721                self.parents.pop()
722                self.children.pop()
723                self.time.pop()
724                self.kind.pop()
725                self.life_lenght.pop()
726            return creature
727
728        file = open(filename)
729
730        # counting the number of expected nodes
731        nodes = 0
732        for line in file:
733            line_arr = line.split(' ', 1)
734            if len(line_arr) == 2:
735                if line_arr[0] == CLI_PREFIX:
736                    line_arr = line_arr[1].split(' ', 1)
737                if line_arr[0] == "[BORN]":
738                    nodes += 1
739
740        nodes = min(nodes, max_nodes if max_nodes != 0 else nodes)+1
741        self.parents = [{} for x in range(nodes)]
742        self.children = [[] for x in range(nodes)]
743        self.time = [0] * nodes
744        self.kind = [0] * nodes
745        self.life_lenght = [0] * nodes
746        self.props = {}
747
748        print("nodes: %d" % len(self.parents))
749
750        file.seek(0)
751        loaded_so_far = 0
752        lasttime = timelib.time()
753
754        get_id("virtual_parent")
755
756        for line in file:
757            line_arr = line.split(' ', 1)
758            if len(line_arr) == 2:
759                if line_arr[0] == CLI_PREFIX:
760                    line_arr = line_arr[1].split(' ', 1)
761                if line_arr[0] == "[BORN]":
762                    creature = try_to_load(line_arr[1])
763                    if not creature:
764                        nodes -= 1
765                        break
766
767                    creature_id = get_id(creature["ID"])
768
769                    # debug
770                    if loaded_so_far%1000 == 0:
771                        #print(". " + str(creature_id) + " " + str(timelib.time() - lasttime))
772                        lasttime = timelib.time()
773
774                    if "Time" in creature:
775                        self.time[creature_id] = creature["Time"]
776
777                    for prop in creature:
778                        if prop not in default_props:
779                            if prop not in self.props:
780                                self.props[prop] = [0 for i in range(nodes)]
781                            self.props[prop][creature_id] = creature[prop]
782
783                    loaded_so_far += 1
784
785                if line_arr[0] == "[OFFSPRING]":
786                    creature = try_to_load(line_arr[1])
787                    if not creature:
788                        nodes -= 1
789                        break
790
791                    if "FromIDs" in creature:
792                        # make sure that ID's of parents are lower than that of their children
793                        for i in range(0, len(creature["FromIDs"])):
794                            if creature["FromIDs"][i] not in self.ids:
795                                get_id("virtual_parent")
796
797                        creature_id = get_id(creature["ID"])#, False)
798
799                        # we assign to each parent its contribution to the genotype of the child
800                        for i in range(0, len(creature["FromIDs"])):
801                            if creature["FromIDs"][i] in self.ids:
802                                parent_id = get_id(creature["FromIDs"][i])
803                            else:
804                                if creature["FromIDs"][i] not in merged_with_virtual_parent:
805                                    merged_with_virtual_parent.append(creature["FromIDs"][i])
806                                parent_id = get_id("virtual_parent")
807                            inherited = (creature["Inherited"][i] if 'Inherited' in creature else 1)
808                            self.parents[creature_id][parent_id] = inherited
809
810                        if "Kind" in creature:
811                            self.kind[creature_id] = creature["Kind"]
812                    else:
813                        raise LoadingError("[OFFSPRING] misses the 'FromIDs' field!")
814
815                if line_arr[0] == "[DIED]" or line_arr[0] == "[BORN]":
816                    creature = try_to_load(line_arr[1])
817                    if not creature:
818                        nodes -= 1
819                        break
820                    creature_id = get_id(creature["ID"], False)
821                    if creature_id is not None:
822                        for prop in creature:
823                            if prop not in default_props:
824                                if prop not in self.props:
825                                    self.props[prop] = [0 for i in range(nodes)]
826                                self.props[prop][creature_id] = creature[prop]
827
828
829            if loaded_so_far >= max_nodes and max_nodes != 0:
830                break
831
832        print("NOTE: all individuals with parent not provided or missing were connected to a single 'virtual parent' node: " + str(merged_with_virtual_parent))
833
834        for c_id in range(1, nodes):
835            if not self.parents[c_id]:
836                self.parents[c_id][get_id("virtual_parent")] = 1
837
838        for k in range(len(self.parents)):
839            v = self.parents[k]
840            for val in self.parents[k]:
841                self.children[val].append(k)
842
843depth = {}
844kind = {}
845
846def main():
847
848    parser = argparse.ArgumentParser(description='Draws a genealogical tree (generates a SVG file) based on parent-child relationship '
849                                                 'information from a text file. Supports files generated by Framsticks experiments.')
850    parser.add_argument('-i', '--in', dest='input', required=True, help='input file name with stuctured evolutionary data')
851    parser.add_argument('-o', '--out', dest='output', required=True, help='output file name for the evolutionary tree (SVG/PNG/JPG/BMP)')
852    parser.add_argument('-c', '--config', dest='config', default="", help='config file name ')
853
854    parser.add_argument('-W', '--width', default=600, type=int, dest='width', help='width of the output image (600 by default)')
855    parser.add_argument('-H', '--height', default=800, type=int, dest='height', help='height of the output image (800 by default)')
856    parser.add_argument('-m', '--multi', default=1, type=int, dest='multi', help='multisampling factor (applicable only for raster images)')
857
858    parser.add_argument('-t', '--time', default='GENERATIONAL', dest='time', help='values on vertical axis (BIRTHS/GENERATIONAL(d)/REAL); '
859                                                                      'BIRTHS: time measured as the number of births since the beginning; '
860                                                                      'GENERATIONAL: time measured as number of ancestors; '
861                                                                      'REAL: real time of the simulation')
862    parser.add_argument('-b', '--balance', default='DENSITY', dest='balance', help='method of placing nodes in the tree (RANDOM/MIN/DENSITY(d))')
863    parser.add_argument('-s', '--scale', default='SIMPLE', dest='scale', help='type of timescale added to the tree (NONE(d)/SIMPLE)')
864    parser.add_argument('-j', '--jitter', dest="jitter", action='store_true', help='draw horizontal positions of children from the normal distribution')
865    parser.add_argument('-p', '--skip', dest="skip", type=int, default=0, help='skip last P levels of the tree (0 by default)')
866    parser.add_argument('-x', '--max-nodes', type=int, default=0, dest='max_nodes', help='maximum number of nodes drawn (starting from the first one)')
867    parser.add_argument('--seed', type=int, dest='seed', help='seed for the random number generator (-1 for random)')
868
869    parser.set_defaults(draw_tree=True)
870    parser.set_defaults(draw_skeleton=False)
871    parser.set_defaults(draw_spine=False)
872
873    parser.set_defaults(seed=-1)
874
875    args = parser.parse_args()
876
877    TIME = args.time.upper()
878    BALANCE = args.balance.upper()
879    SCALE = args.scale.upper()
880    JITTER = args.jitter
881    if not TIME in ['BIRTHS', 'GENERATIONAL', 'REAL']\
882        or not BALANCE in ['RANDOM', 'MIN', 'DENSITY']\
883        or not SCALE in ['NONE', 'SIMPLE']:
884        print("Incorrect value of one of the parameters! (time or balance or scale).") #user has to figure out which parameter is wrong...
885        return
886
887    dir = args.input
888    seed = args.seed
889    if seed == -1:
890        seed = random.randint(0, 10000)
891    random.seed(seed)
892    print("randomseed:", seed)
893
894    tree = TreeData()
895    tree.load(dir, max_nodes=args.max_nodes)
896
897
898    designer = Designer(tree, jitter=JITTER, time=TIME, balance=BALANCE)
899    designer.calculate_measures()
900    designer.calculate_node_positions(ignore_last=args.skip)
901
902    if args.output.endswith(".svg"):
903        drawer = SvgDrawer(designer, args.config, w=args.width, h=args.height)
904    else:
905        drawer = PngDrawer(designer, args.config, w=args.width, h=args.height)
906    drawer.draw_design(args.output, args.input, multi=args.multi, scale=SCALE)
907
908
909main()
Note: See TracBrowser for help on using the repository browser.