diff tools/sample_seqs/sample_seqs.py @ 1:369515d7738b draft

Uploaded v0.0.1 preview 2, correct typo in test
author peterjc
date Thu, 27 Mar 2014 09:39:49 -0400
parents 4c5b37848acb
children dc55e58fa890
line wrap: on
line diff
--- a/tools/sample_seqs/sample_seqs.py	Tue Feb 18 13:06:14 2014 -0500
+++ b/tools/sample_seqs/sample_seqs.py	Thu Mar 27 09:39:49 2014 -0400
@@ -23,16 +23,19 @@
     sys.exit(err)
 
 if "-v" in sys.argv or "--version" in sys.argv:
-    print "v0.1.0"
+    print("v0.1.0")
     sys.exit(0)
 
 #Parse Command Line
 if len(sys.argv) < 5:
     stop_err("Requires at least four arguments: seq_format, in_file, out_file, mode, ...")
 seq_format, in_file, out_file, mode = sys.argv[1:5]
+if in_file != "/dev/stdin" and not os.path.isfile(in_file):
+    stop_err("Missing input file %r" % in_file)
+
 if mode == "everyNth":
     if len(sys.argv) != 6:
-        stop_err("If using everyNth, just need argument N")
+        stop_err("If using everyNth, just need argument N (integer, at least 2)")
     try:
         N = int(sys.argv[5])
     except:
@@ -40,25 +43,41 @@
     if N < 2:
         stop_err("Bad N argument %r" % sys.argv[5])
     if (N % 10) == 1:
-        print("Sampling every %ist sequence" % N)
+        sys.stderr.write("Sampling every %ist sequence\n" % N)
     elif (N % 10) == 2:
-        print("Sampling every %ind sequence" % N)
+        sys.stderr.write("Sampling every %ind sequence\n" % N)
     elif (N % 10) == 3:
-        print("Sampling every %ird sequence" % N)
+        sys.stderr.write("Sampling every %ird sequence\n" % N)
     else:
-        print("Sampling every %ith sequence" % N)
+        sys.stderr.write("Sampling every %ith sequence\n" % N)
+    def sampler(iterator):
+        global N
+        count = 0
+        for record in iterator:
+            count += 1
+            if count % N == 1:
+                yield record
+elif mode == "percentage":
+    if len(sys.argv) != 6:
+        stop_err("If using percentage, just need percentage argument (float, range 0 to 100)")
+    try:
+        percent = float(sys.argv[5]) / 100.0
+    except:
+        stop_err("Bad percent argument %r" % sys.argv[5])
+    if percent <= 0.0 or 1.0 <= percent:
+        stop_err("Bad percent argument %r" % sys.argv[5])
+    sys.stderr.write("Sampling %0.3f%% of sequences\n" % (100.0 * percent))
+    def sampler(iterator):
+        global percent
+        count = 0
+        taken = 0
+        for record in iterator:
+            count += 1
+            if percent * count > taken:
+                taken += 1
+                yield record
 else:
     stop_err("Unsupported mode %r" % mode)
-if not os.path.isfile(in_file):
-    stop_err("Missing input file %r" % in_file)
-
-
-def pick_every_N(iterator, N):
-    count = 0
-    for record in iterator:
-        count += 1
-        if count % N == 1:
-            yield record
 
 def raw_fasta_iterator(handle):
     """Yields raw FASTA records as multi-line strings."""
@@ -94,24 +113,24 @@
         if not line:
             return # StopIteration 
 
-def fasta_filter_every_N(in_file, out_file, N):
+def fasta_filter(in_file, out_file, iterator_filter):
     count = 0
     #Galaxy now requires Python 2.5+ so can use with statements,
     with open(in_file) as in_handle:
         with open(out_file, "w") as pos_handle:
-            for record in pick_every_N(raw_fasta_iterator(in_handle), N):
+            for record in iterator_filter(raw_fasta_iterator(in_handle)):
                 count += 1
                 pos_handle.write(record)
     return count
 
 try:
     from galaxy_utils.sequence.fastq import fastqReader, fastqWriter
-    def fastq_filter_every_N(in_file, out_file, N):
+    def fastq_filter(in_file, out_file, iterator_filter):
         count = 0
         #from galaxy_utils.sequence.fastq import fastqReader, fastqWriter
         reader = fastqReader(open(in_file, "rU"))
         writer = fastqWriter(open(out_file, "w"))
-        for record in pick_every_N(reader, N):
+        for record in iterator_filter(reader):
             count += 1
             writer.write(record)
         writer.close()
@@ -119,16 +138,16 @@
         return count
 except ImportError:
     from Bio.SeqIO.QualityIO import FastqGeneralIterator
-    def fastq_filter_every_N(in_file, out_file, N):
+    def fastq_filter(in_file, out_file, iterator_filter):
         count = 0
         with open(in_file) as in_handle:
             with open(out_file, "w") as pos_handle:
-                for title, seq, qual in pick_every_N(FastqGeneralIterator(in_handle), N):
+                for title, seq, qual in iterator_filter(FastqGeneralIterator(in_handle)):
                     count += 1
                     pos_handle.write("@%s\n%s\n+\n%s\n" % (title, seq, qual))
         return count
 
-def sff_filter_every_N(in_file, out_file, N):
+def sff_filter(in_file, out_file, iterator_filter):
     count = 0
     try:
         from Bio.SeqIO.SffIO import SffIterator, SffWriter
@@ -148,17 +167,17 @@
         with open(out_file, "wb") as out_handle:
             writer = SffWriter(out_handle, xml=manifest)
             in_handle.seek(0) #start again after getting manifest
-            count = writer.write_file(pick_every_N(SffIterator(in_handle), N))
+            count = writer.write_file(iterator_filter(SffIterator(in_handle)))
             #count = writer.write_file(SffIterator(in_handle))
     return count
 
 if seq_format.lower()=="sff":
-    count = sff_filter_every_N(in_file, out_file, N)
+    count = sff_filter(in_file, out_file, sampler)
 elif seq_format.lower()=="fasta":
-    count = fasta_filter_every_N(in_file, out_file, N)
+    count = fasta_filter(in_file, out_file, sampler)
 elif seq_format.lower().startswith("fastq"):
-    count = fastq_filter_every_N(in_file, out_file, N)
+    count = fastq_filter(in_file, out_file, sampler)
 else:
     stop_err("Unsupported file type %r" % seq_format)
 
-print("Sampled %i records" % count)
+sys.stderr.write("Sampled %i records\n" % count)