aq-dulegplot

#!/usr/bin/env python
import pygal
import argparse
import os
import matplotlib
matplotlib.use('SVG')
font = {'size'   : 8}

matplotlib.rc('font', **font)
import matplotlib.pyplot as plt
import numpy as np
#Read off of command line options:
# - input duleg file
# - mapping file
# - output file
# - metadata category name
parser = argparse.ArgumentParser("Tool to plot indicator species analysis results from AXIOME.")
parser.add_argument("-i",metavar="infile",action="store",help="input indicator species .tab file", required=True)
parser.add_argument("-o",metavar="outdir",action="store",help="output directory location", required=False)
parser.add_argument("-m",metavar="mapping",action="store",help="input mapping file location", required=True)
parser.add_argument("-c",metavar="category",action="store",help="metadata category to plot", required=False)
parser.add_argument("-l",metavar="level",action="store",help="comma separated list of taxonomic levels to plot", default="6", required=False)
parser.add_argument("-t",metavar="type",action="store",help="plot type: by sequences (seq) or OTUs (otu)",default="seq", required=False)
args = parser.parse_args()

dulegFile = open(args.i,"rU")
mappingFile = open(args.m,"rU")
levelList = args.l.split(",")
category = args.c
base = os.path.basename(args.i)
if (args.t == "seq"):
	plotType = "seq"
else:
	plotType = "otu"
if (category == None):
	#Attempt to parse out category from the incoming file name
	try:
		category = base.split("_",2)[2].split(".")[0]
	except:
		print "Cannot parse out category name"
		exit(1)
try:
	saveBase = base.split("_",2)[0] + "_" + base.split("_",2)[1]
except:
	saveBase = base
outDir = args.o
clusterList = []
if not os.path.exists(outDir):
    os.makedirs(outDir)


#Read in mapping file
#Find position of metadata of interest
#Return 
def parse_mapping_file(mapping, metadata):
	sampleMemberships = {}
	for line in mapping:
		if line.startswith("#"):
			line = line.split("\t")
			metadataPos = 0
			while (metadataPos < len(line)):
				if line[metadataPos].lower().strip() == metadata.lower():
					break
				metadataPos += 1
			if metadataPos >= len(line):
				print "Error: could not find metadata in mapping file. Is it spelled correctly? Parsed value: " + category
				exit(1)
		else:
			line = line.split("\t")
			sampleMemberships[line[0]] = line[metadataPos]
		clusterList = list(set(sampleMemberships.itervalues()))
	return sampleMemberships, clusterList
	
def clusterSort(line, sumPos, sampleNames, sampleMemberships):
	sortedAbundance = {}
	i = 0
	for sampleAbundance in line[1:sumPos]:
		sampleName = sampleNames[i]
		cluster = sampleMemberships[sampleName]
		if cluster in sortedAbundance:
			sortedAbundance[cluster] += int(float(sampleAbundance))
		else:
			sortedAbundance[cluster] = int(float(sampleAbundance))
		i += 1
	return sortedAbundance

sampleMemberships, clusterList = parse_mapping_file(mappingFile,category)

for level in levelList:
	print "Plotting level " + level
	dulegFile.seek(0)
	taxaDict = {}
	for line in dulegFile:
		if line.startswith("#"):
			line = line.split("\t")
			sumPos = 0
			for word in line:
				if word == "Sum":
					break
				sumPos += 1
			sampleNames = line[1:sumPos]
		else:
			line = line.split("\t")
			abundance = line[sumPos]
			classification = line[sumPos+1]
			classification = classification.split(";")
			ind_cluster = line[sumPos+3]
			#Collapse the columns by cluster
			sortedAbundance = clusterSort(line, sumPos, sampleNames, sampleMemberships)
			try:
				classAtLevel = classification[int(level)-1]
				if classAtLevel not in taxaDict:
					taxaDict[classAtLevel] = {}
				clusterDict = taxaDict[classAtLevel]	
				if (plotType == "seq"):
					for cluster in clusterList:
						if cluster in clusterDict:
							clusterDict[cluster] += int(float(sortedAbundance[cluster]))
						else:
							clusterDict[cluster] = int(float(sortedAbundance[cluster]))
				else:
					if ind_cluster in clusterDict:
						clusterDict[ind_cluster] += 1
					else:
						clusterDict[ind_cluster] = 1
			except:
				pass
	plt.ioff()
	fig = plt.figure(figsize=(15,5),facecolor='w') 
	colournum = 0
	colourlist = [ 'green', 'red', 'blue', 'yellow', 'violet', 'orange', 'cyan', 'purple', 'lime' ]
	N = len(taxaDict.keys())
	ind = np.arange(N)
	width = 0.35
	#Pygal plotting, not super compatible with image editors :(		
	#chart = pygal.StackedBar()
	#chart.title = "Indicator Species by Abundance (Level " + level + ")"
	#chart.x_labels = taxaDict.keys()
	#chart.x_label_rotation = 90
	plt.ylabel('Abundance')
	plt.title('Indicator Species by Abundance (Level ' + level + ')')
	plt.xticks(ind+width/2., taxaDict.keys())
	plt.setp(plt.xticks()[1], rotation=90)
	chartDict = {}
	for taxa in taxaDict.values():
		for cluster in clusterList:
			if cluster not in chartDict:
				chartDict[cluster] = []
			if cluster not in taxa:
				chartDict[cluster].append(0)
			else:
				chartDict[cluster].append(taxa[cluster])
	plotlst = []
	usedcolours = []
	yoff = None
	for cluster in chartDict:
		if (yoff == None):
			pl = plt.bar(ind, chartDict[cluster], width, color=colourlist[colournum%9])
			yoff = chartDict[cluster]
			#Deal with the empty yoff case to prevent cashing matplotlib
			if (len(yoff) == 0):
				yoff = None
		else:
			pl = plt.bar(ind, chartDict[cluster], width, bottom=yoff, color=colourlist[colournum%9])
			yoff = map(sum,zip(chartDict[cluster],yoff))
		plotlst.append(pl)
		usedcolours.append(pl[0])
		colournum += 1
	plt.legend(usedcolours, clusterList)
		#chart.add(str(cluster),chartDict[cluster])
#	chart.render_to_file(os.path.dirname(args.o) + "/" + saveBase + "_" + category + "_L" + level + ".svg")
	plt.savefig(os.path.dirname(args.o) + "/" + saveBase + "_" + category + "_L" + level + '.svg', transparent=False, bbox_inches='tight', pad_inches=0)
	plt.close()