compara-deep-learning/tree_data.py

38 lines
952 B
Python
Raw Normal View History

from ete3 import Tree
import numpy as np
def create_tree_data(treename,df):
t=Tree(treename)
branch_lengths_s=[]
branch_lengths_hs=[]
dist=[]
ns=[]
nhs=[]
for index,row in 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)
return np.array(branch_lengths_s),np.array(branch_lengths_hs),np.array(dist),np.array(ns),np.array(nhs)