compara-deep-learning/tree_data.py at master · EMBLEBIGSOC/compara-deep-learning · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
from ete3 import Tree
import numpy as np
import progressbar


def create_branch_length_padding(bl):
    maxlen = 29
    for x in bl:
        for i in range(len(x), maxlen):
            x.append(0)


def create_tree_data(treename, df):
    t = Tree(treename)
    branch_lengths_s = []
    branch_lengths_hs = []
    dist = []
    ns = []
    nhs = []
    for index, row in progressbar.progressbar(df.iterrows()):
        d = 0
        x = row["species"]
        y = row["homology_species"]
        bl = []
        c = 0
        mca = t.get_common_ancestor(x, y)
        node = t & x
        while node.up != mca:
            d += node.dist
            bl.append(node.dist)
            node = node.up
            c += 1
        ns.append(c)
        c = 0
        branch_lengths_s.append(bl)
        bl = []
        node = t & y
        while node.up != mca:
            d += node.dist
            bl.append(node.dist)
            node = node.up
            c += 1
        nhs.append(c)
        branch_lengths_hs.append(bl)
        dist.append(d)
    create_branch_length_padding(branch_lengths_s)
    create_branch_length_padding(branch_lengths_hs)
    return np.array(branch_lengths_s), np.array(
        branch_lengths_hs), np.array(dist), np.array(ns), np.array(nhs)