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

Last change on this file since 687 was 687, checked in by konrad, 7 years ago

Using "inherited' to decide horizontal distance from the parent

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