import matplotlib.pyplot as plt
import ete3
import numpy as np
np.seterr(divide='ignore', invalid='ignore')
import wquantiles
import itertools,argparse,gzip,sys
def getMidPos_method2(pos1, pos2, depth1, depth2, depth3): return((1.*pos2*(depth2-depth1) + pos1*(depth3-depth1))/(depth3-2*depth1+depth2))
def getNodePos(node, method = 1):
children = node.get_children()
if method == 1:
c1,c2 = children
assert len(children)==2, "Position method one only works for bifurcating nodes."
return((1.*c2.pos*(c1.depth-node.depth)+c1.pos*(c2.depth-node.depth))/(c2.depth-2*node.depth+c1.depth))
elif method == 2:
return(np.mean([c.pos for c in node.get_children()]))
else: raise "Position method can only be 1 or 2."
def drawTree(tree, leafPos = None, depthDict = None, depthRangeDict=None, extendTips=False, rootIsZero=False,
show=True, posMethod=1, col="black",linewidth=2,alpha=1,direction="down",taxColDict=None):
tree = tree.copy("newick")
#get node depths
for node in tree.traverse():
if depthDict is None: node.add_feature("depth", node.get_distance(tree))
else: node.add_feature("depth", depthDict[node.name])
#extend tips to align them if needed
if extendTips:
maxDP = max([l.depth for l in tree.iter_leaves()])
for l in tree.iter_leaves(): l.depth = maxDP
#adjust depths so that they are aligned at zero
if not rootIsZero:
maxDP = max([l.depth for l in tree.iter_leaves()])
for node in tree.traverse():
node.depth -= maxDP
#set leaf positions
if leafPos is None:
leafNames = [l.name for l in tree.get_leaves()]
leafPos = dict(zip(leafNames,range(len(leafNames))))
for leaf in tree.iter_leaves(): leaf.add_feature("pos", leafPos[leaf.name])
#set positions for all other nodes relative to their children
for node in tree.traverse(strategy="postorder"):
if not node.is_leaf():
node.add_feature("pos", getNodePos(node,method=posMethod))
#draw
for node in tree.traverse():
if direction is "down":
plt.setp(plt.gca(),xticks=[])
for child in node.get_children():
plt.plot([node.pos,child.pos],[node.depth,child.depth],color=col,linewidth=linewidth,alpha=alpha, solid_capstyle="round")
if depthRangeDict:
plt.plot([node.pos]*2,depthRangeDict[node.name],color=col,linewidth=1,alpha=alpha, solid_capstyle="round")
plt.plot([node.pos-.1,node.pos+.1],[node.depth]*2,color=col,linewidth=1,alpha=alpha, solid_capstyle="round")
if node.is_leaf(): plt.text(node.pos, node.depth - 0.1, node.name,
horizontalalignment='center', verticalalignment='center',
color=taxColDict[node.name] if taxColDict else "black")
else:
plt.setp(plt.gca(),yticks=[])
for child in node.get_children():
plt.plot([node.depth,child.depth],[node.pos,child.pos],color=col,linewidth=linewidth,alpha=alpha, solid_capstyle="round")
if depthRangeDict:
plt.plot(depthRangeDict[node.name],[node.pos]*2,color=col,linewidth=1,alpha=alpha, solid_capstyle="round")
plt.plot([node.depth]*2,[node.pos-.1,node.pos+.1],color=col,linewidth=1,alpha=alpha, solid_capstyle="round")
if node.is_leaf(): plt.text(node.depth - 0.1, node.pos, node.name,
horizontalalignment='left', verticalalignment='center',
color=taxColDict[node.name] if taxColDict else "black")
if show: plt.show()
def subset(things,subLen):
starts = range(0,len(things),subLen)
ends = [start+subLen for start in starts]
return [things[starts[i]:ends[i]] for i in range(len(starts))]
def asColumn(a):
return a[:,np.newaxis]
def addNodeNames(tree):
for n in tree.traverse():
if n.name is not None: n.name = "_".join(sorted(n.get_leaf_names()))
allIsNaN = lambda x: np.all(np.isnan(x))
def normLength(tree, outgroup):
l = 1.*tree.get_distance(outgroup)
for node in tree.traverse(): node.dist /= l
def addNodeNames(tree):
for n in tree.traverse():
if n.name is "": n.name = "_".join(sorted(n.get_leaf_names()))
def treeToParentChildTable(tree):
return [(n.up.name,n.name,n.dist) for n in tree.traverse() if n.up is not None]
def getLeafPairs(node):
assert len(node.children) == 2, "Node {} does not have two children.".format(node.name)
return itertools.product(node.children[0].get_leaf_names(), node.children[1].get_leaf_names())
def make2DarrayFrom1DupperTriangle(upperTriangle1D,N,includesDiagnol=False):
a = np.zeros([N,N])
n=N if includesDiagnol else N-1
indices = list(np.triu_indices(n))
if not includesDiagnol: indices[1]+=1
a[indices] = a[indices[::-1]] = upperTriangle1D
return a
########################## plot topos with average branch lengths
parser = argparse.ArgumentParser()
parser.add_argument("-w", "--weightsFiles", help="Input weights file(s) from Twisst", action = "store", nargs= "+", required = True)
parser.add_argument("-d", "--distsFiles", help="Input dists file(s) from Twisst", action = "store", nargs = "+", required = True)
parser.add_argument("-f", "--figFile", help="File for output figure", action = "store", required = True)
parser.add_argument("--figFormat", help="Format of figFile", action = "store", default="pdf")
parser.add_argument("--figSize", help="Size of figFile", action = "store", nargs=2, type=float, default=(10,10,))
parser.add_argument("--posMethod", help="Node positioning method", choices=(1,2,), type=int, action = "store", default = 2)
parser.add_argument("--quantiles", help="Add quantiles for each node in tree", type=float, nargs = 2, action = "store")
parser.add_argument("--plotTaxa", help="Prune tree to include on the specifed taxa", nargs = "+", action = "store")
parser.add_argument("--taxOrder", choices = ("levelorder", "preorder", "postorder", "predefined"), action = "store", default="levelorder",
help="How to determine order of taxa in plots")
parser.add_argument("--lineWidth", help="Width for tree lines", type=float, action = "store", default= 4)
parser.add_argument("--scaleLinesByWeights", help="Scale tree lines in figure by weights", action = "store_true")
parser.add_argument("--orderByWeights", help="Order tree plots by weights", action = "store_true")
parser.add_argument("--cols", help="Topology colours", nargs = "+", action = "store")
parser.add_argument("--taxCols", help="Taxon name colours", nargs = "+", action = "store")
parser.add_argument("--alpha", help="Topology alpha", type=float, action = "store", default=1.)
parser.add_argument("--layout", help="Rows and columns to plot", nargs=2, type=int, action = "store")
parser.add_argument("--tight", help="Pading for tight edges", nargs=2, type=float, action = "store")
args = parser.parse_args()
sys.stderr.write("\nReading distances file...")
dists = np.vstack([np.loadtxt(f, skiprows=1) for f in args.distsFiles])
#get topologies
sys.stderr.write("\nGetting topologies...")
topos = []
with gzip.open(args.weightsFiles[0], "r") as wf:
while True:
try: topos.append(ete3.Tree(wf.readline().split()[-1]))
except: break
nTopos = len(topos)
for t in topos: addNodeNames(t)
nTaxa = len(topos[0].get_leaves())
sys.stderr.write("\nThere are {} topologies and {} taxa".format(nTopos,nTaxa))
#make a separate set of topologies for plotting
plotTopos = [t.copy("newick") for t in topos]
if args.plotTaxa:
for t in plotTopos: t.prune(args.plotTaxa)
if args.taxOrder == "predefined":
assert args.plotTaxa is not None, "Predefined taxa order must be given using --plotTaxa."
taxOrder = [args.plotTaxa]*nTopos
else:
taxOrder = [[node.name for node in topo.traverse(strategy=args.taxOrder) if node.is_leaf()] for topo in plotTopos]
if args.taxCols:
try: taxColDict = dict(zip(args.plotTaxa,args.taxCols))
except: raise ValueError("To plot coloured taxon labels, you must specify names of taxa using --plotTaxa")
else:
taxColDict = None
if args.layout: nRow,nCol = args.layout
elif nTopos == 3: nRow,nCol = (1,3,)
elif nTopos == 15: nRow,nCol = (3,5,)
elif nTopos == 105: nRow,nCol = (3,5,)
else: raise ValueError("Please specify number of rows and columns in plot using --layout")
#get pair names in the dists file. The order is essential here.
#We use the order of the first N headers, but assume that the rest follow the same pattern.
#for example, if the taxaare called A, B C and D, the headers should be:
#Topo1_A_B Topo1_A_C Topo1_A_D Topo1_B_C Topo1_B_D Topo1_C_D Topo2_A_B Topo2_A_C ... ect
with gzip.open(args.distsFiles[0], "r") as df: pairs = df.readline().split()
pairs = [pair.split("_")[1:] for pair in pairs][:nTopos]
taxonNames = pairs[0] + [pair[1] for pair in pairs[1:nTaxa]]
#get columns for dists separated by topology
topo_column_indices = subset(range(dists.shape[1]), dists.shape[1]/nTopos)
#split topologies into a third dimension
dists = np.dstack([dists[:,i] for i in topo_column_indices])
dists = np.swapaxes(dists,1,2)
#set all missing dists to zero. Necessary for averaging, and doesnt impact results, because weighting for these is zero
dists = np.nan_to_num(dists)
############# weights
#read weights and convert to proportions
sys.stderr.write("\nReading weights file...")
weights = np.vstack([np.loadtxt(f, skiprows=nTopos+1) for f in args.weightsFiles])
assert weights.shape[0] == dists.shape[0]
assert weights.shape[1] == dists.shape[1]
#convert to proportions
rowSums = np.apply_along_axis(np.sum, 1, weights)
weights = weights / np.reshape(rowSums,[len(rowSums),1])
#convert any nan to zero
weights = np.nan_to_num(weights)
#get means
meanWeights = np.apply_along_axis(np.mean, 0, weights)
# now we need to get the average distance between leaves for each node in each topo
# the first step here is to get the two sets of leaves that descend from each node
nodes_all = [list(tree.traverse()) for tree in topos]
nodeNames = [[n.name for n in nodes] for nodes in nodes_all]
nodeLeafPairs = [[zip(*[(taxonNames.index(x),taxonNames.index(y),) for x,y in getLeafPairs(node)]) if not node.is_leaf() else None for node in nodes] for nodes in nodes_all]
#make a detpths array that gives the depth of each node for each topo at each window
depths = np.zeros([dists.shape[0], dists.shape[1], len(nodeNames[0])])
#now we go line by line, topology by topology and retrieve the depth
# as the avergae pairwise distance between all leaf pairs for each node
# unless the node is a leaf, in which case depth is zero.
sys.stderr.write("\nComputing depth for each node for each topology for each line in input...")
for x in range(depths.shape[0]):
for y in range(nTopos):
distMat = make2DarrayFrom1DupperTriangle(dists[x,y,:], nTaxa)
for z in range(len(nodeNames[y])):
depths[x,y,z] = distMat[nodeLeafPairs[y][z]].mean() if not nodes_all[y][z].is_leaf() else 0.0
#scale depths by dividing by the depth of the root
#the first node in each topo is the root, as the traversal goes to the root first
depths = depths / np.repeat(depths[:,:,0,np.newaxis], depths.shape[2], axis=2)
#anyehere we have nan is where the root depth was zero. This happens where we had missing data. So we can set all these tree depths to zero.
depths = np.nan_to_num(depths)
depths_average = np.average(depths, axis = 0, weights=np.repeat(weights[:,:,np.newaxis], depths.shape[2], axis=2))
depths_median = [[wquantiles.median(depths[:,j,k], weights=weights[:,j]) for k in range(depths.shape[2])] for j in range(depths.shape[1])]
if args.quantiles:
depths_qL = [[wquantiles.quantile(depths[:,j,k], weights[:,j], args.quantiles[0]) for k in range(depths.shape[2])] for j in range(depths.shape[1])]
depths_qU = [[wquantiles.quantile(depths[:,j,k], weights[:,j], args.quantiles[1]) for k in range(depths.shape[2])] for j in range(depths.shape[1])]
#cols = np.array([
#"#2BCE48", #Green
#"#005C31", #Forest
#"#94FFB5", #Jade
#"#9DCC00", #Lime
#"#426600", #Quagmire
#"#00998F", #Turquoise
#"#5EF1F2", #Sky
#"#0075DC", #Blue
#"#003380", #Navy
#"#740AFF", #Violet
#"#FF5005", #Zinnia
#"#F0A3FF", #Amethyst
#"#FFA405", #Orpiment
#"#FF0010", #Red
#"#C20088"]) #Mallow
cols = args.cols if args.cols else ["#000000"]*nTopos
lineWidths = np.array([args.lineWidth]*nTopos, dtype=float)
if args.scaleLinesByWeights: lineWidths *= (meanWeights/meanWeights.max())
plotOrder = np.argsort(meanWeights)[::-1] if args.orderByWeights else range(nTopos)
sys.stderr.write("\nMaking plot.")
plt.figure(figsize=args.figSize, frameon=False)
for i in range(len(plotOrder)):
plt.subplot(nRow,nCol,i+1)
x = plotOrder[i]
for y in np.arange(0,1.1,0.1): plt.plot([0,len(taxOrder[x])+1],[y,y],color="#CCCCCC")
drawTree(plotTopos[x], leafPos = dict(zip(taxOrder[x],np.arange(1,len(taxOrder[x])+1))),
depthDict = dict(zip(nodeNames[x],depths_median[x])),
depthRangeDict = dict(zip(nodeNames[x],zip(depths_qL[x],depths_qU[x]))) if args.quantiles else None,
show=False, alpha = args.alpha, posMethod = args.posMethod,
linewidth = lineWidths[x], col=cols[x], taxColDict=taxColDict)
axes = plt.gca()
axes.set_ylim([-0.1,1.1])
axes.set_xlim([0.5,len(taxOrder[x])+.5])
axes.spines["top"].set_visible(False)
axes.spines["right"].set_visible(False)
axes.spines["bottom"].set_visible(False)
axes.spines["left"].set_visible(False)
plt.text(1,.95,"T"+str(x+1),color=cols[x],
horizontalalignment='left', verticalalignment='center', size=12)
#plt.show()
if args.tight: plt.tight_layout(h_pad=args.tight[0], w_pad = args.tight[1])
plt.savefig(args.figFile, format=args.figFormat, figsize=args.figSize, frameon=False)
plt.close()
sys.stderr.write("\nDone.\n")
