WIP: Verified metric support to lloyd

Author: PyDataBlog

src/kmeans.jl

@@ -132,7 +132,6 @@ A Float type representing the computed metric is returned.
 """
 function sum_of_squares(containers, x, labels, centre, weights, metric, r, idx)
     s = zero(eltype(x))
-
     @inbounds for i in r
         s += isnothing(weights) ? distance(metric, x, centre, i, labels[i]) : weights[i] * distance(metric, x, centre, i, labels[i])
     end
@@ -181,7 +180,7 @@ function kmeans(alg::AbstractKMeansAlg, design_matrix, k;
 
     return kmeans!(alg, containers, design_matrix, k, weights, n_threads = n_threads,
                     k_init = k_init, max_iters = max_iters, tol = tol,
-                    verbose = verbose, init = init, rng = rng, metric = Euclidean())
+                    verbose = verbose, init = init, rng = rng, metric = metric)
 
 end
 

src/lloyd.jl

@@ -48,7 +48,6 @@ function kmeans!(alg::Lloyd, containers, X, k, weights;
         J_previous = J
         niters += 1
     end
-
     @parallelize n_threads ncol sum_of_squares(containers, X, containers.labels, centroids, weights, metric)
     totalcost = sum(containers.sum_of_squares)
 

test/test03_lloyd.jl

@@ -4,6 +4,8 @@ using ParallelKMeans
 using Test
 using StableRNGs
 using StatsBase
+using Distances
+
 
 @testset "basic kmeans" begin
     X = [1. 2. 4.;]
@@ -101,4 +103,23 @@ end
     @test res.iterations == 6
 end
 
+@testset "Lloyd metric support" begin
+    Random.seed!(2020)
+    X = [1. 2. 4.;]
+
+    res = kmeans(Lloyd(), X, 2; tol = 1e-16, metric=Cityblock())
+
+    @test res.assignments == [1, 1, 2]
+    @test res.centers == [1.5 4.0]
+    @test res.totalcost == 1.0
+    @test res.converged
+
+    Random.seed!(2020)
+    X = rand(3, 100)
+
+    res = kmeans(X, 2, tol = 1e-16, metric=Cityblock())
+    @test res.totalcost ≈ 62.040452528953736
+    @test res.converged
+end
+
 end # module