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

Last change on this file since 685 was 685, checked in by Maciej Komosinski, 7 years ago

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