#!/usr/bin/env python
from sys import argv, exit, stdout, stderr
from random import randint
method = 0
global n
global dataset_filename
subset_filename = ""
rest_filename = ""
def exit_with_help():
print("""\
Usage: {0} [options] dataset number [output1] [output2]
This script selects a subset of the given dataset.
options:
-s method : method of selection (default 0)
0 -- stratified selection (classification only)
1 -- random selection
output1 : the subset (optional)
output2 : rest of the data (optional)
If output1 is omitted, the subset will be printed on the screen.""".format(argv[0]))
exit(1)
def process_options():
global method, n
global dataset_filename, subset_filename, rest_filename
argc = len(argv)
if argc < 3:
exit_with_help()
i = 1
while i < len(argv):
if argv[i][0] != "-":
break
if argv[i] == "-s":
i = i + 1
method = int(argv[i])
if method < 0 or method > 1:
print("Unknown selection method {0}".format(method))
exit_with_help()
i = i + 1
dataset_filename = argv[i]
n = int(argv[i+1])
if i+2 < argc:
subset_filename = argv[i+2]
if i+3 < argc:
rest_filename = argv[i+3]
def main():
class Label:
def __init__(self, label, index, selected):
self.label = label
self.index = index
self.selected = selected
process_options()
# get labels
i = 0
labels = []
f = open(dataset_filename, 'r')
for line in f:
labels.append(Label(float((line.split())[0]), i, 0))
i = i + 1
f.close()
l = i
# determine where to output
if subset_filename != "":
file1 = open(subset_filename, 'w')
else:
file1 = stdout
split = 0
if rest_filename != "":
split = 1
file2 = open(rest_filename, 'w')
# select the subset
warning = 0
if method == 0: # stratified
labels.sort(key = lambda x: x.label)
label_end = labels[l-1].label + 1
labels.append(Label(label_end, l, 0))
begin = 0
label = labels[begin].label
for i in range(l+1):
new_label = labels[i].label
if new_label != label:
nr_class = i - begin
k = i*n//l - begin*n//l
# at least one instance per class
if k == 0:
k = 1
warning = warning + 1
for j in range(nr_class):
if randint(0, nr_class-j-1) < k:
labels[begin+j].selected = 1
k = k - 1
begin = i
label = new_label
elif method == 1: # random
k = n
for i in range(l):
if randint(0,l-i-1) < k:
labels[i].selected = 1
k = k - 1
i = i + 1
# output
i = 0
if method == 0:
labels.sort(key = lambda x: int(x.index))
f = open(dataset_filename, 'r')
for line in f:
if labels[i].selected == 1:
file1.write(line)
else:
if split == 1:
file2.write(line)
i = i + 1
if warning > 0:
stderr.write("""\
Warning:
1. You may have regression data. Please use -s 1.
2. Classification data unbalanced or too small. We select at least 1 per class.
The subset thus contains {0} instances.
""".format(n+warning))
# cleanup
f.close()
file1.close()
if split == 1:
file2.close()
main()