from .correlate import correlate, CCG
from .convolve import convolve

@@ -62,8 +62,1446 @@
 onvolve%0A
+import torch%0Aimport numpy as np%0A%0A%0Adef get_n_vq(y, N=500, method='proportional'):%0A    nclu = y.long().bincount().shape%5B0%5D%0A    if method == 'proportional':%0A        dist = y.long().bincount()/float(y.shape%5B0%5D)%0A    elif method == 'equal':%0A        dist = torch.ones((nclu,)) / nclu%0A    n_vq = torch.round(dist*N)%0A    err = n_vq.sum() - N%0A    n_vq%5B0%5D -= err%0A    assert(n_vq.sum()==N)%0A    return n_vq.long()%0A  %0Adef VQ(X, y, N=500, method='proportional', n_vq=None):%0A    n_vq = get_n_vq(y, N)%0A    vq = %5B%5D%0A    nclu = y.long().bincount().shape%5B0%5D%0A    from sklearn.cluster import MiniBatchKMeans%0A    for _clu_id in range(nclu):%0A        km = MiniBatchKMeans(n_vq%5B_clu_id%5D)%0A        km.fit(X%5By==_clu_id%5D.numpy())%0A        vq.append(km.cluster_centers_)%0A    vq = np.vstack(vq)%0A    vq_labels = torch.repeat_interleave(torch.arange(nclu), n_vq)%0A    return vq, vq_labels%0A  %0A%0Aclass VQ_KNN(object):%0A    '''%0A    model = VQ_KNN(X,y,N=6,k=1)%0A    model.predict(X_test)%0A    '''%0A    def __init__(self, X, y, N=500, k=1, method='proportional', n_vq=None):%0A        if type(X) == np.ndarray:%0A            X = torch.from_numpy(X)%0A        if type(y) == np.ndarray:%0A            y = torch.from_numpy(y)%0A        self.vq, self.vq_label = VQ(X, y, N)%0A        from sklearn.neighbors import KNeighborsClassifier as KNN%0A        self.knn = KNN(n_neighbors=k)%0A        self.knn.fit(self.vq, self.vq_label)%0A    %0A    def predict(self, X_test):%0A        return self.knn.predict(X_test)%0A

Update __init__.py