diff --git a/pecok/pecok_clustering.py b/pecok/pecok_clustering.py index f99b55d..75a3757 100644 --- a/pecok/pecok_clustering.py +++ b/pecok/pecok_clustering.py @@ -15,3 +15,12 @@ def cluster(X, K): X : array-like or sparse matrix, shape=(n_samples, n_features) Training instances to cluster.""" return pecok_admm(X.dot(X.T)-Gamma_hat4(X), K) + +def cluster_sbm(A, K): + """Implementation of PECOK estimator of B* + + Parameters + ---------- + A : adjacency matrix for network, shape=(n_samples, n_samples) + Training instances to cluster.""" + return pecok_admm(A.dot(A), K) diff --git a/tests/test_sbm.py b/tests/test_sbm.py index 852118a..0bdd8c4 100644 --- a/tests/test_sbm.py +++ b/tests/test_sbm.py @@ -17,12 +17,12 @@ expA[:n_samples//2,:n_samples//2] += p-q expA[n_samples//2:,n_samples//2:] += p-q A = np.vectorize(lambda p : np.random.binomial(1, p))(expA) -A.flat[::n_samples+1] = 10 +A.flat[::n_samples+1] = 0 i_lower = np.tril_indices(n_samples, -1) A[i_lower] = A.T[i_lower] print(A) -Bhat = pecok_clustering.cluster(A, 2) +Bhat = pecok_clustering.cluster_sbm(A, 2) kMeans = cluster.KMeans(n_clusters=2, init='k-means++', n_init=100, copy_x=True) print("truth:".ljust(10), truth) print("pecok:".ljust(10), kMeans.fit(Bhat).labels_)