update sbm test

martinroyer · Oct 18, 2017 · dc9efa0 · dc9efa0
1 parent 58120d2
commit dc9efa0
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git 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
@@ -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_)