/
kahip_solver.py
54 lines (46 loc) · 1.55 KB
/
kahip_solver.py
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
import os, sys, math
import numpy as np
#from sklearn.cluster import MiniBatchKMeans, KMeans
import utils
import pickle
import time
#import json
#from collections import defaultdict
#import kahip
import os.path as osp
import pdb
data_dir = 'data'
#n_clusters = 64
km_method = 'km' #mbkm km
max_loyd = 10
'''
Implements solver fit and transform functionality.
'''
class KahipSolver():
def __init__(self):
#self.n_clusters = n_clusters
#kahip partition top level result
#64 for now!!
#-loads partition data and prepares to make predictions
self.kahip_path = osp.join(utils.data_dir, 'cache_partition64strong_0ht2')
classes_l = utils.load_lines(self.kahip_path) ##########
self.classes_l = [int(c) for c in classes_l]
'''
predict data
Input:
-dataset_idx: indices of data. List of ints.
Output:
- classes for index, as numpy array
'''
def predict(self, dataset_idx):
#needs to predict classes: d_cls_idx = solver.predict(dataset)
#should use numpy array for efficiency!
pred_ar = np.zeros(len(dataset_idx))
for i, idx in enumerate(dataset_idx):
pred_ar[i] = self.classes_l[idx]
return pred_ar
if __name__ == '__main__':
opt = utils.parse_args()
n_clusters = opt.n_clusters
n_clusters = 2
KahipSolver(n_clusters, opt)