bioscripts/SimContigs.py at master · lizhao007/bioscripts · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
#!/usr/bin/env python
import sys
import argparse
import random
import copy


def GetOpts():
	group = argparse.ArgumentParser()
	group.add_argument('--min', help='minimum length of contig, default: 15k, you can use both number or string end with k,m', default="15k")
	group.add_argument('--max', help='minimum length of contig, default: 5m, you can use both number or string end with k,m', default="5m")
	group.add_argument('-n', '--n50', help='size of N50, default: 500k, you can use both number or string end with k,m', default="500k")
	group.add_argument('-i', '--input', help='origin fasta file of genome', required=True)
	group.add_argument('-o', '--output', help='filename of simulated data', required=True)

	return group.parse_args()


def ReadFasta(inFasta):
	fastaDB = {}
	with open(inFasta, 'r') as fIn:
		id = ''
		seq = ''
		for line in fIn:
			if line[0] == '>':
				if seq != '':
					fastaDB[id] = seq
				id = line.strip()[1:]
				seq = ''
			else:
				seq += line.strip()
		fastaDB[id] = seq
	return fastaDB


def GenCtgLen(fastaLenDB, cntLowerDB, cntHigherDB, n50, minLen, maxLen):
	ctgLenDB = {}
	for chrn in cntLowerDB:
		ctgLenDB[chrn] = []
		totalLower = 0
		totalHigher = 0
		totalLen = 0
		for i in range(0, cntLowerDB[chrn]):
			tmpLen = random.randint(minLen, n50)
			totalLower += tmpLen
			if totalLower > fastaLenDB[chrn]/2:
				break
			ctgLenDB[chrn].append(tmpLen)
			totalLen += tmpLen
		for i in range(0, cntHigherDB[chrn]):
			tmpLen = random.randint(n50, maxLen)
			totalHigher += tmpLen
			if totalHigher > fastaLenDB[chrn]/2:
				break
			ctgLenDB[chrn].append(tmpLen)
			totalLen += tmpLen
		cntN50 = int((fastaLenDB[chrn]-totalLen)/n50)
		for i in range(0, cntN50):
			tmpLen = random.randint(int(n50-n50*0.1), int(n50+n50*0.1))
			totalLen += tmpLen
			if totalLen > fastaLenDB[chrn]:
				totalLen -= tmpLen
				break
			ctgLenDB[chrn].append(tmpLen)
		ctgLenDB[chrn].append(fastaLenDB[chrn]-totalLen)
	return ctgLenDB


def GenCtgRegions(fastaLenDB, ctgLenDB):
	ctgRegionsDB = {}
	for chrn in fastaLenDB:
		ctgRegionsDB[chrn] = []
		totalCtgLen = 0
		for ctgLen in ctgLenDB[chrn]:
			totalCtgLen += ctgLen
		cntCtg = len(ctgLenDB[chrn])
		lastPos = 0
		ctgLenList = copy.deepcopy(ctgLenDB[chrn])
		for i in range(0, cntCtg):
			index = random.randint(0, cntCtg-1)
			ctgRegionsDB[chrn].append([lastPos, lastPos+ctgLenList[index]])
			lastPos += ctgLenList[index]
			del ctgLenList[index]
			cntCtg -= 1
	return ctgRegionsDB


def SimGenomeCtg(inFasta, outFasta, n50, minLen, maxLen):
	random.seed()
	print("Reading fasta")
	fastaDB = ReadFasta(inFasta)
	fastaLenDB = {}
	cntLowerDB = {}
	cntHigherDB = {}
	for chrn in fastaDB:
		fastaLenDB[chrn] = len(fastaDB[chrn])
		cntLowerDB[chrn] = int(fastaLenDB[chrn]/(minLen+n50))
		cntHigherDB[chrn] = int(fastaLenDB[chrn]/(maxLen+n50))

	print("\nGenerating contigs")
	ctgLenDB = GenCtgLen(fastaLenDB, cntLowerDB, cntHigherDB, n50, minLen, maxLen)
	ctgRegionsDB = GenCtgRegions(fastaLenDB, ctgLenDB)

	print("\nStatistics")
	for chrn in sorted(fastaDB):
		print("\tChromosome:\t%s"%(chrn))
		print("\tChromosome size:\t%d"%(fastaLenDB[chrn]))
		print("\tContig counts:\t%d"%(len(ctgLenDB[chrn])))
		tmpLen = 0
		n50Len = 0
		for ctgLen in sorted(ctgLenDB[chrn], reverse=True):
			tmpLen += ctgLen
			if tmpLen >= fastaLenDB[chrn]/2 and n50Len == 0:
				n50Len = ctgLen
		print("\tContig total size:\t%d"%(tmpLen))
		print("\tN50 size:\t%d\n"%(n50Len))

	print("\nWriting contigs")
	with open(outFasta, 'w') as fOut:
		for chrn in sorted(fastaDB):
			base = 100
			for region in ctgRegionsDB[chrn]:
				s = region[0]
				e = region[1]
				ctgName = "tig%07d"%(base)
				fOut.write(">%s %s %d:%d length=%d\n%s\n"%(ctgName, chrn, s+1, e+1, e-s+1, fastaDB[chrn][s: e]))
				base += 100
	print("\nFinished")


if __name__ == "__main__":
	opts = GetOpts()
	inFasta = opts.input
	outFasta = opts.output
	n50 = opts.n50
	n50 = n50.lower()
	n50 = n50.replace('m', '000000')
	n50 = n50.replace('k', '000')
	n50 = int(n50)
	minLen = opts.min
	minLen = minLen.lower()
	minLen = minLen.replace('m', '000000')
	minLen = minLen.replace('k', '000')
	minLen = int(minLen)
	maxLen = opts.max
	maxLen = maxLen.lower()
	maxLen = maxLen.replace('m', '000000')
	maxLen = maxLen.replace('k', '000')
	maxLen = int(maxLen)
	SimGenomeCtg(inFasta, outFasta, n50, minLen, maxLen)