def drawTreeOnly(tree, plotLinesAttr, stree = None,
allTipsZero = True, xroot = None, keepPositions = False,
coalAttr = None, txt = False) :
assert not stree or allTipsZero
# A mapping from node id to (positive) node height (with !allTipsZero, tips
# may have > 0 height, at least one tip has 0 height).
nh = nodeHeights(tree, allTipsZero = allTipsZero)
if stree is not None :
snh = nodeHeights(stree)
rl = _embSetLeftRight(tree, tree.root, nh, stree, snh)
# position root in the middle
if xroot is None :
xroot = stree.node(stree.root).data.plotAux.center.center()
_embDrawTreeOnly(tree, tree.root, xroot, nh, plotLinesAttr,
keepPositions, coalAttr, txt)
else :
# set auxiliary info per node
rl = _setLeftRight(tree, tree.root, nh)
# position root in the middle
if xroot is None :
xroot = (rl.right + rl.left)/2
_drawTreeOnly(tree, tree.root, xroot, nh, plotLinesAttr, keepPositions, txt)
for i in tree.all_ids() :
del tree.node(i).data.cladeInfo
def _adjustTips(tree, refTree) :
nhs = nodeHeights(refTree, allTipsZero = False)
nhsd = dict([(refTree.node(i).data.taxon,nhs[i])
for i in nhs if not refTree.node(i).succ])
for ni in tree.get_terminals():
node = tree.node(ni)
h = nhsd[node.data.taxon]
if h > 0 :
node.data.branchlength -= h
if node.data.branchlength < 0 :
#node branch too short for adjustment.
#chop branches on the path to root.
#import pdb; pdb.set_trace()
while node.id != tree.root :
pr = tree.node(node.prev)
b = node.data.branchlength
node.data.branchlength = 0
pr.data.branchlength += b
# Other descendants should keep their height
for i in pr.succ :
if i != node.id :
tree.node(i).data.branchlength -= b
if pr.data.branchlength >= 0 :
break
node = pr
def taxaPartitionsSummaryTree(trees, summaryType = "median", atz = False) :
if summaryType not in ["mean", "median", "both"] :
raise ValueError("summaryType should be one of mean, median, both")
torder = _taxOrder(trees)
hs = [list() for k in torder[:-1]]
nhsd = None
if not atz:
tree = trees[0]
nhs = nodeHeights(tree, allTipsZero = False)
nhsd = [(tree.node(i).data.taxon, nhs[i])
for i in nhs if not tree.node(i).succ]
if any([h > 1e-13 for t,h in nhsd]) :
nhsd = dict(nhsd)
else :
nhsd = None
txp = dict(zip(torder,itertools.count()))
for t in trees :
_getCladeEstimates(t, txp, hs, nhsd)
if summaryType in ["median","both"] :
md = mau.mau2Tree((torder, [median(x) if len(x) else 0 for x in hs],
[(0,False) for x in hs]))
if not atz :
_adjustTips(md, trees[0])
if summaryType in ["mean","both"] :
mn = mau.mau2Tree((torder, [mean(x) if len(x) else 0 for x in hs],
[(0,False) for x in hs]))
if not atz :
_adjustTips(mn, trees[0])
return (mn,md) if summaryType=="both" else mn if summaryType == "mean" else md
def drawTree(tree, nid = None, cladesDict = None, positioning = None,
fill = None, generalPlotAttributes = None, splitPoints = False,
keepAux = False) :
if fill is None and generalPlotAttributes is None :
generalPlotAttributes = dict()
if fill is not None and 'ec' not in fill :
fill['ec'] = 'none'
if generalPlotAttributes is not None and 'color' not in generalPlotAttributes :
generalPlotAttributes['color'] = 'blue'
if nid is None :
nid = tree.root
if positioning is not None :
h = _drawTree(tree, nid, cladesDict, positioning,
fill, generalPlotAttributes, splitPoints, keepAux)
else :
nh = nodeHeights(tree, allTipsZero = True)
rl = _setLeftRight(tree, tree.root, nh)
xroot = (rl.right + rl.left)/2
h = _drawTreeTopDown(tree, nid, xroot,
fill, nh, generalPlotAttributes, splitPoints, keepAux)
return h
def simulateGeneTree(sTree) :
""" Simulate one gene tree under species (spp) tree C{sTree}.
Return a pair of tree and root height.
"""
nh = nodeHeights(sTree)
simTree = TreeBuilder()
t,rootHeight = _simulateGeneTreeForNode(sTree, sTree.root, simTree, nh)[0]
t = simTree.finalize(t)
return (t,rootHeight)
def __init__(self, tree) :
self.tree = tree
self.nhts = nodeHeights(tree)
assert all([isinstance(tree.node(n).data.demographic,
(demographic.StepFunctionPopulation,
demographic.ConstantPopulation,
demographic.LinearPiecewisePopulation))
for n in tree.all_ids()])
nlp = sum([isinstance(tree.node(n).data.demographic,demographic.LinearPiecewisePopulation)
for n in tree.all_ids()])
assert nlp == 0 or nlp == len(tree.all_ids())
self.isStepDemographic = nlp == 0
def _collectCladeTaxaNonATZ(tree, taxa, partitions) :
nhs = nodeHeights(tree, allTipsZero = False)
for n in getPostOrder(tree):
data = n.data
if not n.succ:
data.clade = [taxa.index(data.taxon),]
else :
p = []
for s in n.succ :
d = tree.node(s).data
p.extend(d.clade)
del d.clade
data.clade = p
partitions[frozenset(data.clade)] = (n.id, nhs[n.id])
return nhs[tree.root]
def setIMrates(stree, pim = 0.05, spr = 0.15, bdGrowth = None,
restrictToTree = True) :
# growth rate = lambda - mu
# Distribution of times from split until complete separation
if bdGrowth is not None :
mm = spr*(0.5/bdGrowth)
else :
# Crude: take mean from tree
mm = mean([stree.node(n).data.branchlength
for n in stree.all_ids() if n != stree.root])
e1 = randomDistributions.LogNormal(mm, .25)
nh = nodeHeights(stree)
# internals, leaves to root order
internals = set(stree.all_ids()) - set(stree.get_terminals())
internals = [z for y,z in sorted([(nh[x], x) for x in internals])]
for x in internals :
n = stree.node(x)
h = nh[x]
tOrig = t = e1.sample()
if restrictToTree:
# Clip migration "stop" times if they go beyond stop time of children
for ch in n.succ:
nch = stree.node(ch)
if not nch.data.taxon:
assert nh[n.id] >= nh[nch.id]
bound = nh[ch] - nch.data.imc
if h - t < bound:
t = h - bound
if h > t :
d = demographic.LinearPiecewisePopulation([0, 0, pim], [h-t, h])
n.data.imc = t
else :
if h == t :
assert tOrig > h
t = tOrig
d = demographic.LinearPiecewisePopulation([(1 - h/t)*pim , pim], [h])
n.data.imc = h
n.data.ima = (d, d)
def summaryTreeUsingCA(tree, xtrees, atz = False) :
tree = copy.deepcopy(tree)
sClades = getTreeClades(tree)
# Target clades in tree to set height of
cladeSets = [frozenset(c) for c,n in sClades]
#stats = [[0.0,0.0] for c in cladeSets]
stats = [[0.0] for c in cladeSets]
NS = 0
for t in xtrees:
nhs = nodeHeights(t, allTipsZero = atz)
# Store height estimate from tree per target clade
h = [0.0]*len(cladeSets)
# We depend on those being in pre order
for c,n in getTreeClades(t) :
c = frozenset(c)
for k,x in enumerate(cladeSets) :
if x.issubset(c) :
h[k] = nhs[n.id]
NS += 1
for st,hk in zip(stats,h) :
st[0] += hk
#st[1] += hk**2
for (c,n),s in zip(sClades,stats) :
n.data.hh = s[0]/NS
# use t, nhs from last iteration
for tx in tree.get_terminals():
if atz :
h1 = 0.0
else :
n2 = t.search_taxon(tree.node(tx).data.taxon)
h1 = nhs[n2]
tree.node(tx).data.hh = h1
for nid in tree.all_ids() :
if nid != tree.root :
n = tree.node(nid)
ph = tree.node(n.prev).data.hh
h = n.data.hh
assert ph >= h
n.data.branchlength = ph - h
return tree
def setIMrates(stree, mSpec, sSpec, balanced = None, restrictToTree = True) :
nh = nodeHeights(stree)
# internals, leaves to root order
internals = set(stree.all_ids()) - set(stree.get_terminals())
internals = [z for y,z in sorted([(nh[x], x) for x in internals])]
for x in internals :
n = stree.node(x)
pim = mSpec.sample()
if sSpec is not None:
h = nh[x]
tOrig = t = sSpec.sample()
if restrictToTree:
# Clip migration "stop" times if they go beyond stop time of children
for ch in n.succ:
nch = stree.node(ch)
if not nch.data.taxon:
assert nh[n.id] >= nh[nch.id]
bound = nh[ch] - nch.data.imc
if h - t < bound:
t = h - bound
if h > t :
d = demographic.LinearPiecewisePopulation([0, 0, pim], [h-t, h])
n.data.imc = t
else :
if h == t :
assert tOrig > h
t = tOrig
d = demographic.LinearPiecewisePopulation([(1 - h/t)*pim , pim], [h])
n.data.imc = h
n.data.ima = (d, d)
else: # sSpec is None
d = demographic.ConstantPopulation(pim)
if balanced :
n.data.ima = (d, d)
else :
d1 = demographic.ConstantPopulation(mSpec.sample())
n.data.ima = (d, d1)
def _setTreeHeights(tree, opts,fctr) :
order = getPostOrder(tree)
nhs = nodeHeights(tree, allTipsZero = False)
for node in order :
if not node.succ:
node.data.height = nhs[node.id]*fctr
else :
hs = [tree.node(x).data.height for x in node.succ]
mn = sum([h + opts[x] for x,h in zip(node.succ,hs)])/len(hs)
node.data.height = max(*(hs + [mn]))
for n in tree.all_ids() :
node = tree.node(n)
if node.prev is not None:
p = tree.node(node.prev)
node.data.branchlength = p.data.height - node.data.height
assert node.data.branchlength >= 0
return tree
def allPartitions(referenceTree, trees, func = None, withHeights = False,
withRoot = False) :
""" Clade information summary for a set of trees.
Summerize clades from all trees in one mapping. All trees must be on the
same taxa as the reference tree. Return a mapping whose key is the clade, and
the value is a sequence of (tree,node) pairs. If func is given, the values are
the results of applying func to the (tree,node) pair. The clade is a (frozen)
set of integers, each integer is the index of the taxon in the reference tree.
"""
if withRoot :
rootHeights = []
taxa = referenceTree.get_taxa()
ultrametric = True
if withHeights :
nhs = nodeHeights(referenceTree, allTipsZero = False)
nhsd = [(referenceTree.node(i).data.taxon, nhs[i])
for i in nhs if not referenceTree.node(i).succ]
ultrametric = all([h <= 1e-13 for t,h in nhsd])
p = dict()
for tree in trees:
p1 = dict()
if ultrametric :
h = _collectCladeTaxa(tree, tree.root, taxa, p1, withHeights)
else :
h = _collectCladeTaxaNonATZ(tree, taxa, p1)
for k,nd in p1.iteritems() :
pk = p.get(k)
v = (tree, nd)
if func :
v = func(*v)
if pk is None :
p[k] = [v]
else :
pk.append(v)
if withRoot :
rootHeights.append(h)
return (p,rootHeights) if withRoot else p
def mauCanonical(tree, internalNodes = True) :
""" Convert tree to Mau representation.
Result is a tuple containing two lists [taxa-names, internal-nodes-heights].
if 'internalNodes' the tuple contains a third list containing the node id of
each internal node and a boolean indication if the childrens have been swaped
of not.
>>> mauCanonical(Trees.Tree('((chimp:1,human:1):1,gorilla:2)')) in \
[[['gorilla', 'chimp', 'human'], [2.0, 1.0]], \
[['chimp', 'human', 'gorilla'], [1.0, 2.0]], \
[['human', 'chimp', 'gorilla'], [1.0, 2.0]]]
True
@param tree:
@type tree: ITrees.Tree
@param internalNodes: when True, return low level information about the
location of internal nodes and if their siblings has been swapped.
"""
nh = nodeHeights(tree)
return _mauCanonicalSub(tree, tree.node(tree.root), nh, internalNodes)
def _setTreeHeightsForTargets(tree, ftargets, fctr) :
for i,h in ftargets() :
tree.node(i).data.height = h
nhs = nodeHeights(tree, allTipsZero = False)
order = getPostOrder(tree)
for node in order :
if not node.succ:
node.data.height = nhs[node.id]*fctr
else :
node.data.height = max([node.data.height]+
[tree.node(x).data.height for x in node.succ])
for n in tree.all_ids() :
node = tree.node(n)
if node.prev is not None:
p = tree.node(node.prev)
node.data.branchlength = p.data.height - node.data.height
assert node.data.branchlength >= 0
for i in tree.all_ids() :
del tree.node(i).data.height
def summaryTreeUsingMedianHeights(tree, xtrees) :
tree = copy.deepcopy(tree)
func = lambda t,(n,h) : h
posteriorParts,rhs = allPartitions(tree, xtrees, func = func,
withHeights = True, withRoot = True)
treeParts = allPartitions(tree, [tree])
for k in treeParts :
# Node id
nn = treeParts[k][0][1]
if k in posteriorParts :
tree.node(nn).data.height = median(posteriorParts[k])
else :
raise RuntimeError("tree incompatible with trees")
# Assume all trees share same tip heights (not checked)
nh = nodeHeights(xtrees[0])
tree.node(tree.root).data.height = median(rhs)
for n in getPostOrder(tree):
if not len(n.succ) :
n.data.height = nh[xtrees[0].search_taxon(n.data.taxon)]
else :
# Make sure node is heigher than descendants
n.data.height = max([n.data.height] + [tree.node(x).data.height for x in n.succ])
for n in tree.all_ids() :
node = tree.node(n)
if node.prev is not None:
p = tree.node(node.prev)
node.data.branchlength = p.data.height - node.data.height
assert node.data.branchlength >= 0
return tree
def getGTorder(gtree, stree, nTries = 3, perTreeTries = 3, tryPairs=True) :
""" get genes order to plot inside a species tree.
preliminary version.
"""
gtax = []
gtx = defaultdict(lambda : [])
for n in gtree.get_terminals() :
gn = gtree.node(n)
gtx[gn.data.snode.id].append(gn)
gtax.append(gn)
nh = nodeHeights(gtree)
for x in gtree.all_ids() :
gtree.node(x).data.ht = nh[x]
ms, mp = _getGTorderFixedStree(gtree, stree, gtax, gtx, tryPairs)
sint = [stree.node(n) for n in stree.all_ids() if stree.node(n).succ]
for nk in range(nTries) :
cont = True
while cont :
random.shuffle(sint)
cont = False
for n in sint :
for ll in range(perTreeTries) :
sc = n.succ
n.succ = [sc[1],sc[0]]
tms, tmp = _getGTorderFixedStree(gtree, stree, gtax, gtx, tryPairs)
if tms < ms :
ms, mp = tms, tmp
cont = True
else :
n.succ = sc
for g,o in zip(gtax,mp) :
g.data.o = o
return ms, mp, gtx
## if f(0) * f(1) < 0 :
## sol = brentq(f, 0, 1)
## else :
## sol = thc/2
return sol/3
if 0:
k = len(seqs)//nMax
if max(k,ns) > len(seqs)//50 :
return thc/2
sq = random.sample(seqs, max(k,ns))
t,ds = treeFromSeqs(sq, matchScores = scores)
del ds
th = sorted(nodeHeights(t).values())[-(ns//k)]
#import pdb; pdb.set_trace()
return th
if matches is None :
#print "matches"
matches = _buildLookupC(seqs)
#print "done matches"
# lseqs = [len(x) for x in seqs]
fdis = lambda s,j : calign.globalAlign(s, seqs[j], scores = scores,
report = calign.JCcorrection)
nInClade = max((len(seqs)//nMax) + 1, nMax)
lim = 2*nInClade
v = []
li = random.sample(range(len(seqs)), min(ns,len(seqs)))
def guesstimateTH(seqs, nMax, thc, matches = None, scores = None, ns = 100) :
d = [[calign.globalAlign(x1,x2, report = calign.JCcorrection, scores = scores) for (x1,x2) in
[random.sample(seqs, 2)]][0] for k in range(2000)]
q = 1- 1./(len(seqs)/nMax)
p = ([int(nMax * q**n) for n in range(0,100)])
th = sum([nPairs(x) for x in p])/nPairs(sum(p))/2
f = lambda t : (sum([x < t for x in d])/len(d)) - th
#import pdb; pdb.set_trace()
if f(0) * f(1) < 0 :
sol = brentq(f, 0, 1)
else :
sol = thc/2
return sol
k = len(seqs)//nMax
if max(k,ns) > len(seqs)//50 :
return thc/2
sq = random.sample(seqs, max(k,ns))
t,ds = treeFromSeqs(sq, matchScores = scores)
del ds
th = sorted(nodeHeights(t).values())[-(ns//k)]
#import pdb; pdb.set_trace()
return th
if matches is None :
#print "matches"
matches = _buildLookupC(seqs)
#print "done matches"
# lseqs = [len(x) for x in seqs]
fdis = lambda s,j : calign.globalAlign(s, seqs[j], scores = scores,
report = calign.JCcorrection)
nInClade = max((len(seqs)//nMax) + 1, nMax)
lim = 2*nInClade
v = []
li = random.sample(range(len(seqs)), min(ns,len(seqs)))
for i in li:
p = getMates(seqs[i], seqs, .3, fdis, matches, lim = lim)
v.append((i,sorted(p[1])))
f3 = lambda x : mean(x) if len(x) else 1
f1 = lambda th : f3([1./x for x in [sum([x < th for x in u[1]]) for u in v] if x > 0])
#import pdb; pdb.set_trace()
#f1 = lambda th : mean([1/max(sum([x < th for x in u[1]]),0.0001) for u in v])
if f1(1) >= 1./nInClade :
f = lambda x : f1(1) * 1.01 - f1(x)
else :
f = lambda x : f1(x) - 1./nInClade
assert f(0) * f(1) < 0
sol = brentq(f, 0, 1)
f2 = lambda th : 1 - (sum([sum([x < th for x in u[1]]) for u in v])/sum([len(u[1]) for u in v]))
f = lambda x : f2(x) - 1./nInClade
#import pdb; pdb.set_trace()
assert f(0) * f(1) <= 0
sol1 = brentq(f, 0, 1)
th = max(sol,sol1)
th = int(th*1000+.5)/1000
return th,matches
def minPosteriorHSDistanceTree(tree, trees, limit = scipy.inf, norm = True,
nodesMinHeight = None, withDerivative = False,
withInit = True, factr=10000000.0,
warnings = True) :
""" Find a branch length assignment for tree which minimizes the total
distance to the set of trees.
limit is an upper bound (presumably from a prior call here with another tree).
If the distance is known to be larger, the optimization for this tree can be
skipped.
"""
#assert not withDerivative
# not correct for tip/node lower bounds
assert nodesMinHeight is None
treeParts = allPartitions(tree, [tree])
posteriorParts = allPartitions(tree, trees,
func = lambda t,(n,h) : (h, t.node(n).data.branchlength),
withHeights = True)
# For numerical stability sake in computing gradients, scale trees
# so that mean root height is 1
fctr = len(trees)/ sum([treeHeight(x) for x in trees]) if norm else 1
if verbose: print fctr
# Text of expression to compute the total distance. The variables are the
# branch lengths.
ee = ""
dee = ""
pee = ""
for r,k in enumerate(treeParts) :
nn = treeParts[k][0][1]
if k in posteriorParts :
# A tree clade which is in some posterior trees
# Branchs from posterior for this clade
br = [(h*fctr,b*fctr) for h, b in posteriorParts[k]]
# Number of posterior trees without the clade
a1 = (len(trees) - len(posteriorParts[k]))
assert len(trees) == a1 + len(br)
if not tree.node(nn).data.taxon :
if withDerivative :
pee += " ab%d,abd%d = _absDiffBranchDer(b%d,%s)\n" % (nn,nn,nn,_prepare(br))
pee += " abd%d += %d\n" % (nn,a1)
ee += "+ ab%d" % (nn,)
else :
ee += "+ _absDiffBranch(bs%d,%s)" % (nn,_prepare([x[0] for x in br]))
ee += "+ %d * b%d" % (a1,nn)
if withDerivative:
dee += "+(abd%d * d_b%d_x[k])" % (nn,nn)
else :
# A tree clade not appearing in posterior: contributes the full branch
# length for each posterior tree.
ee += "+(%d * b%d)" % (len(trees), nn)
if withDerivative:
dee += "+(%d * d_b%d_x[k])" % (len(trees), nn)
# Constant term of total distance (independent from branch lengths)
c0 = 0
for k in posteriorParts :
if k not in treeParts:
c0 += sum([b * fctr for b in posteriorParts[k]])
# Total distance of branches terminating at a clade which is missing in tree.
# This is (not necessarily good) lower bound on the total distance.
z0 = c0
del posteriorParts
# Tuck in the constant
ee = ("%.15g " % c0) + ee
if z0 >= limit :
return (None, 0.0)
# Get code which transforms the heights encoding to branch lengths
# A descendant height is specified as a fraction in [0,1] of its ancestor
# height (but leading number is the root height).
ba,minRoot,htox = _treeBranchAssignmentExprs(tree, treeParts, fctr,
nodesMinHeight = nodesMinHeight,
withDerivative = withDerivative,
withInit = True,
withHeights = True)
# Define the posterior distance function on the fly.
cod = ("def v1score(x):\n " + "\n ".join(ba) + "\n" + pee + "\n return " +
(('(' + ee + ", array([(" + dee + ") for k in range(nDer)]) )")
if withDerivative else ee))
exec cod
#v1score = v1score1
if verbose: print cod
# Number of variables (heights)
nx = len(tree.get_terminals())-1
xcod = "def htox():\n x = [0]*%d\n " % nx + "\n ".join(htox) \
+ "\n return x"
exec xcod
#global code2branches
# Function to get the branch lengths from optimized heights
codb = "def code2branches(x):\n " + "\n ".join(ba) + "\n " + \
"return (" + ",".join(['(%d,b%d)' % ((treeParts[k][0][1],)*2) for k in treeParts]) + ")"
exec codb in globals()
if verbose: print cod
#code2branches = code2branchesa
if verbose :
print "@@",nx, minRoot, treeHeight(tree) * fctr
print cod
if 0 :
cod8 = "def tv1score(tree):\n"
for r,k in enumerate(treeParts) :
cod8 += " b%d = tree.node(%d).data.branchlength\n" % ((treeParts[k][0][1],)*2)
cod8 += " return " + ee
exec cod8
nhs = nodeHeights(tree, allTipsZero = False)
cod9 = "def tv1scoreh(hs):\n"
hs = _getNodeIDsDescendingHeight(tree, tree.root, 0)
for b in _getNodeIDsDescendingHeight(tree, tree.root, 1)[1:] :
nd = tree.node(b)
if not nd.succ:
nh = "%.15g" % nhs[b]
else :
nh = "hs[%d]" % hs.index(b)
cod9 += " b%d = hs[%d] - %s\n" % (b, hs.index(tree.node(b).prev), nh)
cod9 += " return " + ee
exec cod9
#ssx = copy.deepcopy(tree)
#global xtrees
#xtrees = trees
#exec """def v1scorex(zz0) : return v1score_ck(v1score, zz0, totr(zz0), xtrees)"""
maxfun = 15000
if 1 :
# small protection against disasters
while True:
if withInit :
x0 = htox()
if norm:
x0[0] = 1
else :
x0 = [1 if norm else treeHeight(tree)] + \
[random.random() for k in range(nx-1)]
initialVal = v1score(x0)
assert x0[0] >= minRoot
#pdb.set_trace()
#global mcalls, dcalls
#mcalls, dcalls = 0,0
zz = scipy.optimize.fmin_l_bfgs_b(v1score, x0,
approx_grad=0 if withDerivative else 1,
bounds = [[minRoot,None]] + [[0,1]]*(nx-1),
factr = factr,
iprint=-1, maxfun=maxfun)
if warnings and zz[2]['warnflag'] != 0 :
print "WARNING:", zz[2]['task']
finaleVal = v1score(zz[0])
if finaleVal < initialVal :
break
withInit = False
factr /= 10
if factr < 1e6 :
# failed, leave as is
zz = htox()
finaleVal = v1score(zz[0])
else :
zz = scipy.optimize.fmin_tnc(v1score,
[treeHeight(tree)*fctr] +
[random.uniform(.8,.9) for k in range(nx-1)],
#[.8 for k in range(nx-1)],
approx_grad=1,
bounds = [[minRoot,None]] + [[0,1]]*(nx-1),
maxfun=maxfun,
messages=0)
assert zz[2] == 1
# Do not change tree passed as argument. Copy tree and set branch lengths of
# the copy.
ss = copy.deepcopy(tree)
brs = code2branches(zz[0])
for nn,br in brs:
ss.node(nn).data.branchlength = br/fctr
## for r,k in enumerate(treeParts) :
## ss.node(treeParts[k][0][1]).data.branchlength = brs[r]/fctr
val = finaleVal
if withDerivative :
val = val[0]
return (ss, val/fctr if norm else val)
def _treeBranchAssignmentExprs(tree, clades, fctr, nodesMinHeight = None,
withDerivative = False, withInit = False,
withHeights = False, paranoid = False) :
""" nodesMinHeight: minimum height (lower bound) for internal nodes
paranoid: add internal consistency checks
"""
# tree should be scaled by fctr already
allid = set(tree.all_ids())
# taxa
terms = set(tree.get_terminals())
# internal nodes
allint = allid - terms
# Works for dated tips as well
nh = nodeHeights(tree, allTipsZero = False)
# Node before its descendants
nInOrder = getPreOrder(tree, includeTaxa=False)
# Reversed, child before parent
nhInOrder = [(nh[x],x) for x in reversed(nInOrder)]
# Mapping (per node) of minimum height of node, which is the max among all of
# its descendants
mh = dict()
for n in terms:
mh[n] = nh[n]
for h,n in nhInOrder:
mh[n] = max([mh[c] for c in tree.node(n).succ])
if nodesMinHeight is not None :
for n in nodesMinHeight:
mh[n] = max(mh[n], nodesMinHeight[n]*fctr)
#if fctr != 1 :
# for n in mh :
# mh[n] = mh[n] * fctr
# x[0] is root
if withInit :
htox = []
sr = []
if paranoid:
# solver can send values out of range
sr.append("x = [max(x[0],%f)] + [min(max(z,0),1.0) for z in x[1:]]" % mh[tree.root])
sr.append("h%d = x[0]" % nInOrder[0])
if withInit:
htox.append("x[0] = %.15g" % nh[nInOrder[0]])
if withDerivative:
sr.append("d_h%d_x = [1] + [0]*%d" % (nInOrder[0],len(nInOrder)-1))
sr.append("nDer = %d" % len(nInOrder))
for i,k in enumerate(nInOrder[1:]):
h = tree.node(k).prev
if mh[k] != 0 :
m = "%.15g" % mh[k]
sr.append("h%d = x[%d] * (h%d - %s) + %s" % (k, i+1, h, m, m))
if withInit:
htox.append("x[%d] = %.15g" % (i+1, hs2ratio(nh[k], nh[h], mh[k])))
else :
sr.append("h%d = x[%d] * h%d" % (k, i+1, h))
if withInit:
htox.append("x[%d] = %.15g" % (i+1, hs2ratio(nh[k], nh[h], 0)))
if withDerivative:
sr.append("d_h%d_x = [%s]" % (k,der(i+1, len(nInOrder), h)))
if mh[k] != 0 :
sr.append("d_h%d_x[%d] -= %s" % (k,i+1,m))
if paranoid:
sr.append("assert h%d >= 0" % k)
for r,k in enumerate(clades) :
n = clades[k][0][1] ; assert n != tree.root
p = tree.node(n).prev
if n in terms:
if mh[n] != 0 :
sr.append("b%d=(h%d - %.15g) # %d" % (n, p, mh[n], n))
if withHeights: sr.append("bs%d = %.15g" % (n, mh[n]))
else :
sr.append("b%d=(h%d) # %d" % (n, p, n))
if withHeights: sr.append("bs%d = 0" % (n))
if withDerivative:
sr.append("d_b%d_x = d_h%d_x" % (n,p))
else :
sr.append("b%d=(h%d - h%d) # %d" % (n, p, n, n))
if withHeights: sr.append("bs%d = h%d" % (n,n))
if withDerivative:
sr.append("d_b%d_x = [u-v for u,v in zip(d_h%d_x, d_h%d_x)]" % (n,p,n))
if paranoid:
sr.append("assert b%d >= 0, (%d,b%d,h%d,x)" % (n,r,r,p))
if withInit:
return sr, mh[tree.root], htox
return sr, mh[tree.root]
