File: subset.py

package info (click to toggle)
libsvm 3.24%2Bds-6
  • links: PTS, VCS
  • area: main
  • in suites: bookworm, bullseye
  • size: 992 kB
  • sloc: java: 3,680; cpp: 3,146; ansic: 2,253; python: 1,270; makefile: 153; sh: 41
file content (120 lines) | stat: -rwxr-xr-x 3,196 bytes parent folder | download | duplicates (2)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
#!/usr/bin/env python

import os, sys, math, random
from collections import defaultdict

if sys.version_info[0] >= 3:
	xrange = range

def exit_with_help(argv):
	print("""\
Usage: {0} [options] dataset subset_size [output1] [output2]

This script randomly selects a subset of the 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(argv):
	argc = len(argv)
	if argc < 3:
		exit_with_help(argv)

	# default method is stratified selection
	method = 0
	subset_file = sys.stdout
	rest_file = None

	i = 1
	while i < argc:
		if argv[i][0] != "-":
			break
		if argv[i] == "-s":
			i = i + 1
			method = int(argv[i])
			if method not in [0,1]:
				print("Unknown selection method {0}".format(method))
				exit_with_help(argv)
		i = i + 1

	dataset = argv[i]
	subset_size = int(argv[i+1])
	if i+2 < argc:
		subset_file = open(argv[i+2],'w')
	if i+3 < argc:
		rest_file = open(argv[i+3],'w')

	return dataset, subset_size, method, subset_file, rest_file

def random_selection(dataset, subset_size):
	l = sum(1 for line in open(dataset,'r'))
	return sorted(random.sample(xrange(l), subset_size))

def stratified_selection(dataset, subset_size):
	labels = [line.split(None,1)[0] for line in open(dataset)]
	label_linenums = defaultdict(list)
	for i, label in enumerate(labels):
		label_linenums[label] += [i]

	l = len(labels)
	remaining = subset_size
	ret = []

	# classes with fewer data are sampled first; otherwise
	# some rare classes may not be selected
	for label in sorted(label_linenums, key=lambda x: len(label_linenums[x])):
		linenums = label_linenums[label]
		label_size = len(linenums)
		# at least one instance per class
		s = int(min(remaining, max(1, math.ceil(label_size*(float(subset_size)/l)))))
		if s == 0:
			sys.stderr.write('''\
Error: failed to have at least one instance per class
    1. You may have regression data.
    2. Your classification data is unbalanced or too small.
Please use -s 1.
''')
			sys.exit(-1)
		remaining -= s
		ret += [linenums[i] for i in random.sample(xrange(label_size), s)]
	return sorted(ret)

def main(argv=sys.argv):
	dataset, subset_size, method, subset_file, rest_file = process_options(argv)
	#uncomment the following line to fix the random seed
	#random.seed(0)
	selected_lines = []

	if method == 0:
		selected_lines = stratified_selection(dataset, subset_size)
	elif method == 1:
		selected_lines = random_selection(dataset, subset_size)

	#select instances based on selected_lines
	dataset = open(dataset,'r')
	prev_selected_linenum = -1
	for i in xrange(len(selected_lines)):
		for cnt in xrange(selected_lines[i]-prev_selected_linenum-1):
			line = dataset.readline()
			if rest_file:
				rest_file.write(line)
		subset_file.write(dataset.readline())
		prev_selected_linenum = selected_lines[i]
	subset_file.close()

	if rest_file:
		for line in dataset:
			rest_file.write(line)
		rest_file.close()
	dataset.close()

if __name__ == '__main__':
	main(sys.argv)