Fix Kmeans (#115)

* Fix kmeans

Fix while loop convergence, kmeans++ init, tolerance value

* Update test_metrics.py

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update _kmeans.py

* Update test_metrics.py

* Update test_metrics.py

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: adamgayoso

src/scib_metrics/utils/_kmeans.py

@@ -12,10 +12,17 @@
 from ._utils import get_ndarray, validate_seed
 
 
+def _tolerance(X: jnp.ndarray, tol: float) -> float:
+    """Return a tolerance which is dependent on the dataset."""
+    variances = np.var(X, axis=0)
+    return np.mean(variances) * tol
+
+
 def _initialize_random(X: jnp.ndarray, n_clusters: int, key: jax.random.KeyArray) -> jnp.ndarray:
     """Initialize cluster centroids randomly."""
     n_obs = X.shape[0]
-    indices = jax.random.choice(key, n_obs, (n_clusters,), replace=False)
+    key, subkey = jax.random.split(key)
+    indices = jax.random.choice(subkey, n_obs, (n_clusters,), replace=False)
     initial_state = X[indices]
     return initial_state
 
@@ -53,13 +60,14 @@ def _step(state, _):
         return state, state["centroid"]
 
     _, centroids = jax.lax.scan(_step, initial_state, jnp.arange(n_clusters - 1))
+    centroids = jnp.concatenate([initial_centroid[jnp.newaxis, :], centroids])
     return centroids
 
 
 @jax.jit
 def _get_dist_labels(X: jnp.ndarray, centroids: jnp.ndarray) -> jnp.ndarray:
     """Get the distance and labels for each observation."""
-    dist = cdist(X, centroids)
+    dist = jnp.square(cdist(X, centroids))
     labels = jnp.argmin(dist, axis=1)
     return dist, labels
 
@@ -94,15 +102,15 @@ def __init__(
         self,
         n_clusters: int = 8,
         init: Literal["k-means++", "random"] = "k-means++",
-        n_init: int = 10,
+        n_init: int = 1,
         max_iter: int = 300,
         tol: float = 1e-4,
         seed: IntOrKey = 0,
     ):
         self.n_clusters = n_clusters
         self.n_init = n_init
         self.max_iter = max_iter
-        self.tol = tol
+        self.tol_scale = tol
         self.seed: jax.random.KeyArray = validate_seed(seed)
 
         if init not in ["k-means++", "random"]:
@@ -115,6 +123,7 @@ def __init__(
     def fit(self, X: np.ndarray):
         """Fit the model to the data."""
         X = check_array(X, dtype=np.float32, order="C")
+        self.tol = _tolerance(X, self.tol_scale)
         # Subtract mean for numerical accuracy
         mean = X.mean(axis=0)
         X -= mean
@@ -136,8 +145,7 @@ def _fit(self, X: np.ndarray):
     @partial(jax.jit, static_argnums=(0,))
     def _kmeans_full_run(self, X: jnp.ndarray, key: jnp.ndarray) -> jnp.ndarray:
         def _kmeans_step(state):
-            old_inertia = state[1]
-            centroids, _, _, n_iter = state
+            centroids, old_inertia, _, n_iter = state
             # TODO(adamgayoso): Efficiently compute argmin and min simultaneously.
             dist, new_labels = _get_dist_labels(X, centroids)
             # From https://colab.research.google.com/drive/1AwS4haUx6swF82w3nXr6QKhajdF8aSvA?usp=sharing
@@ -159,19 +167,22 @@ def _kmeans_step(state):
                 )
                 / counts
             )
-            new_inertia = jnp.mean(jnp.min(dist, axis=1))
+            new_inertia = jnp.sum(jnp.min(dist, axis=1))
             n_iter = n_iter + 1
             return new_centroids, new_inertia, old_inertia, n_iter
 
         def _kmeans_convergence(state):
             _, new_inertia, old_inertia, n_iter = state
-            cond1 = jnp.abs(old_inertia - new_inertia) < self.tol
-            cond2 = n_iter > self.max_iter
+            cond1 = jnp.abs(old_inertia - new_inertia) > self.tol
+            cond2 = n_iter < self.max_iter
             return jnp.logical_or(cond1, cond2)[0]
 
         centroids = self._initialize(X, self.n_clusters, key)
         # centroids, new_inertia, old_inertia, n_iter
         state = (centroids, jnp.inf, jnp.inf, jnp.array([0.0]))
-        state = _kmeans_step(state)
         state = jax.lax.while_loop(_kmeans_convergence, _kmeans_step, state)
-        return state[0], state[1]
+        # Compute final inertia
+        centroids = state[0]
+        dist, _ = _get_dist_labels(X, centroids)
+        final_intertia = jnp.sum(jnp.min(dist, axis=1))
+        return centroids, final_intertia

tests/test_metrics.py

@@ -7,7 +7,10 @@
 from scipy.sparse import csr_matrix
 from scipy.spatial.distance import cdist as sp_cdist
 from scipy.spatial.distance import pdist, squareform
+from sklearn.cluster import KMeans as SKMeans
+from sklearn.datasets import make_blobs
 from sklearn.metrics import silhouette_samples as sk_silhouette_samples
+from sklearn.metrics.pairwise import pairwise_distances_argmin
 from sklearn.neighbors import NearestNeighbors
 
 import scib_metrics
@@ -115,11 +118,30 @@ def test_isolated_labels():
 
 
 def test_kmeans():
-    X, _ = dummy_x_labels()
-    kmeans = scib_metrics.utils.KMeans(2)
+    centers = [[1, 1], [-1, -1], [1, -1]]
+    len(centers)
+    X, labels_true = make_blobs(n_samples=3000, centers=centers, cluster_std=0.7)
+    kmeans = scib_metrics.utils.KMeans(n_clusters=3)
     kmeans.fit(X)
     assert kmeans.labels_.shape == (X.shape[0],)
 
+    skmeans = SKMeans(n_clusters=3)
+    skmeans.fit(X)
+    sk_inertia = np.array([skmeans.inertia_])
+    jax_inertia = np.array([kmeans.inertia_])
+    np.testing.assert_allclose(sk_inertia, jax_inertia, atol=4e-2)
+
+    # Reorder cluster centroids between methods and measure accuracy
+    k_means_cluster_centers = kmeans.cluster_centroids_
+    order = pairwise_distances_argmin(kmeans.cluster_centroids_, skmeans.cluster_centers_)
+    sk_means_cluster_centers = skmeans.cluster_centers_[order]
+
+    k_means_labels = pairwise_distances_argmin(X, k_means_cluster_centers)
+    sk_means_labels = pairwise_distances_argmin(X, sk_means_cluster_centers)
+
+    accuracy = (k_means_labels == sk_means_labels).sum() / len(k_means_labels)
+    assert accuracy > 0.999
+
 
 def test_kbet():
     X, _, batch = dummy_x_labels_batch(x_is_neighbors_graph=True)