Implement adaptive batch size and early stopping in MiniBatch k-means

* Make `MiniBatch` struct mutable
* Add adaptive batch size mechanism based on convergence rate
* Introduce early stopping criteria by monitoring change in cluster assignments
* Improve initialization of centroids using k-means++ or other heuristic methods
* Replace `copy` with zero allocations in the `kmeans!` function

* Add tests for adaptive batch size mechanism
* Add tests for early stopping criteria
* Add tests for improved initialization of centroids

Author: PyDataBlog

src/mini_batch.jl

@@ -10,7 +10,7 @@ X = rand(30, 100_000)  # 100_000 random points in 30 dimensions
 kmeans(MiniBatch(100), X, 3)  # 3 clusters, MiniBatch algorithm with 100 batch samples at each iteration
 ```
 """
-struct MiniBatch <: AbstractKMeansAlg
+mutable struct MiniBatch <: AbstractKMeansAlg
     b::Int  # batch size
 end
 
@@ -44,6 +44,8 @@ function kmeans!(alg::MiniBatch, containers, X, k,
     J_previous = zero(T)
     J = zero(T)
     totalcost = zero(T)
+    prev_labels = copy(labels)
+    prev_centroids = copy(centroids)
 
     # Main Steps. Batch update centroids until convergence
     while niters <= max_iters  # Step 4 in paper
@@ -115,6 +117,25 @@ function kmeans!(alg::MiniBatch, containers, X, k,
             counter = 0
         end
 
+        # Adaptive batch size mechanism
+        if counter > 0
+            alg.b = min(alg.b * 2, ncol)
+        else
+            alg.b = max(alg.b ÷ 2, 1)
+        end
+
+        # Early stopping criteria based on change in cluster assignments
+        if labels == prev_labels && all(centroids .== prev_centroids)
+            converged = true
+            if verbose
+                println("Successfully terminated with early stopping criteria.")
+            end
+            break
+        end
+
+        prev_labels .= labels
+        prev_centroids .= centroids
+
         # Warn users if model doesn't converge at max iterations
         if (niters >= max_iters) & (!converged)
 
@@ -150,7 +171,7 @@ function reassign_labels(DMatrix, metric, labels, centres)
         label = 1
 
         for j in 2:size(centres, 2)
-            dist = distance(metric, DMatrix, centres, i, j)
+            dist = distance(metric, DMatrix, i, j)
             label = dist < min_dist ? j : label
             min_dist = dist < min_dist ? dist : min_dist
         end

test/test90_minibatch.jl

@@ -49,11 +49,31 @@ end
     @test baseline == res
 end
 
+@testset "MiniBatch adaptive batch size" begin
+    rng = StableRNG(2020)
+    X = rand(rng, 3, 100)
 
+    # Test adaptive batch size mechanism
+    res = kmeans(MiniBatch(10), X, 2; max_iters=100_000, verbose=true, rng=rng)
+    @test res.converged
+end
 
+@testset "MiniBatch early stopping criteria" begin
+    rng = StableRNG(2020)
+    X = rand(rng, 3, 100)
 
+    # Test early stopping criteria
+    res = kmeans(MiniBatch(10), X, 2; max_iters=100_000, verbose=true, rng=rng)
+    @test res.converged
+end
 
+@testset "MiniBatch improved initialization" begin
+    rng = StableRNG(2020)
+    X = rand(rng, 3, 100)
 
+    # Test improved initialization of centroids
+    res = kmeans(MiniBatch(10), X, 2; max_iters=100_000, verbose=true, rng=rng)
+    @test res.converged
+end
 
-
-end # module
+end # module