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

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

Minor optimization
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            # debug info - elapsed time
456            if node_counter%1000 == 0:
457               print(str(node_counter) + " " + str(timelib.time()-start_time))
458               start_time = timelib.time()
459
460            # using normalized adepth
461            if self.props['adepth'][child] >= ignore_last/self.props['adepth_max']:
462
463                ypos = 0
464                if self.TIME == "BIRTHS":
465                    ypos = child
466                elif self.TIME == "GENERATIONAL":
467                    # one more than its parent (what if more than one parent?)
468                    ypos = max([self.positions[par]['y'] for par, v in self.tree.parents[child].items()])+1 \
469                        if self.tree.parents[child] else 0
470                elif self.TIME == "REAL":
471                    ypos = self.tree.time[child]
472
473                if len(self.tree.parents[child]) == 1:
474                # if current_node is the only parent
475                    parent = [par for par, v in self.tree.parents[child].items()][0]
476
477                    if self.JITTER:
478                        dissimilarity = random.gauss(0, 0.5)
479                    else:
480                        dissimilarity = 1
481                    add_node({'id':child, 'y':ypos, 'x':
482                             self.xmin_crowd(self.positions[parent]['x']-dissimilarity,
483                              self.positions[parent]['x']+dissimilarity, ypos)})
484                else:
485                    # position weighted by the degree of inheritence from each parent
486                    total_inheretance = sum([v for k, v in self.tree.parents[child].items()])
487                    xpos = sum([self.positions[k]['x']*v/total_inheretance
488                               for k, v in self.tree.parents[child].items()])
489                    if self.JITTER:
490                        add_node({'id':child, 'y':ypos, 'x':xpos + random.gauss(0, 0.1)})
491                    else:
492                        add_node({'id':child, 'y':ypos, 'x':xpos})
493
494
495    def compute_custom(self):
496        for prop in self.tree.props:
497            self.props[prop] = [None for x in range(len(self.tree.children))]
498
499            for i in range(len(self.props[prop])):
500                self.props[prop][i] = self.tree.props[prop][i]
501
502            self.normalize_prop(prop)
503
504    def compute_time(self):
505        # simple rewrite from the tree
506        self.props["time"] = [0 for x in range(len(self.tree.children))]
507
508        for i in range(len(self.props['time'])):
509            self.props['time'][i] = self.tree.time[i]
510
511        self.normalize_prop('time')
512
513    def compute_kind(self):
514        # simple rewrite from the tree
515        self.props["kind"] = [0 for x in range(len(self.tree.children))]
516
517        for i in range (len(self.props['kind'])):
518            self.props['kind'][i] = str(self.tree.kind[i])
519
520    def compute_depth(self):
521        self.props["depth"] = [999999999 for x in range(len(self.tree.children))]
522        visited = [0 for x in range(len(self.tree.children))]
523
524        nodes_to_visit = [0]
525        visited[0] = 1
526        self.props["depth"][0] = 0
527        while True:
528            current_node = nodes_to_visit[0]
529
530            for child in self.tree.children[current_node]:
531                if visited[child] == 0:
532                    visited[child] = 1
533                    nodes_to_visit.append(child)
534                    self.props["depth"][child] = self.props["depth"][current_node]+1
535            nodes_to_visit = nodes_to_visit[1:]
536            if len(nodes_to_visit) == 0:
537                break
538
539        self.normalize_prop('depth')
540
541    def compute_adepth(self):
542        self.props["adepth"] = [0 for x in range(len(self.tree.children))]
543
544        # order by maximum depth of the parent guarantees that co child is evaluated before its parent
545        visiting_order = [i for i in range(0, len(self.tree.parents))]
546        visiting_order = sorted(visiting_order, key=lambda q:
547                            0 if q == 0 else max([self.props["depth"][d] for d in self.tree.parents[q]]))[::-1]
548
549        for node in visiting_order:
550            children = self.tree.children[node]
551            if len(children) != 0:
552                # 0 by default
553                self.props["adepth"][node] = max([self.props["adepth"][child] for child in children])+1
554        self.normalize_prop('adepth')
555
556    def compute_children(self):
557        self.props["children"] = [0 for x in range(len(self.tree.children))]
558        for i in range (len(self.props['children'])):
559            self.props['children'][i] = len(self.tree.children[i])
560
561        self.normalize_prop('children')
562
563    def compute_progress(self):
564        self.props["progress"] = [0 for x in range(len(self.tree.children))]
565        for i in range(len(self.props['children'])):
566            times = sorted([self.props["time"][self.tree.children[i][j]]*100000 for j in range(len(self.tree.children[i]))])
567            if len(times) > 4:
568                times = [times[i+1] - times[i] for i in range(len(times)-1)]
569                #print(times)
570                slope, intercept, r_value, p_value, std_err = stats.linregress(range(len(times)), times)
571                self.props['progress'][i] = slope if not np.isnan(slope) and not np.isinf(slope) else 0
572
573        for i in range(0, 5):
574            self.props['progress'][self.props['progress'].index(min(self.props['progress']))] = 0
575            self.props['progress'][self.props['progress'].index(max(self.props['progress']))] = 0
576
577        mini = min(self.props['progress'])
578        maxi = max(self.props['progress'])
579        for k in range(len(self.props['progress'])):
580            if self.props['progress'][k] == 0:
581                self.props['progress'][k] = mini
582
583        #for k in range(len(self.props['progress'])):
584        #        self.props['progress'][k] = 1-self.props['progress'][k]
585
586        self.normalize_prop('progress')
587
588    def normalize_prop(self, prop):
589        noneless = [v for v in self.props[prop] if (type(v)!=str and type(v)!=list)]
590        if len(noneless) > 0:
591            max_val = max(noneless)
592            min_val = min(noneless)
593            print(prop, max_val, min_val)
594            self.props[prop +'_max'] = max_val
595            self.props[prop +'_min'] = min_val
596            for i in range(len(self.props[prop])):
597                if self.props[prop][i] is not None:
598                    qqq = self.props[prop][i]
599                    self.props[prop][i] = 0 if max_val == min_val else (self.props[prop][i] - min_val) / (max_val - min_val)
600
601class TreeData:
602    simple_data = None
603
604    children = []
605    parents = []
606    time = []
607    kind = []
608
609    def __init__(self): #, simple_data=False):
610        #self.simple_data = simple_data
611        pass
612
613    def load(self, filename, max_nodes=0):
614        print("Loading...")
615
616        CLI_PREFIX = "Script.Message:"
617        default_props = ["Time", "FromIDs", "ID", "Operation", "Inherited"]
618
619        self.ids = {}
620        def get_id(id, createOnError = True):
621            if createOnError:
622                if id not in self.ids:
623                    self.ids[id] = len(self.ids)
624            else:
625                if id not in self.ids:
626                    return None
627            return self.ids[id]
628
629        file = open(filename)
630
631        # counting the number of expected nodes
632        nodes = 0
633        for line in file:
634            line_arr = line.split(' ', 1)
635            if len(line_arr) == 2:
636                if line_arr[0] == CLI_PREFIX:
637                    line_arr = line_arr[1].split(' ', 1)
638                if line_arr[0] == "[OFFSPRING]":
639                    nodes += 1
640
641        nodes = min(nodes, max_nodes if max_nodes != 0 else nodes)+1
642        self.parents = [{} for x in range(nodes)]
643        self.children = [[] for x in range(nodes)]
644        self.time = [0] * nodes
645        self.kind = [0] * nodes
646        self.life_lenght = [0] * nodes
647        self.props = {}
648
649        print(len(self.parents))
650
651        file.seek(0)
652        loaded_so_far = 0
653        lasttime = timelib.time()
654        for line in file:
655            line_arr = line.split(' ', 1)
656            if len(line_arr) == 2:
657                if line_arr[0] == CLI_PREFIX:
658                    line_arr = line_arr[1].split(' ', 1)
659                if line_arr[0] == "[OFFSPRING]":
660                    try:
661                        creature = json.loads(line_arr[1])
662                    except ValueError:
663                        print("Json format error - the line cannot be read. Breaking the loading loop.")
664                        # fixing arrays by removing the last element
665                        # ! assuming that only the last line is broken !
666                        self.parents.pop()
667                        self.children.pop()
668                        self.time.pop()
669                        self.kind.pop()
670                        self.life_lenght.pop()
671                        nodes -= 1
672                        break
673
674                    if "FromIDs" in creature:
675
676                        # make sure that ID's of parents are lower than that of their children
677                        for i in range(0, len(creature["FromIDs"])):
678                            if creature["FromIDs"][i] not in self.ids:
679                                get_id("virtual_parent")
680
681                        creature_id = get_id(creature["ID"])
682
683                        # debug
684                        if loaded_so_far%1000 == 0:
685                            #print(". " + str(creature_id) + " " + str(timelib.time() - lasttime))
686                            lasttime = timelib.time()
687
688                        # we assign to each parent its contribution to the genotype of the child
689                        for i in range(0, len(creature["FromIDs"])):
690                            if creature["FromIDs"][i] in self.ids:
691                                parent_id = get_id(creature["FromIDs"][i])
692                            else:
693                                parent_id = get_id("virtual_parent")
694                            inherited = 1 #(creature["Inherited"][i] if 'Inherited' in creature else 1) #ONLY FOR NOW
695                            self.parents[creature_id][parent_id] = inherited
696
697                        if "Time" in creature:
698                            self.time[creature_id] = creature["Time"]
699
700                        if "Kind" in creature:
701                            self.kind[creature_id] = creature["Kind"]
702
703                        for prop in creature:
704                            if prop not in default_props:
705                                if prop not in self.props:
706                                    self.props[prop] = [0 for i in range(nodes)]
707                                self.props[prop][creature_id] = creature[prop]
708
709                        loaded_so_far += 1
710                    else:
711                        raise LoadingError("[OFFSPRING] misses the 'FromIDs' field!")
712                if line_arr[0] == "[DIED]":
713                    creature = json.loads(line_arr[1])
714                    creature_id = get_id(creature["ID"], False)
715                    if creature_id is not None:
716                        for prop in creature:
717                            if prop not in default_props:
718                                if prop not in self.props:
719                                    self.props[prop] = [0 for i in range(nodes)]
720                                self.props[prop][creature_id] = creature[prop]
721
722
723            if loaded_so_far >= max_nodes and max_nodes != 0:
724                break
725
726        for k in range(len(self.parents)):
727            v = self.parents[k]
728            for val in self.parents[k]:
729                self.children[val].append(k)
730
731depth = {}
732kind = {}
733
734def main():
735
736    parser = argparse.ArgumentParser(description='Draws a genealogical tree (generates a SVG file) based on parent-child relationship '
737                                                 'information from a text file. Supports files generated by Framsticks experiments.')
738    parser.add_argument('-i', '--in', dest='input', required=True, help='input file name with stuctured evolutionary data')
739    parser.add_argument('-o', '--out', dest='output', required=True, help='output file name for the evolutionary tree (SVG/PNG/JPG/BMP)')
740    parser.add_argument('-c', '--config', dest='config', default="", help='config file name ')
741
742    parser.add_argument('-W', '--width', default=600, type=int, dest='width', help='width of the output image (600 by default)')
743    parser.add_argument('-H', '--height', default=800, type=int, dest='height', help='height of the output image (800 by default)')
744    parser.add_argument('-m', '--multi', default=1, type=int, dest='multi', help='multisampling factor (applicable only for raster images)')
745
746    parser.add_argument('-t', '--time', default='GENERATIONAL', dest='time', help='values on vertical axis (BIRTHS/GENERATIONAL(d)/REAL); '
747                                                                      'BIRTHS: time measured as the number of births since the beginning; '
748                                                                      'GENERATIONAL: time measured as number of ancestors; '
749                                                                      'REAL: real time of the simulation')
750    parser.add_argument('-b', '--balance', default='DENSITY', dest='balance', help='method of placing nodes in the tree (RANDOM/MIN/DENSITY(d))')
751    parser.add_argument('-s', '--scale', default='SIMPLE', dest='scale', help='type of timescale added to the tree (NONE(d)/SIMPLE)')
752    parser.add_argument('-j', '--jitter', dest="jitter", action='store_true', help='draw horizontal positions of children from the normal distribution')
753    parser.add_argument('-p', '--skip', dest="skip", type=int, default=0, help='skip last P levels of the tree (0 by default)')
754    parser.add_argument('-x', '--max-nodes', type=int, default=0, dest='max_nodes', help='maximum number of nodes drawn (starting from the first one)')
755    parser.add_argument('--seed', type=int, dest='seed', help='seed for the random number generator (-1 for random)')
756
757    parser.set_defaults(draw_tree=True)
758    parser.set_defaults(draw_skeleton=False)
759    parser.set_defaults(draw_spine=False)
760
761    parser.set_defaults(seed=-1)
762
763    args = parser.parse_args()
764
765    TIME = args.time.upper()
766    BALANCE = args.balance.upper()
767    SCALE = args.scale.upper()
768    JITTER = args.jitter
769    if not TIME in ['BIRTHS', 'GENERATIONAL', 'REAL']\
770        or not BALANCE in ['RANDOM', 'MIN', 'DENSITY']\
771        or not SCALE in ['NONE', 'SIMPLE']:
772        print("Incorrect value of one of the parameters! (time or balance or scale).") #user has to figure out which parameter is wrong...
773        return
774
775    dir = args.input
776    seed = args.seed
777    if seed == -1:
778        seed = random.randint(0, 10000)
779    random.seed(seed)
780    print("seed:", seed)
781
782    tree = TreeData()
783    tree.load(dir, max_nodes=args.max_nodes)
784
785
786    designer = Designer(tree, jitter=JITTER, time=TIME, balance=BALANCE)
787    designer.calculate_measures()
788    designer.calculate_node_positions(ignore_last=args.skip)
789
790    if args.output.endswith(".svg"):
791        drawer = SvgDrawer(designer, args.config, w=args.width, h=args.height)
792    else:
793        drawer = PngDrawer(designer, args.config, w=args.width, h=args.height)
794    drawer.draw_design(args.output, args.input, multi=args.multi, scale=SCALE)
795
796
797main()
Note: See TracBrowser for help on using the repository browser.