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

Last change on this file since 710 was 710, checked in by konrad, 3 years ago

Added multi-input-file support (eg. ./tree-genealogy.py -i file1.out file2.out file3.out -o tree.png)

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