def _treeFromTimes(times, order = None) : """ Build a biopython tree from intra-coalescent intervals in times. Without an order, the joining is random. With order, use the given ordering supplied as a sequence of pairs, each pair is two indices picked to join. """ if order is not None : order = list(order) assert len(times) == len(order) n = len(times) + 1 tb = TreeBuilder() taxa = [[tb.createLeaf("ID%d" % k), 0.0] for k in range(n)] height = 0.0 for t in reversed(times) : i,j = order.pop(0) if order else \ random.sample(range(len(taxa)), 2) l1,l2 = taxa[i],taxa[j] height += t nd = tb.mergeNodes([[x[0], height - x[1]] for x in (l1,l2)]) taxa.pop(max(i,j)) taxa[min(i,j)] = [nd, height] return tb.finalize(taxa[0][0])
def _build(nodes, weight=1.0, rooted=True, name='') : # name/support ignored tb = TreeBuilder(weight=weight, rooted = rooted, name=name) t = [None]*len(nodes) for k, x in enumerate(nodes): if x[2] is None or len(x[2]) == 0: t[k] = tb.createLeaf(x[0]) else : t[k] = tb.mergeNodes([ [t[l], nodes[l][1]] for l in x[2]]) if x[3] is not None: t[k].data.attributes = dict(x[3]) return tb.finalize(t[-1])
def sampleCoalescentTree(demog, labels) : """ Build a tree under the coalescent model using given demographic.""" if isinstance(labels, int) : labels = ["tip%03d" % k for k in range(int(labels))] else : assert all([isinstance(s, str) for s in labels]) n = len(labels) times= [0,]*n lineageInfo = len(labels) # times = reduce(operator.add, [[t,]*n for n,t in lineageInfo]) #assert len(labels) == len(times) at = getArrivalTimes(demog, lineageInfo) tb = TreeBuilder() nodes = [tb.createLeaf(l) for l in labels] strees = zip(nodes, times) for t in at : # find 2 labels with times < t and coalesce available = [k for k in range(len(strees)) if strees[k][1] < t] assert len(available) > 1,(available,t,at,strees) i,j = random.sample(available, 2) l1,l2 = strees[i],strees[j] n = tb.mergeNodes([[x[0], t - x[1]] for x in (l1,l2)]) strees.pop(max(i,j)) strees.pop(min(i,j)) strees.insert(0, (n, t)) assert len(strees) == 1 return tb.finalize(strees[0][0])