Script to split a dataset into an 80/20 training/test set by sampling from a binomial distribution

split.py

#!/usr/bin/python

import sys
import numpy.random as rand


'''
Splits a dataset into an 80/20 train/test set by sampling from a binomial
distribution with each datapoint deciding which set it belongs to.
'''

if __name__ == '__main__':
    if len(sys.argv) != 4:
        print("Usage: splitdata inputFname trainFname testFname\n" +
              "   Creates a random 80/20 train/test set")
        sys.exit(1)
    (fname, train_fname, test_fname) = sys.argv[1:]

    try:
        train = open(train_fname, 'w')
        test = open(test_fname, 'w')
        f = open(fname, 'r')

        for l in f:
            if rand.binomial(1, 0.8):
                train.write(l)
            else:
                test.write(l)
    finally:
        train.close()
        test.close()
        f.close()