Metric support for Coreset & Yinyang

Author: PyDataBlog

src/coreset.jl

@@ -35,18 +35,20 @@ Coreset(; m = 100, alg = Lloyd()) = Coreset(m, alg)
 Coreset(m::Int) = Coreset(m, Lloyd())
 Coreset(alg::AbstractKMeansAlg) = Coreset(100, alg)
 
-function kmeans!(alg::Coreset, containers, X, k, weights;
+function kmeans!(alg::Coreset, containers, X, k, weights, metric=Euclidean();
                 n_threads = Threads.nthreads(),
                 k_init = "k-means++", max_iters = 300,
                 tol = eltype(design_matrix)(1e-6), verbose = false,
                 init = nothing, rng = Random.GLOBAL_RNG)
+
     nrow, ncol = size(X)
+
     centroids = isnothing(init) ? smart_init(X, k, n_threads, weights, rng, init=k_init).centroids : deepcopy(init)
 
     T = eltype(X)
     # Steps 2-4 of the paper's algorithm 3
     # We distribute points over the centers and calculate weights of each cluster
-    @parallelize n_threads ncol chunk_fit(alg, containers, centroids, X, weights)
+    @parallelize n_threads ncol chunk_fit(alg, containers, centroids, X, weights, metric)
 
     # after this step, containers.centroids_new
     collect_containers(alg, containers, n_threads)
@@ -62,15 +64,16 @@ function kmeans!(alg::Coreset, containers, X, k, weights;
 
     # run usual kmeans for new set with new weights.
     res = kmeans(alg.alg, coreset, k, weights = coreset_weights, tol = tol, max_iters = max_iters,
-        verbose = verbose, init = centroids, n_threads = n_threads, rng = rng)
+        verbose = verbose, init = centroids, n_threads = n_threads, rng = rng, metric = metric)
 
-    @parallelize n_threads ncol chunk_apply(alg, containers, res.centers, X, weights)
+    @parallelize n_threads ncol chunk_apply(alg, containers, res.centers, X, weights, metric)
 
     totalcost = sum(containers.totalcost)
 
     return KmeansResult(res.centers, containers.labels, T[], Int[], T[], totalcost, res.iterations, res.converged)
 end
 
+
 function create_containers(alg::Coreset, X, k, nrow, ncol, n_threads)
     T = eltype(X)
 
@@ -109,7 +112,8 @@ function create_containers(alg::Coreset, X, k, nrow, ncol, n_threads)
     )
 end
 
-function chunk_fit(alg::Coreset, containers, centroids, X, weights, r, idx)
+
+function chunk_fit(alg::Coreset, containers, centroids, X, weights, metric, r, idx)
     centroids_cnt = containers.centroids_cnt[idx]
     centroids_dist = containers.centroids_dist[idx]
     labels = containers.labels
@@ -118,10 +122,10 @@ function chunk_fit(alg::Coreset, containers, centroids, X, weights, r, idx)
 
     J = zero(T)
     for i in r
-        dist = distance(X, centroids, i, 1)
+        dist = distance(metric, X, centroids, i, 1)
         label = 1
         for j in 2:size(centroids, 2)
-            new_dist = distance(X, centroids, i, j)
+            new_dist = distance(metric, X, centroids, i, j)
 
             # calculation of the closest center (steps 2-3 of the paper's algorithm 3)
             label = new_dist < dist ? j : label
@@ -144,6 +148,7 @@ function chunk_fit(alg::Coreset, containers, centroids, X, weights, r, idx)
     containers.J[idx] = J
 end
 
+
 function collect_containers(::Coreset, containers, n_threads)
     # Here we transform formula of the step 6
     # By multiplying both sides of equation on $c_\phi / \alpha$ we obtain
@@ -172,6 +177,7 @@ function collect_containers(::Coreset, containers, n_threads)
     end
 end
 
+
 function chunk_update_sensitivity(alg::Coreset, containers, r, idx)
     labels = containers.labels
     centroids_const = containers.centroids_const
@@ -182,18 +188,19 @@ function chunk_update_sensitivity(alg::Coreset, containers, r, idx)
     end
 end
 
-function chunk_apply(alg::Coreset, containers, centroids, X, weights, r, idx)
+
+function chunk_apply(alg::Coreset, containers, centroids, X, weights, metric, r, idx)
     centroids_cnt = containers.centroids_cnt[idx]
     centroids_dist = containers.centroids_dist[idx]
     labels = containers.labels
     T = eltype(X)
 
     J = zero(T)
     for i in r
-        dist = distance(X, centroids, i, 1)
+        dist = distance(metric, X, centroids, i, 1)
         label = 1
         for j in 2:size(centroids, 2)
-            new_dist = distance(X, centroids, i, j)
+            new_dist = distance(metric, X, centroids, i, j)
 
             # calculation of the closest center (steps 2-3 of the paper's algorithm 3)
             label = new_dist < dist ? j : label

src/yinyang.jl

@@ -47,22 +47,24 @@ Yinyang(; group_size = 7, auto = true) = Yinyang(auto, group_size)
 阴阳(group_size::Int) = Yinyang(true, group_size)
 阴阳(; group_size = 7, auto = true) = Yinyang(auto, group_size)
 
-function kmeans!(alg::Yinyang, containers, X, k, weights;
+function kmeans!(alg::Yinyang, containers, X, k, weights, metric=Euclidean();
                 n_threads = Threads.nthreads(),
                 k_init = "k-means++", max_iters = 300,
                 tol = 1e-6, verbose = false,
                 init = nothing, rng = Random.GLOBAL_RNG)
+
     nrow, ncol = size(X)
+
     centroids = init == nothing ? smart_init(X, k, n_threads, weights, rng, init=k_init).centroids : deepcopy(init)
 
     # create initial groups of centers, step 1 in original paper
     initialize(alg, containers, centroids, rng, n_threads)
     # construct initial bounds, step 2
-    @parallelize n_threads ncol chunk_initialize(alg, containers, centroids, X, weights)
+    @parallelize n_threads ncol chunk_initialize(alg, containers, centroids, X, weights, metric)
     collect_containers(alg, containers, n_threads)
 
     # update centers and calculate drifts. Step 3.1 of the original paper.
-    calculate_centroids_movement(alg, containers, centroids)
+    calculate_centroids_movement(alg, containers, centroids, metric)
 
     T = eltype(X)
     converged = false
@@ -87,14 +89,14 @@ function kmeans!(alg::Yinyang, containers, X, k, weights;
 
         # push!(containers.debug, [0, 0, 0])
         # Core calculation of the Yinyang, 3.2-3.3 steps of the original paper
-        @parallelize n_threads ncol chunk_update_centroids(alg, containers, centroids, X, weights)
+        @parallelize n_threads ncol chunk_update_centroids(alg, containers, centroids, X, weights, metric)
         collect_containers(alg, containers, n_threads)
 
         # update centers and calculate drifts. Step 3.1 of the original paper.
-        calculate_centroids_movement(alg, containers, centroids)
+        calculate_centroids_movement(alg, containers, centroids, metric)
     end
 
-    @parallelize n_threads ncol sum_of_squares(containers, X, containers.labels, centroids, weights)
+    @parallelize n_threads ncol sum_of_squares(containers, X, containers.labels, centroids, weights, metric)
     totalcost = sum(containers.sum_of_squares)
 
     # Terminate algorithm with the assumption that K-means has converged
@@ -170,6 +172,7 @@ function create_containers(alg::Yinyang, X, k, nrow, ncol, n_threads)
     )
 end
 
+
 function initialize(alg::Yinyang, containers, centroids, rng, n_threads)
     groups = containers.groups
     indices = containers.indices
@@ -186,21 +189,23 @@ function initialize(alg::Yinyang, containers, centroids, rng, n_threads)
     end
 end
 
-function chunk_initialize(alg::Yinyang, containers, centroids, X, weights, r, idx)
+
+function chunk_initialize(alg::Yinyang, containers, centroids, X, weights, metric, r, idx)
     T = eltype(X)
     centroids_cnt = containers.centroids_cnt[idx]
     centroids_new = containers.centroids_new[idx]
 
     @inbounds for i in r
-        label = point_all_centers!(alg, containers, centroids, X, i)
+        label = point_all_centers!(alg, containers, centroids, X, i, metric)
         centroids_cnt[label] += isnothing(weights) ? one(T) : weights[i]
         for j in axes(X, 1)
             centroids_new[j, label] += isnothing(weights) ? X[j, i] : weights[i] * X[j, i]
         end
     end
 end
 
-function calculate_centroids_movement(alg::Yinyang, containers, centroids)
+
+function calculate_centroids_movement(alg::Yinyang, containers, centroids, metric)
     p = containers.p
     groups = containers.groups
     gd = containers.gd
@@ -210,7 +215,7 @@ function calculate_centroids_movement(alg::Yinyang, containers, centroids)
     @inbounds for (gi, ri) in enumerate(groups)
         max_drift = T(-Inf)
         for i in ri
-            p[i] = sqrt(distance(centroids, centroids_new, i, i))
+            p[i] = sqrt(distance(metric, centroids, centroids_new, i, i))
             max_drift = p[i] > max_drift ? p[i] : max_drift
 
             # Should do it more elegantly
@@ -222,7 +227,8 @@ function calculate_centroids_movement(alg::Yinyang, containers, centroids)
     end
 end
 
-function chunk_update_centroids(alg::Yinyang, containers, centroids, X, weights, r, idx)
+
+function chunk_update_centroids(alg::Yinyang, containers, centroids, X, weights, metric, r, idx)
     # unpack containers for easier manipulations
     centroids_new = containers.centroids_new[idx]
     centroids_cnt = containers.centroids_cnt[idx]
@@ -256,7 +262,7 @@ function chunk_update_centroids(alg::Yinyang, containers, centroids, X, weights,
 
         # tighten upper bound
         label = labels[i]
-        ubx = sqrt(distance(X, centroids, i, label))
+        ubx = sqrt(distance(metric, X, centroids, i, label))
         ub[i] = ubx
         ubx <= lbx && continue
 
@@ -275,7 +281,7 @@ function chunk_update_centroids(alg::Yinyang, containers, centroids, X, weights,
                 ((c == old_label) | (ubx < old_lb - p[c])) && continue
                 mask[c] = true
                 # containers.debug[end][2] += 1 # local filter update
-                dist = distance(X, centroids, i, c)
+                dist = distance(metric, X, centroids, i, c)
                 if dist < ubx2
                     new_lb2 = ubx2
                     ubx2 = dist
@@ -290,7 +296,7 @@ function chunk_update_centroids(alg::Yinyang, containers, centroids, X, weights,
                 mask[c] && continue
                 new_lb < old_lb - p[c] && continue
                 # containers.debug[end][3] += 1 # lower bound update
-                dist = distance(X, centroids, i, c)
+                dist = distance(metric, X, centroids, i, c)
                 if dist < new_lb2
                     new_lb2 = dist
                     new_lb = sqrt(new_lb2)
@@ -314,7 +320,7 @@ function chunk_update_centroids(alg::Yinyang, containers, centroids, X, weights,
                 ubx < old_lb - p[c] && continue
                 # containers.debug[end][2] += 1 # local filter update
                 mask[c] = true
-                dist = distance(X, centroids, i, c)
+                dist = distance(metric, X, centroids, i, c)
                 if dist < ubx2
                     # closest center was in previous cluster
                     if indices[label] != gi
@@ -336,7 +342,7 @@ function chunk_update_centroids(alg::Yinyang, containers, centroids, X, weights,
                 mask[c] && continue
                 new_lb < old_lb - p[c] && continue
                 # containers.debug[end][3] += 1 # lower bound update
-                dist = distance(X, centroids, i, c)
+                dist = distance(metric, X, centroids, i, c)
                 if dist < new_lb2
                     new_lb2 = dist
                     new_lb = sqrt(new_lb2)
@@ -360,12 +366,13 @@ function chunk_update_centroids(alg::Yinyang, containers, centroids, X, weights,
     end
 end
 
+
 """
     point_all_centers!(containers, centroids, X, i)
 
 Calculates new labels and upper and lower bounds for all points.
 """
-function point_all_centers!(alg::Yinyang, containers, centroids, X, i)
+function point_all_centers!(alg::Yinyang, containers, centroids, X, i, metric)
     ub = containers.ub
     lb = containers.lb
     labels = containers.labels
@@ -381,7 +388,7 @@ function point_all_centers!(alg::Yinyang, containers, centroids, X, i)
         group_min_distance2 = T(Inf)
         group_label = ri[1]
         for k in ri
-            dist = distance(X, centroids, i, k)
+            dist = distance(metric, X, centroids, i, k)
             if group_min_distance > dist
                 group_label = k
                 group_min_distance2 = group_min_distance
@@ -407,6 +414,7 @@ function point_all_centers!(alg::Yinyang, containers, centroids, X, i)
     return label
 end
 
+
 # I believe there should be oneliner for it
 function rangify(x)
     res = UnitRange{Int}[]

test/test05_hamerly.jl

@@ -7,6 +7,7 @@ using StableRNGs
 using Random
 using Distances
 
+
 @testset "initialize" begin
     X = permutedims([1.0 2; 2 1; 4 5; 6 6])
     centroids = permutedims([1.0 2; 4 5; 6 6])

test/test06_yinyang.jl

@@ -3,6 +3,8 @@ module TestYinyang
 using ParallelKMeans
 using Test
 using StableRNGs
+using Distances
+
 
 @testset "basic kmeans Yinyang" begin
     X = [1. 2. 4.;]
@@ -196,4 +198,29 @@ end
     @test !alg.auto
 end
 
+@testset "Yinyang metric support" begin
+    rng = StableRNG(2020)
+    X = [1. 2. 4.;]
+
+    res = kmeans(Yinyang(), X, 2; tol = 1e-16, metric=Cityblock(), rng = rng)
+
+    @test res.assignments == [2, 2, 1]
+    @test res.centers == [4.0 1.5]
+    @test res.totalcost == 1.0
+    @test res.converged
+
+    rng = StableRNG(2020)
+    X = rand(3, 100)
+    rng_orig = deepcopy(rng)
+
+    baseline = kmeans(Lloyd(), X, 2, tol = 1e-16, metric=Cityblock(), rng = rng)
+
+    rng = deepcopy(rng_orig)
+    res = kmeans(Yinyang(), X, 2; tol = 1e-16, metric=Cityblock(), rng = rng)
+
+    @test res.totalcost ≈ baseline.totalcost
+    @test res.converged == baseline.converged
+    @test res.iterations == baseline.iterations
+end
+
 end # module

test/test07_coreset.jl

@@ -3,6 +3,8 @@ module TestCoreset
 using ParallelKMeans
 using Test
 using StableRNGs
+using Distances
+
 
 @testset "basic coresets" begin
     rng = StableRNG(2020)
@@ -45,4 +47,16 @@ end
     @test alg.alg == Hamerly()
 end
 
+@testset "Coreset metric support" begin
+    rng = StableRNG(2020)
+    X = [1. 2. 4.;]
+
+    res = kmeans(Coreset(), X, 2; tol = 1e-16, metric=Cityblock(), rng = rng)
+
+    @test res.assignments == [2, 2, 1]
+    @test res.centers == [4.0 1.4865168535972686]
+    @test res.totalcost == 1.0
+    @test res.converged
+end
+
 end # module