annotate mytools/sampline.py @ 7:f0dc65e7f6c0

Uploaded
author xuebing
date Fri, 09 Mar 2012 19:59:07 -0500
parents
children
Ignore whitespace changes - Everywhere: Within whitespace: At end of lines:
rev   line source
7
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
1 #!/usr/bin/env python
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
2
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
3 """
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
4 Sampling random records from a file. Each record is defined by a fixed number of lines.
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
5
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
6 Usage: sampline.py [options]
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
7
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
8 Options:
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
9 -h, --help show this help message and exit
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
10 -r, --replacement Sampling with replacement
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
11 -i INPUT, --input=INPUT
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
12 Input file
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
13 -o OUTPUT, --output=OUTPUT
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
14 Output file
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
15 -k NSAMPLE, --nSample=NSAMPLE
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
16 (required) number of records to be sampled/output
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
17 -m RECSIZE, --recSize=RECSIZE
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
18 (default=1) number of lines spanned by each record
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
19 -n NSKIP, --nSkip=NSKIP
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
20 (default=0) number of comment lines to skip at the
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
21 beginning
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
22
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
23 example:
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
24 python sampline.py -i test10000.fastq -o out.txt --nSample=5 --recSize=4 --nSkip=0 --replacement
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
25 """
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
26
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
27 import optparse, string, random,sys,math,itertools
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
28
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
29 assert sys.version_info[:2] >= ( 2, 4 )
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
30
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
31 def main():
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
32
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
33 # Parse command line
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
34 parser = optparse.OptionParser( usage="%prog [options] " )
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
35 parser.add_option( "-r", "--replacement", action="store_true", dest="replacement",default=False,
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
36 help="Sampling with replacement" )
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
37 parser.add_option( "-i", "--input", dest="input", default=None,
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
38 help="Input file" )
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
39 parser.add_option( "-o", "--output", dest="output", default=None,
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
40 help="Output file" )
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
41 parser.add_option("-k","--nSample", type='int',dest="nSample",default=None,
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
42 help="(required) number of records to be sampled/output" )
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
43 parser.add_option("-m","--recSize", type='int',dest="recSize",default=1,
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
44 help="(default=1) number of lines spanned by each record" )
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
45 parser.add_option("-n","--nSkip", type='int',dest="nSkip",default=0,
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
46 help="(default=0) number of comment lines to skip at the beginning")
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
47 options, args = parser.parse_args()
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
48 #assert options.region in ( 'coding', 'utr3', 'utr5', 'transcribed' ), "Invalid region argument"
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
49
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
50 sampline(options.input,options.output,options.nSample,options.recSize,options.nSkip,options.replacement)
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
51
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
52 def sample_wr(population, k):
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
53 "Chooses k random elements (with replacement) from a population"
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
54 n = len(population)
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
55 _random, _int = random.random, int # speed hack
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
56 return [_int(_random() * n) for i in itertools.repeat(None, k)]
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
57
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
58 # num of lines
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
59 def readinput(filename):
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
60 try:
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
61 f = open (filename)
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
62 except:
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
63 print >> sys.stderr, "can't open file "+str(filename)
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
64 sys.exit(0)
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
65
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
66 nline = 0
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
67 for line in f:
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
68 nline = nline + 1
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
69 f.close()
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
70 return nline
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
71
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
72 def sampline(infile,outfile,nSample,recSize,nSkip,replacement):
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
73 # sample nSample records from file
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
74 # each record contains recSize lines
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
75 # skip the top nSkip lines
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
76
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
77 nLine = readinput(infile)
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
78 print 'num of lines in input: ',nLine
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
79 print 'avoid sampling the first ',nSkip,' lines'
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
80 print 'lines per record: ',recSize
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
81
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
82 if (nLine-nSkip) % recSize:
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
83 print >> sys.stderr, "the number of lines is not dividable by record size!"
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
84 sys.exit(0)
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
85
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
86 nTotalRecords = (nLine-nSkip) / recSize
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
87 print "total number of records: ",nTotalRecords
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
88
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
89 if replacement or nTotalRecords < nSample:
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
90 sel = sample_wr(range(nTotalRecords),nSample)
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
91 else:
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
92 sel = random.sample(range(nTotalRecords),nSample)
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
93
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
94 #print len(sel), sorted(sel)
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
95
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
96 # output
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
97 try:
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
98 fout = open (outfile,'w')
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
99 except:
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
100 print >> sys.stderr, "can't open file "+str(outfile)
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
101 sys.exit(0)
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
102 fin = open(infile)
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
103 n = 0 # index of line
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
104 rec = "" # to store all content of a record
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
105 nrepeat = 0 # number of times a record is sampled
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
106 curr_rec = -1
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
107 for line in fin:
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
108 if n < nSkip:
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
109 n = n + 1
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
110 fout.write(line)
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
111 continue
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
112
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
113 if not (n-nSkip) % recSize:# a new record
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
114 # print the previous sampled record
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
115 for i in range(nrepeat):
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
116 fout.write(rec)
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
117 curr_rec = (n-nSkip)/recSize
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
118 nrepeat = sel.count(curr_rec)
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
119 if nrepeat: # sampled
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
120 rec = line
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
121 #print curr_rec,nrepeat
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
122 elif (n-nSkip)/recSize == curr_rec:
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
123 rec = rec + line
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
124 n = n + 1
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
125 # if the last record is selected
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
126 if curr_rec == nTotalRecords-1:
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
127 for i in range(nrepeat):
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
128 fout.write(rec)
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
129 fin.close()
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
130 fout.close()
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
131
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
132
f0dc65e7f6c0 Uploaded
xuebing
parents:
diff changeset
133 if __name__ == "__main__": main()