genevenn.py

#Make venn diagram out of two or three lists
#wdecoster

import sys, os, collections, matplotlib, subprocess
matplotlib.use('Agg')
from matplotlib_venn import venn3, venn2
import matplotlib.pyplot as plt

if '-h=1' in sys.argv:
	header=True
	print("Creating lists, assuming a header.")
else:
	header=False
	print("Creating lists, assuming no header.")
for arg in sys.argv:
	if arg.startswith('-names='):
		namesB = True #boolean for names
		namesV = arg[7:].split(',') #values for names
		break
else:
	namesB=False
	
sets, namesF = [], [] #all genesets and the filenames
for infile in sys.argv[1:]:
	if not infile.startswith('-'):
		if os.path.isfile(infile):
			with open(infile) as file:
				if header:
					sets.append(set([line.strip() for line in file.readlines()[1:] if not line == ""])) #Expecting files to be not too big
				else:
					sets.append(set([line.strip() for line in file.readlines() if not line == ""]))
				namesF.append(infile.replace('.txt', '').replace('.tsv', ''))
		else:
			sys.exit("Path to file {} is not correct".format(infile))

if namesB and len(sets) == len(namesV):
	names = namesV
else:
	names = namesF
			
if len(sets) == 2:
	venn2([sets[0], sets[1]], (names[0], names[1]))
else:
	venn3([sets[0], sets[1], sets[2]], (names[0], names[1], names[2]))

with open('SetsGeneVenn.txt', 'w') as output:
	sharednessdict = collections.defaultdict(list)
	for (subset, name) in zip(sets, names):
		for item in subset:
			sharednessdict[item].append(name)
	for key, value in sharednessdict.items():
		sharednessdict[key] = '&'.join(sorted(value))
	inv_sharednessdict = {}
	for key, value in sharednessdict.items():
		inv_sharednessdict[value] = inv_sharednessdict.get(value, [])
		inv_sharednessdict[value].append(key)
	for key, value in inv_sharednessdict.items():
		output.write('Group: {}:\t{}\n'.format(key, ','.join(sorted(value))))
outfig = "GeneVenn.jpeg"
plt.savefig(outfig)
subprocess.call(['eog', outfig])