fasta-dinucleotide-shuffle_typeEF.py

#!/usr/local/bin/python

import sys, string, random
import sequence_met

# 
# turn on psyco to speed up by 3X
#
if __name__=='__main__': 
	try:
		import psyco
		#psyco.log()
		psyco.full()
		psyco_found = True
	except ImportError:

#    psyco_found = False
		pass
#  print >> sys.stderr, "psyco_found", psyco_found


# altschulEriksonDinuclShuffle.py
# P. Clote, Oct 2003

def computeCountAndLists(s):

	#Initialize lists and mono- and dinucleotide dictionaries
	List = {} #List is a dictionary of lists
	List['A'] = []; List['C'] = [];
	List['G'] = []; List['T'] = [];
	# FIXME: is this ok?
	List['N'] = []
	List['E'] = []
	List['F'] = []
	nuclList   = ["A","C","G","T","N","E","F"]
	s       = s.upper()
	#s       = s.replace("U","T")
	nuclCnt    = {}  #empty dictionary
	dinuclCnt  = {}  #empty dictionary
	for x in nuclList:
		nuclCnt[x]=0
		dinuclCnt[x]={}
		for y in nuclList:
			dinuclCnt[x][y]=0

	#Compute count and lists
	nuclCnt[s[0]] = 1
	nuclTotal     = 1
	dinuclTotal   = 0
	for i in range(len(s)-1):
		x = s[i]; y = s[i+1]
		List[x].append( y )
		nuclCnt[y] += 1; nuclTotal  += 1
		dinuclCnt[x][y] += 1; dinuclTotal += 1
	assert (nuclTotal==len(s))
	assert (dinuclTotal==len(s)-1)
	return nuclCnt,dinuclCnt,List
 
 
def chooseEdge(x,dinuclCnt):
	z = random.random()
	denom=dinuclCnt[x]['A']+dinuclCnt[x]['C']+dinuclCnt[x]['G']+dinuclCnt[x]['T']+dinuclCnt[x]['N']+dinuclCnt[x]['E']+dinuclCnt[x]['F']
	numerator = dinuclCnt[x]['A']
	if z < float(numerator)/float(denom):
		dinuclCnt[x]['A'] -= 1
		return 'A'
	numerator += dinuclCnt[x]['C']
	if z < float(numerator)/float(denom):
		dinuclCnt[x]['C'] -= 1
		return 'C'
	numerator += dinuclCnt[x]['G']
	if z < float(numerator)/float(denom):
		dinuclCnt[x]['G'] -= 1
		return 'G'
	numerator += dinuclCnt[x]['T']
	if z < float(numerator)/float(denom):
		dinuclCnt[x]['T'] -= 1
		return 'T'
	numerator += dinuclCnt[x]['E']
	if z < float(numerator)/float(denom):
		dinuclCnt[x]['E'] -= 1
		return 'E'
	numerator += dinuclCnt[x]['F']
	if z < float(numerator)/float(denom):
		dinuclCnt[x]['F'] -= 1
		return 'F'
		
	dinuclCnt[x]['N'] -= 1
	
	return 'N'

def connectedToLast(edgeList,nuclList,lastCh):
	D = {}
	for x in nuclList: D[x]=0
	for edge in edgeList:
		a = edge[0]; b = edge[1]
		if b==lastCh: D[a]=1
	for i in range(3):
		for edge in edgeList:
			a = edge[0]; b = edge[1]
			if D[b]==1: D[a]=1
	ok = 0
	for x in nuclList:
		if x!=lastCh and D[x]==0: return 0
	return 1

def eulerian(s):
	nuclCnt,dinuclCnt,List = computeCountAndLists(s)
	#compute nucleotides appearing in s
	nuclList = []
	for x in ["A","C","G","T","N","E","F"]:
		if x in s: nuclList.append(x)
	#create dinucleotide shuffle L 
	firstCh = s[0]  #start with first letter of s
	lastCh  = s[-1]
	edgeList = []
	for x in nuclList:
		if x!= lastCh: edgeList.append( [x,chooseEdge(x,dinuclCnt)] )
	ok = connectedToLast(edgeList,nuclList,lastCh)
	return ok,edgeList,nuclList,lastCh


def shuffleEdgeList(L):
	n = len(L); barrier = n
	for i in range(n-1):
		z = int(random.random() * barrier)
		tmp = L[z]
		L[z]= L[barrier-1]
		L[barrier-1] = tmp
		barrier -= 1
	return L

def dinuclShuffle(s):
	if len(s)==0:
		return s
	ok = 0
	while not ok:
		ok,edgeList,nuclList,lastCh = eulerian(s)
	nuclCnt,dinuclCnt,List = computeCountAndLists(s)

	#remove last edges from each vertex list, shuffle, then add back
	#the removed edges at end of vertex lists.
	for [x,y] in edgeList: List[x].remove(y)
	for x in nuclList: shuffleEdgeList(List[x])
	for [x,y] in edgeList: List[x].append(y)

	#construct the eulerian path
	L = [s[0]]; prevCh = s[0]
	for i in range(len(s)-2):
		ch = List[prevCh][0] 
		L.append( ch )
		del List[prevCh][0]
		prevCh = ch
	L.append(s[-1])
	t = string.join(L,"")
	return t
 
def main():

	#
	# defaults
	#
	file_name = None
	seed = 1
	copies = 1
	#print "running"

	#
	# get command line arguments
	#
	usage = """USAGE: %s [options] 
				-f <filename>   file name (required)
				-t <tag>        added to shuffled sequence names
				-s <seed>	random seed; default: %d
				-c <n>    make <n> shuffled copies of each sequence; default: %d
				-h              print this usage message
				""" % (sys.argv[0], seed, copies)

	# no arguments: print usagei

	if len(sys.argv) == 1:
		#print "arg too short"
		print >> sys.stderr, usage; sys.exit(1)

	tag = ""

	# parse command line
	#print "lenthsys",len(sys.argv)
	i = 1
	#print i
	while i < len(sys.argv):
		arg = sys.argv[i]
		#print i, arg
		if (arg == "-f"):
			i += 1
			try: 
				file_name = sys.argv[i]
				#print "filename", file_name
			except: 
				print >> sys.stderr, usage; sys.exit(1)
		elif (arg == "-t"):
			i += 1
			try: 
				tag = sys.argv[i]
				#print "tag",tag
			except: 
				print >> sys.stderr, usage; sys.exit(1)
		elif (arg == "-s"):
			i += 1
			try: 
				seed = string.atoi(sys.argv[i])
				#print "seed",seed
			except: 
				print >> sys.stderr, usage; sys.exit(1)
		elif (arg == "-c"):
			i += 1
			try: 
				copies = string.atoi(sys.argv[i])
				#print "copies",copies
			except: 
				print >> sys.stderr, usage; sys.exit(1)
		elif (arg == "-h"):
			print >> sys.stderr, usage; sys.exit(1)
			#print "help"
		else:
			print >> sys.stderr, "Unknown command line argument: " + arg
			sys.exit(1)
		i += 1

	# check that required arguments given
	if (file_name == None):
		print >> sys.stderr, usage; sys.exit(1)

	random.seed(seed)

	# read sequences
	seqs = sequence_met.readFASTA(file_name,'Methylation DNA')

	for s in seqs:
		str = s.getString()
		#FIXME altschul can't handle ambigs
		name = s.getName()
		#print >> sys.stderr, ">%s" % name
		for i in range(copies):
			shuffledSeq = dinuclShuffle(str)
			if (copies == 1):
				print >> sys.stdout, ">%s\n%s" % (name+tag, shuffledSeq)
			else:
				print >> sys.stdout, ">%s_%d\n%s" % (name+tag, i, shuffledSeq)
if __name__ == '__main__': 
	main()