subset_matrix.py

   1 #!/usr/bin/python3
   2 """Select a fraction of rows from a file
   3 """
   4 # Copyright (2015) Diane Trout & California Institute of Technology
   5
   6 # This program is free software; you can redistribute it and/or modify
   7 # it under the terms of the GNU General Public License as published by
   8 # the Free Software Foundation; either version 2 of the License, or
   9 # (at your option) any later version.
  10 #
  11 # This program is distributed in the hope that it will be useful,
  12 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  13 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  14 # GNU General Public License for more details.
  15 #
  16 # You should have received a copy of the GNU General Public License along
  17 # with this program; if not, write to the Free Software Foundation, Inc.,
  18 # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
  19
  20 import argparse
  21 import random
  22 import os
  23 import sys
  24
  25 def main(cmdline=None):
  26     parser = make_parser()
  27     args = parser.parse_args(cmdline)
  28
  29     if args.seed:
  30         random.seed(args.seed)
  31
  32     if args.filename:
  33         instream = open(args.filename[0], 'rt')
  34     else:
  35         parser.error("Please specify input filename")
  36
  37     if args.output:
  38         outstream = open(args.output, 'wt')
  39     else:
  40         outstream = sys.stdout
  41
  42     for line in subset(instream, args.header, args.include):
  43         outstream.write(line)
  44
  45     if args.output:
  46         outstream.close()
  47
  48 def make_parser():
  49     parser = argparse.ArgumentParser()
  50     parser.add_argument("filename", nargs=1,
  51                         help="filename to read from")
  52     parser.add_argument("-o", "--output",
  53                         help="output filename")
  54     parser.add_argument("-i", "--include", default=0.10, type=float,
  55                         help="probability to include a line [0..1]")
  56     parser.add_argument("--header", default=1, type=int,
  57                         help="number of header lines to include")
  58     parser.add_argument("-s", "--seed",
  59                         help="specify seed")
  60     return parser
  61
  62
  63 def subset(instream, header_lines, include_fraction):
  64     """Subset lines from a file
  65
  66     Always include the first specified number of 'header_lines'
  67     then after that include lines if they meet the random threshold
  68     """
  69
  70     while header_lines > 0:
  71         yield next(instream)
  72         header_lines -= 1
  73
  74     for line in instream:
  75         if random.random() < include_fraction:
  76             yield(line)
  77
  78 if __name__ == "__main__":
  79     main()
  80