added gzip compression support

2 files changed, 22 insertions, 4 deletions

diff --git a/quikr.py b/quikr.py
index 478c3bf..401b5ca 100755
--- a/quikr.py
+++ b/quikr.py
@@ -4,13 +4,14 @@ import sys
 import scipy.optimize.nnls
 import scipy.sparse
 import numpy as np
+import quikr_util as qu
 from subprocess import *
 import argparse
 import platform
+import gzip
 
 def main():
 
-
     parser = argparse.ArgumentParser(description=
     "Quikr returns the estimated frequencies of batcteria present when given a \
     input FASTA file. \n \
@@ -44,12 +45,20 @@ def main():
         lamb = args.lamb
     
     xstar = quikr_load_trained_matrix_from_file(args.fasta, args.trained_matrix, args.kmer, lamb)
+
     np.savetxt(args.output, xstar, delimiter=",", fmt="%f")
     return 0
 
 def quikr_load_trained_matrix_from_file(input_fasta_location, trained_matrix_location, kmer, default_lambda):
   
-  trained_matrix = np.load(trained_matrix_location)
+  if qu.isCompressed(trained_matrix_location):
+    print "compressed!"
+    trained_matrix_file = gzip.open(trained_matrix_location, "rb")
+  else:
+    trained_matrix_file = open(trained_matrix_location, "rb")
+  
+  trained_matrix = np.load(trained_matrix_file)
+
   xstar = quikr(input_fasta_location, trained_matrix, kmer, default_lambda)
   return xstar
   
diff --git a/quikr_train.py b/quikr_train.py
index 5089e2a..a427436 100755
--- a/quikr_train.py
+++ b/quikr_train.py
@@ -2,6 +2,7 @@
 import numpy as np
 import os
 import sys
+import gzip
 from subprocess import *
 import platform
 import argparse
@@ -19,7 +20,10 @@ def main():
 
   parser.add_argument("-i", "--input", help="training database of sequences (fasta format)", required=True)
   parser.add_argument("-o", "--output", help="sensing matrix (text file)", required=True)
-  parser.add_argument("-k", "--kmer", type=int, help="kmer size (integer)", required=False )
+  parser.add_argument("-k", "--kmer", help="kmer size (integer)", 
+                      type=int, required=False )
+  parser.add_argument("-z", "--compress", help="compress output (integer)", 
+                      action='store_true', required=False)
 
   args = parser.parse_args()
 
@@ -29,7 +33,12 @@ def main():
   # call the quikr train function, save the output with np.save
   matrix = quikr_train(args.input, args.kmer)
 
-  np.save(args.output, matrix)
+  if args.compress: 
+    output_file = gzip.open(args.output, "wb")
+  else:
+    output_file = open(args.output, "wb")
+
+  np.save(output_file, matrix)
 
   return 0