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

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

Using matplotlib's color package to interpret colors. Color can be now specified in one of the following ways: 1. a letter from the set ‘rgbcmykw’; 2.a hex color string, like ‘#00FFFF’; 3.a standard name, like ‘aqua’; 4. a float, like ‘0.4’, indicating gray on a 0-1 scale

File size: 34.0 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
9from matplotlib import colors
10import numpy as np
11
12class LoadingError(Exception):
13    pass
14
15class Drawer:
16
17    def __init__(self, design, config_file, w=600, h=800, w_margin=10, h_margin=20):
18        self.design = design
19        self.width = w
20        self.height = h
21        self.w_margin = w_margin
22        self.h_margin = h_margin
23        self.w_no_margs = w - 2* w_margin
24        self.h_no_margs = h - 2* h_margin
25
26        self.color_converter = colors.ColorConverter()
27
28        self.settings = {
29            'colors_of_kinds': ['red', 'green', 'blue', 'magenta', 'yellow', 'cyan', 'orange', 'purple'],
30            'dots': {
31                'color': {
32                    'meaning': 'Lifespan',
33                    'start': 'red',
34                    'end': 'green',
35                    'bias': 1
36                    },
37                'size': {
38                    'meaning': 'EnergyEaten',
39                    'start': 1,
40                    'end': 6,
41                    'bias': 0.5
42                    },
43                'opacity': {
44                    'meaning': 'EnergyEaten',
45                    'start': 0.2,
46                    'end': 1,
47                    'bias': 1
48                    }
49            },
50            'lines': {
51                'color': {
52                    'meaning': 'adepth',
53                    'start': 'black',
54                    'end': 'red',
55                    'bias': 3
56                    },
57                'width': {
58                    'meaning': 'adepth',
59                    'start': 0.1,
60                    'end': 4,
61                    'bias': 3
62                    },
63                'opacity': {
64                    'meaning': 'adepth',
65                    'start': 0.1,
66                    'end': 0.8,
67                    'bias': 5
68                    }
69            }
70        }
71
72        def merge(source, destination):
73            for key, value in source.items():
74                if isinstance(value, dict):
75                    node = destination.setdefault(key, {})
76                    merge(value, node)
77                else:
78                    destination[key] = value
79
80            return destination
81
82        if config_file != "":
83            with open(config_file) as config:
84                c = json.load(config)
85            self.settings = merge(c, self.settings)
86            #print(json.dumps(self.settings, indent=4, sort_keys=True))
87
88    def draw_dots(self, file, min_width, max_width, max_height):
89        for i in range(len(self.design.positions)):
90            node = self.design.positions[i]
91            if 'x' not in node:
92                continue
93            dot_style = self.compute_dot_style(node=i)
94            self.add_dot(file, (self.w_margin+self.w_no_margs*(node['x']-min_width)/(max_width-min_width),
95                               self.h_margin+self.h_no_margs*node['y']/max_height), dot_style)
96
97    def draw_lines(self, file, min_width, max_width, max_height):
98        for parent in range(len(self.design.positions)):
99            par_pos = self.design.positions[parent]
100            if not 'x' in par_pos:
101                continue
102            for child in self.design.tree.children[parent]:
103                chi_pos = self.design.positions[child]
104                if 'x' not in chi_pos:
105                    continue
106                line_style = self.compute_line_style(parent, child)
107                self.add_line(file, (self.w_margin+self.w_no_margs*(par_pos['x']-min_width)/(max_width-min_width),
108                                  self.h_margin+self.h_no_margs*par_pos['y']/max_height),
109                                  (self.w_margin+self.w_no_margs*(chi_pos['x']-min_width)/(max_width-min_width),
110                                  self.h_margin+self.h_no_margs*chi_pos['y']/max_height), line_style)
111
112    def draw_scale(self, file, filename):
113        self.add_text(file, "Generated from " + filename.split("\\")[-1], (5, 5), "start")
114
115        start_text = ""
116        end_text = ""
117        if self.design.TIME == "BIRTHS":
118           start_text = "Birth #0"
119           end_text = "Birth #" + str(len(self.design.positions)-1)
120        if self.design.TIME == "REAL":
121           start_text = "Time " + str(min(self.design.tree.time))
122           end_text = "Time " + str(max(self.design.tree.time))
123        if self.design.TIME == "GENERATIONAL":
124           start_text = "Depth " + str(self.design.props['adepth_min'])
125           end_text = "Depth " + str(self.design.props['adepth_max'])
126
127        self.add_dashed_line(file, (self.width*0.7, self.h_margin), (self.width, self.h_margin))
128        self.add_text(file, start_text, (self.width, self.h_margin), "end")
129        self.add_dashed_line(file, (self.width*0.7, self.height-self.h_margin), (self.width, self.height-self.h_margin))
130        self.add_text(file, end_text, (self.width, self.height-self.h_margin), "end")
131
132    def compute_property(self, part, prop, node):
133        start = self.settings[part][prop]['start']
134        end = self.settings[part][prop]['end']
135        value = (self.design.props[self.settings[part][prop]['meaning']][node]
136                 if self.settings[part][prop]['meaning'] in self.design.props else 0 )
137        bias = self.settings[part][prop]['bias']
138        if prop == "color":
139            return self.compute_color(start, end, value, bias)
140        else:
141            return self.compute_value(start, end, value, bias)
142
143    def compute_color(self, start, end, value, bias=1):
144        if isinstance(value, str):
145            value = int(value)
146            r, g, b = self.color_converter.to_rgb(self.settings['colors_of_kinds'][value])
147        else:
148            start_color = self.color_converter.to_rgb(start)
149            end_color = self.color_converter.to_rgb(end)
150            value = 1 - (1-value)**bias
151            r = start_color[0]*(1-value)+end_color[0]*value
152            g = start_color[1]*(1-value)+end_color[1]*value
153            b = start_color[2]*(1-value)+end_color[2]*value
154        return (100*r, 100*g, 100*b)
155
156    def compute_value(self, start, end, value, bias=1):
157        value = 1 - (1-value)**bias
158        return start*(1-value) + end*value
159
160class PngDrawer(Drawer):
161
162    def scale_up(self):
163        self.width *= self.multi
164        self.height *= self.multi
165        self.w_margin *= self.multi
166        self.h_margin *= self.multi
167        self.h_no_margs *= self.multi
168        self.w_no_margs *= self.multi
169
170    def scale_down(self):
171        self.width /= self.multi
172        self.height /= self.multi
173        self.w_margin /= self.multi
174        self.h_margin /= self.multi
175        self.h_no_margs /= self.multi
176        self.w_no_margs /= self.multi
177
178    def draw_design(self, filename, input_filename, multi=1, scale="SIMPLE"):
179        print("Drawing...")
180
181        self.multi=multi
182        self.scale_up()
183
184        back = Image.new('RGBA', (self.width, self.height), (255,255,255,0))
185
186        min_width = min([x['x'] for x in self.design.positions if 'x' in x])
187        max_width = max([x['x'] for x in self.design.positions if 'x' in x])
188        max_height = max([x['y'] for x in self.design.positions if 'y' in x])
189
190        self.draw_lines(back, min_width, max_width, max_height)
191        self.draw_dots(back, min_width, max_width, max_height)
192
193        if scale == "SIMPLE":
194            self.draw_scale(back, input_filename)
195
196        #back.show()
197        self.scale_down()
198
199        back.thumbnail((self.width, self.height), Image.ANTIALIAS)
200
201        back.save(filename)
202
203    def add_dot(self, file, pos, style):
204        x, y = int(pos[0]), int(pos[1])
205        r = style['r']*self.multi
206        offset = (int(x - r), int(y - r))
207        size = (2*int(r), 2*int(r))
208
209        c = style['color']
210
211        img = Image.new('RGBA', size)
212        ImageDraw.Draw(img).ellipse((1, 1, size[0]-1, size[1]-1),
213                                    (int(2.55*c[0]), int(2.55*c[1]), int(2.55*c[2]), int(255*style['opacity'])))
214        file.paste(img, offset, mask=img)
215
216    def add_line(self, file, from_pos, to_pos, style):
217        fx, fy, tx, ty = int(from_pos[0]), int(from_pos[1]), int(to_pos[0]), int(to_pos[1])
218        w = int(style['width'])*self.multi
219
220        offset = (min(fx-w, tx-w), min(fy-w, ty-w))
221        size = (abs(fx-tx)+2*w, abs(fy-ty)+2*w)
222
223        c = style['color']
224
225        img = Image.new('RGBA', size)
226        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),
227                                  (int(2.55*c[0]), int(2.55*c[1]), int(2.55*c[2]), int(255*style['opacity'])), w)
228        file.paste(img, offset, mask=img)
229
230    def add_dashed_line(self, file, from_pos, to_pos):
231        style = {'color': (0,0,0), 'width': 1, 'opacity': 1}
232        sublines = 50
233        # TODO could be faster: compute delta and only add delta each time (but currently we do not use it often)
234        normdiv = 2*sublines-1
235        for i in range(sublines):
236            from_pos_sub = (self.compute_value(from_pos[0], to_pos[0], 2*i/normdiv, 1),
237                            self.compute_value(from_pos[1], to_pos[1], 2*i/normdiv, 1))
238            to_pos_sub = (self.compute_value(from_pos[0], to_pos[0], (2*i+1)/normdiv, 1),
239                          self.compute_value(from_pos[1], to_pos[1], (2*i+1)/normdiv, 1))
240            self.add_line(file, from_pos_sub, to_pos_sub, style)
241
242    def add_text(self, file, text, pos, anchor, style=''):
243        font = ImageFont.truetype("Vera.ttf", 16*self.multi)
244
245        img = Image.new('RGBA', (self.width, self.height))
246        draw = ImageDraw.Draw(img)
247        txtsize = draw.textsize(text, font=font)
248        pos = pos if anchor == "start" else (pos[0]-txtsize[0], pos[1])
249        draw.text(pos, text, (0,0,0), font=font)
250        file.paste(img, (0,0), mask=img)
251
252    def compute_line_style(self, parent, child):
253        return {'color': self.compute_property('lines', 'color', child),
254                'width': self.compute_property('lines', 'width', child),
255                'opacity': self.compute_property('lines', 'opacity', child)}
256
257    def compute_dot_style(self, node):
258        return {'color': self.compute_property('dots', 'color', node),
259                'r': self.compute_property('dots', 'size', node),
260                'opacity': self.compute_property('dots', 'opacity', node)}
261
262class SvgDrawer(Drawer):
263    def draw_design(self, filename, input_filename, multi=1, scale="SIMPLE"):
264        print("Drawing...")
265        file = open(filename, "w")
266
267        min_width = min([x['x'] for x in self.design.positions if 'x' in x])
268        max_width = max([x['x'] for x in self.design.positions if 'x' in x])
269        max_height = max([x['y'] for x in self.design.positions if 'y' in x])
270
271        file.write('<svg xmlns:svg="http://www.w3.org/2000/svg" xmlns="http://www.w3.org/2000/svg" '
272                   'xmlns:xlink="http://www.w3.org/1999/xlink" version="1.0" '
273                   'width="' + str(self.width) + '" height="' + str(self.height) + '">')
274
275        self.draw_lines(file, min_width, max_width, max_height)
276        self.draw_dots(file, min_width, max_width, max_height)
277
278        if scale == "SIMPLE":
279            self.draw_scale(file, input_filename)
280
281        file.write("</svg>")
282        file.close()
283
284    def add_text(self, file, text, pos, anchor, style=''):
285        style = (style if style != '' else 'style="font-family: Arial; font-size: 12; fill: #000000;"')
286        # assuming font size 12, it should be taken from the style string!
287        file.write('<text ' + style + ' text-anchor="' + anchor + '" x="' + str(pos[0]) + '" y="' + str(pos[1]+12) + '" >' + text + '</text>')
288
289    def add_dot(self, file, pos, style):
290        file.write('<circle ' + style + ' cx="' + str(pos[0]) + '" cy="' + str(pos[1]) + '" />')
291
292    def add_line(self, file, from_pos, to_pos, style):
293        file.write('<line ' + style + ' x1="' + str(from_pos[0]) + '" x2="' + str(to_pos[0]) +
294                       '" y1="' + str(from_pos[1]) + '" y2="' + str(to_pos[1]) + '"  fill="none"/>')
295
296    def add_dashed_line(self, file, from_pos, to_pos):
297        style = 'stroke="black" stroke-width="0.5" stroke-opacity="1" stroke-dasharray="5, 5"'
298        self.add_line(file, from_pos, to_pos, style)
299
300    def compute_line_style(self, parent, child):
301        return self.compute_stroke_color('lines', child) + ' ' \
302               + self.compute_stroke_width('lines', child) + ' ' \
303               + self.compute_stroke_opacity(child)
304
305    def compute_dot_style(self, node):
306        return self.compute_dot_size(node) + ' ' \
307               + self.compute_fill_opacity(node) + ' ' \
308               + self.compute_dot_fill(node)
309
310    def compute_stroke_color(self, part, node):
311        color = self.compute_property(part, 'color', node)
312        return 'stroke="rgb(' + str(color[0]) + '%,' + str(color[1]) + '%,' + str(color[2]) + '%)"'
313
314    def compute_stroke_width(self, part, node):
315        return 'stroke-width="' + str(self.compute_property(part, 'width', node)) + '"'
316
317    def compute_stroke_opacity(self, node):
318        return 'stroke-opacity="' + str(self.compute_property('lines', 'opacity', node)) + '"'
319
320    def compute_fill_opacity(self, node):
321        return 'fill-opacity="' + str(self.compute_property('dots', 'opacity', node)) + '"'
322
323    def compute_dot_size(self, node):
324        return 'r="' + str(self.compute_property('dots', 'size', node)) + '"'
325
326    def compute_dot_fill(self, node):
327        color = self.compute_property('dots', 'color', node)
328        return 'fill="rgb(' + str(color[0]) + '%,' + str(color[1]) + '%,' + str(color[2]) + '%)"'
329
330class Designer:
331
332    def __init__(self, tree, jitter=False, time="GENERATIONAL", balance="DENSITY"):
333        self.props = {}
334
335        self.tree = tree
336
337        self.TIME = time
338        self.JITTER = jitter
339
340        if balance == "RANDOM":
341            self.xmin_crowd = self.xmin_crowd_random
342        elif balance == "MIN":
343            self.xmin_crowd = self.xmin_crowd_min
344        elif balance == "DENSITY":
345            self.xmin_crowd = self.xmin_crowd_density
346        else:
347            raise ValueError("Error, the value of BALANCE does not match any expected value.")
348
349    def calculate_measures(self):
350        print("Calculating measures...")
351        self.compute_depth()
352        self.compute_adepth()
353        self.compute_children()
354        self.compute_kind()
355        self.compute_time()
356        self.compute_progress()
357        self.compute_custom()
358
359    def xmin_crowd_random(self, x1, x2, y):
360        return (x1 if random.randrange(2) == 0 else x2)
361
362    def xmin_crowd_min(self, x1, x2, y):
363        x1_closest = 999999
364        x2_closest = 999999
365        miny = y-3
366        maxy = y+3
367        i = bisect.bisect_left(self.y_sorted, miny)
368        while True:
369            if len(self.positions_sorted) <= i or self.positions_sorted[i]['y'] > maxy:
370                break
371            pos = self.positions_sorted[i]
372
373            x1_closest = min(x1_closest, abs(x1-pos['x']))
374            x2_closest = min(x2_closest, abs(x2-pos['x']))
375
376            i += 1
377        return (x1 if x1_closest > x2_closest else x2)
378
379    def xmin_crowd_density(self, x1, x2, y):
380        # TODO experimental - requires further work to make it less 'jumpy' and more predictable
381        CONST_LOCAL_AREA_RADIUS = 5
382        CONST_GLOBAL_AREA_RADIUS = 10
383        CONST_WINDOW_SIZE = 20000 #TODO should depend on the maxY ?
384        x1_dist_loc = 0
385        x2_dist_loc = 0
386        count_loc = 1
387        x1_dist_glob = 0
388        x2_dist_glob = 0
389        count_glob = 1
390        miny = y-CONST_WINDOW_SIZE
391        maxy = y+CONST_WINDOW_SIZE
392        i_left = bisect.bisect_left(self.y_sorted, miny)
393        i_right = bisect.bisect_right(self.y_sorted, maxy)
394        #TODO test: maxy=y should give the same results, right?
395
396        def include_pos(pos):
397            nonlocal x1_dist_loc, x2_dist_loc, x1_dist_glob, x2_dist_glob, count_loc, count_glob
398
399            dysq = (pos['y']-y)**2 + 1 #+1 so 1/dysq is at most 1
400            dx1 = math.fabs(pos['x']-x1)
401            dx2 = math.fabs(pos['x']-x2)
402
403            d = math.fabs(pos['x'] - (x1+x2)/2)
404
405            if d < CONST_LOCAL_AREA_RADIUS:
406                x1_dist_loc += math.sqrt(dx1/dysq + dx1**2)
407                x2_dist_loc += math.sqrt(dx2/dysq + dx2**2)
408                count_loc += 1
409            elif d > CONST_GLOBAL_AREA_RADIUS:
410                x1_dist_glob += math.sqrt(dx1/dysq + dx1**2)
411                x2_dist_glob += math.sqrt(dx2/dysq + dx2**2)
412                count_glob += 1
413
414        # optimized to draw from all the nodes, if less than 10 nodes in the range
415        if len(self.positions_sorted) > i_left:
416            if i_right - i_left < 10:
417                for j in range(i_left, i_right):
418                    include_pos(self.positions_sorted[j])
419            else:
420                for j in range(10):
421                    pos = self.positions_sorted[random.randrange(i_left, i_right)]
422                    include_pos(pos)
423
424        return (x1 if (x1_dist_loc-x2_dist_loc)/count_loc-(x1_dist_glob-x2_dist_glob)/count_glob > 0  else x2)
425        #return (x1 if x1_dist +random.gauss(0, 0.00001) > x2_dist +random.gauss(0, 0.00001)  else x2)
426        #print(x1_dist, x2_dist)
427        #x1_dist = x1_dist**2
428        #x2_dist = x2_dist**2
429        #return x1 if x1_dist+x2_dist==0 else (x1*x1_dist + x2*x2_dist) / (x1_dist+x2_dist) + random.gauss(0, 0.01)
430        #return (x1 if random.randint(0, int(x1_dist+x2_dist)) < x1_dist else x2)
431
432    def calculate_node_positions(self, ignore_last=0):
433        print("Calculating positions...")
434
435        def add_node(node):
436            index = bisect.bisect_left(self.y_sorted, node['y'])
437            self.y_sorted.insert(index, node['y'])
438            self.positions_sorted.insert(index, node)
439            self.positions[node['id']] = node
440
441        self.positions_sorted = [{'x':0, 'y':0, 'id':0}]
442        self.y_sorted = [0]
443        self.positions = [{} for x in range(len(self.tree.parents))]
444        self.positions[0] = {'x':0, 'y':0, 'id':0}
445
446        # order by maximum depth of the parent guarantees that co child is evaluated before its parent
447        visiting_order = [i for i in range(0, len(self.tree.parents))]
448        visiting_order = sorted(visiting_order, key=lambda q:
449                            0 if q == 0 else max([self.props["depth"][d] for d in self.tree.parents[q]]))
450
451        start_time = timelib.time()
452
453        # for each child of the current node
454        for node_counter,child in enumerate(visiting_order, start=1):
455            # debug info - elapsed time
456            if node_counter % 100000 == 0:
457               print("%d%%\t%d\t%g" % (node_counter*100/len(self.tree.parents), node_counter, 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, similarity = [(par, v) for par, v in self.tree.parents[child].items()][0]
476
477                    if self.JITTER:
478                        dissimilarity = (1-similarity) + random.gauss(0, 0.01) + 0.001
479                    else:
480                        dissimilarity = (1-similarity) + 0.001
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("nodes: %d" % 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 = (creature["Inherited"][i] if 'Inherited' in creature else 1)
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("randomseed:", 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.