85 lines
3.2 KiB
Python
85 lines
3.2 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
class Node:
|
||
|
|
def __init__(self, idx, id='', length=None):
|
||
|
|
self.idx=idx; self.id=id; self.length=length; self.children=[]; self.parent=None
|
||
|
|
class Tree:
|
||
|
|
def __init__(self, root): self.root=root
|
||
|
|
def depthfirst(self):
|
||
|
|
out=[]
|
||
|
|
def rec(n):
|
||
|
|
out.append(n)
|
||
|
|
for c in n.children: rec(c)
|
||
|
|
rec(self.root); return out
|
||
|
|
def leaves(self): return [n for n in self.depthfirst() if not n.children]
|
||
|
|
def count_leaves(self):
|
||
|
|
d={}
|
||
|
|
def rec(n):
|
||
|
|
d[n]=1 if not n.children else sum(rec(c) for c in n.children); return d[n]
|
||
|
|
rec(self.root); return d
|
||
|
|
def leaf_sets(self):
|
||
|
|
d={}
|
||
|
|
def rec(n):
|
||
|
|
if not n.children: d[n]=frozenset([n.id]); return d[n]
|
||
|
|
s=frozenset().union(*(rec(c) for c in n.children)); d[n]=s; return s
|
||
|
|
rec(self.root); return d
|
||
|
|
def canonical_splits(self):
|
||
|
|
leafsets=self.leaf_sets(); all_leaves=leafsets[self.root]; splits={}
|
||
|
|
for n,s in leafsets.items():
|
||
|
|
if n is self.root or len(s)==0 or len(s)==len(all_leaves): continue
|
||
|
|
other=all_leaves-s
|
||
|
|
key=s if (len(s), sorted(s)) <= (len(other), sorted(other)) else other
|
||
|
|
splits[key]=n
|
||
|
|
return splits
|
||
|
|
def compare(self, other):
|
||
|
|
a=self.canonical_splits(); b=other.canonical_splits(); return {n:b[s] for s,n in a.items() if s in b}
|
||
|
|
def to_newick(self):
|
||
|
|
def fmt(n):
|
||
|
|
label=n.id or ''
|
||
|
|
text='('+','.join(fmt(c) for c in n.children)+')'+label if n.children else label
|
||
|
|
if n.length not in (None,''): text+=':'+str(n.length)
|
||
|
|
return text
|
||
|
|
return fmt(self.root)+';'
|
||
|
|
def split_newick_records(text):
|
||
|
|
records=[]; buf=[]; depth=0; quote=None
|
||
|
|
for ch in text:
|
||
|
|
buf.append(ch)
|
||
|
|
if quote:
|
||
|
|
if ch==quote: quote=None
|
||
|
|
elif ch in "'\"": quote=ch
|
||
|
|
elif ch=='(': depth+=1
|
||
|
|
elif ch==')': depth-=1
|
||
|
|
elif ch==';' and depth==0:
|
||
|
|
records.append(''.join(buf).strip()); buf=[]
|
||
|
|
return [r for r in records if r]
|
||
|
|
def parse_newick(text):
|
||
|
|
s=text.strip()
|
||
|
|
if s.endswith(';'): s=s[:-1]
|
||
|
|
i=0; idx=0
|
||
|
|
def read_label_len():
|
||
|
|
nonlocal i
|
||
|
|
label=[]
|
||
|
|
while i<len(s) and s[i] not in ',():': label.append(s[i]); i+=1
|
||
|
|
lab=''.join(label).strip(); length=None
|
||
|
|
if i<len(s) and s[i]==':':
|
||
|
|
i+=1; val=[]
|
||
|
|
while i<len(s) and s[i] not in ',()': val.append(s[i]); i+=1
|
||
|
|
length=''.join(val).strip()
|
||
|
|
return lab,length
|
||
|
|
def sub():
|
||
|
|
nonlocal i,idx
|
||
|
|
n=Node(idx); idx+=1
|
||
|
|
if i<len(s) and s[i]=='(':
|
||
|
|
i+=1
|
||
|
|
while True:
|
||
|
|
c=sub(); c.parent=n; n.children.append(c)
|
||
|
|
if i<len(s) and s[i]==',': i+=1; continue
|
||
|
|
if i<len(s) and s[i]==')': i+=1; break
|
||
|
|
break
|
||
|
|
n.id,n.length=read_label_len(); return n
|
||
|
|
return Tree(sub())
|
||
|
|
def read_trees(path):
|
||
|
|
with open(path, encoding='utf-8', errors='replace') as fh: return [parse_newick(x) for x in split_newick_records(fh.read())]
|
||
|
|
def trim_leaf_underscores(tree):
|
||
|
|
for n in tree.leaves():
|
||
|
|
if n.id is not None: n.id=n.id.strip('_')
|