Merge pull request #75 from PyDataBlog/improved-distance

Support for distance metrics

Author: PyDataBlog

Project.toml

@@ -8,12 +8,14 @@ Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
 MLJModelInterface = "e80e1ace-859a-464e-9ed9-23947d8ae3ea"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
+UnsafeArrays = "c4a57d5a-5b31-53a6-b365-19f8c011fbd6"
 
 [compat]
 Distances = "0.8.2"
 MLJModelInterface = "0.2.1"
 StatsBase = "0.32, 0.33"
 julia = "1.3"
+UnsafeArrays = "1"
 
 [extras]
 MLJBase = "a7f614a8-145f-11e9-1d2a-a57a1082229d"

README.md

@@ -5,12 +5,13 @@
 [![Build Status](https://www.travis-ci.org/PyDataBlog/ParallelKMeans.jl.svg?branch=master)](https://www.travis-ci.org/PyDataBlog/ParallelKMeans.jl)
 [![Coverage Status](https://coveralls.io/repos/github/PyDataBlog/ParallelKMeans.jl/badge.svg?branch=master)](https://coveralls.io/github/PyDataBlog/ParallelKMeans.jl?branch=master)
 [![FOSSA Status](https://app.fossa.com/api/projects/git%2Bgithub.com%2FPyDataBlog%2FParallelKMeans.jl.svg?type=shield)](https://app.fossa.com/projects/git%2Bgithub.com%2FPyDataBlog%2FParallelKMeans.jl?ref=badge_shield)
+[![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/PyDataBlog/ParallelKMeans.jl/master)
 _________________________________________________________________________________________________________
 **Authors:** [Bernard Brenyah](https://www.linkedin.com/in/bbrenyah/) & [Andrey Oskin](https://www.linkedin.com/in/andrej-oskin-b2b03959/)
 _________________________________________________________________________________________________________
 
 <div align="center">
-    <b>Classic & Contemporary Variants Of K-Means In Sonic Mode<b>
+    <b>Classic & Contemporary Variants Of K-Means In Sonic Mode</b>
 </div>
 
 <p align="center">
@@ -64,7 +65,9 @@ ________________________________________________________________________________
 - Lightening fast implementation of K-Means clustering algorithm even on a single thread in native Julia.
 - Support for multi-theading implementation of K-Means clustering algorithm.
 - Kmeans++ initialization for faster and better convergence.
-- Implementation of all the variants of the K-Means algorithm.
+- Implementation of all available variants of the K-Means algorithm.
+- Support for all distance metrics available at [Distances.jl](https://github.com/JuliaStats/Distances.jl)
+- Supported interface as an [MLJ](https://github.com/alan-turing-institute/MLJ.jl#available-models) model.
 
 _________________________________________________________________________________________________________
 

benchmark/bench01_distance.jl

@@ -9,13 +9,13 @@ suite = BenchmarkGroup()
 Random.seed!(2020)
 X = rand(3, 100_000)
 centroids = rand(3, 2)
-d = Vector{Float64}(undef, 100_000)
-suite["100kx3"] = @benchmarkable ParallelKMeans.chunk_colwise($d, $X, $centroids, 1, nothing, 1:100_000, 1)
+d = fill(-Inf, 100_000)
+suite["100kx3"] = @benchmarkable ParallelKMeans.chunk_colwise(d1, $X, $centroids, 1, nothing, Euclidean(), 1:100_000, 1) setup=(d1 = copy(d))
 
 X = rand(10, 100_000)
 centroids = rand(10, 2)
-d = Vector{Float64}(undef, 100_000)
-suite["100kx10"] = @benchmarkable ParallelKMeans.chunk_colwise($d, $X, $centroids, 1, nothing, 1:100_000, 1)
+d = fill(-Inf, 100_000)
+suite["100kx10"] = @benchmarkable ParallelKMeans.chunk_colwise(d1, $X, $centroids, 1, nothing, Euclidean(), 1:100_000, 1) setup=(d1 = copy(d))
 
 end # module
 

docs/src/index.md

@@ -51,13 +51,15 @@ pkg> add ParallelKMeans
 The few (and selected) brave ones can simply grab the current experimental features by simply adding the experimental branch to your development environment after invoking the package manager with `]`:
 
 ```julia
-dev git@github.com:PyDataBlog/ParallelKMeans.jl.git
+pkg> add ParallelKMeans#experimental
 ```
 
-Don't forget to checkout the experimental branch and you are good to go with bleeding edge features and breakages!
+You are good to go with bleeding edge features and breakages!
 
-```bash
-git checkout experimental
+To revert to a stable version, you can simply run:
+
+```julia
+pkg> free ParallelKMeans
 ```
 
 ## Features
@@ -207,7 +209,7 @@ ________________________________________________________________________________
 - 0.1.4 Bug fixes.
 - 0.1.5 Added `Yinyang` algorithm.
 - 0.1.6 Added support for weighted k-means; Added `Coreset` algorithm; improved support for different types of the design matrix.
-- 0.1.7 Added `Yinyang` and `Coreset` support in MLJ interface; added `weights` support in MLJ; added RNG seed support in MLJ interface and through all algorithms.
+- 0.1.7 Added `Yinyang` and `Coreset` support in MLJ interface; added `weights` support in MLJ; added RNG seed support in MLJ interface and through all algorithms; added metric support.
 
 ## Contributing
 

src/ParallelKMeans.jl

@@ -2,9 +2,11 @@ module ParallelKMeans
 
 using StatsBase
 using Random
+using UnsafeArrays
+using Distances
 import MLJModelInterface
 import Base.Threads: @spawn
-import Distances
+
 
 include("kmeans.jl")
 include("seeding.jl")

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 = 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/elkan.jl

@@ -18,16 +18,18 @@ kmeans(Elkan(), X, 3) # 3 clusters, Elkan algorithm
 """
 struct Elkan <: AbstractKMeansAlg end
 
-function kmeans!(alg::Elkan, containers, X, k, weights;
+
+function kmeans!(alg::Elkan, containers, X, k, weights=nothing, metric=Euclidean();
                 n_threads = Threads.nthreads(),
                 k_init = "k-means++", max_iters = 300,
                 tol = eltype(X)(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)
 
-    update_containers(alg, containers, centroids, n_threads)
-    @parallelize n_threads ncol chunk_initialize(alg, containers, centroids, X, weights)
+    update_containers(alg, containers, centroids, n_threads, metric)
+    @parallelize n_threads ncol chunk_initialize(alg, containers, centroids, X, weights, metric)
 
     T = eltype(X)
     converged = false
@@ -38,7 +40,7 @@ function kmeans!(alg::Elkan, containers, X, k, weights;
     while niters < max_iters
         niters += 1
         # Core iteration
-        @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 distributed containers (such as centroids_new, centroids_cnt)
         # in paper it is step 4
@@ -47,10 +49,10 @@ function kmeans!(alg::Elkan, containers, X, k, weights;
         J = sum(containers.ub)
 
         # auxiliary calculation, in paper it's d(c, m(c))
-        calculate_centroids_movement(alg, containers, centroids)
+        calculate_centroids_movement(alg, containers, centroids, metric)
 
         # lower and ounds update, in paper it's steps 5 and 6
-        @parallelize n_threads ncol chunk_update_bounds(alg, containers, centroids)
+        @parallelize n_threads ncol chunk_update_bounds(alg, containers, centroids, metric)
 
         # Step 7, final assignment of new centroids
         centroids .= containers.centroids_new[end]
@@ -67,11 +69,11 @@ function kmeans!(alg::Elkan, containers, X, k, weights;
         end
 
         # Step 1 in original paper, calulation of distance d(c, c')
-        update_containers(alg, containers, centroids, n_threads)
+        update_containers(alg, containers, centroids, n_threads, metric)
         J_previous = J
     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
@@ -85,6 +87,7 @@ function kmeans!(alg::Elkan, containers, X, k, weights;
     return KmeansResult(centroids, containers.labels, T[], Int[], T[], totalcost, niters, converged)
 end
 
+
 function create_containers(alg::Elkan, X, k, nrow, ncol, n_threads)
     T = eltype(X)
     lng = n_threads + 1
@@ -128,7 +131,8 @@ function create_containers(alg::Elkan, X, k, nrow, ncol, n_threads)
     )
 end
 
-function chunk_initialize(::Elkan, containers, centroids, X, weights, r, idx)
+
+function chunk_initialize(::Elkan, containers, centroids, X, weights, metric, r, idx)
     ub = containers.ub
     lb = containers.lb
     centroids_dist = containers.centroids_dist
@@ -138,15 +142,15 @@ function chunk_initialize(::Elkan, containers, centroids, X, weights, r, idx)
     T = eltype(X)
 
     @inbounds for i in r
-        min_dist = distance(X, centroids, i, 1)
+        min_dist = distance(metric, X, centroids, i, 1)
         label = 1
         lb[label, i] = min_dist
         for j in 2:size(centroids, 2)
             # triangular inequality
             if centroids_dist[j, label] > min_dist
                 lb[j, i] = min_dist
             else
-                dist = distance(X, centroids, i, j)
+                dist = distance(metric, X, centroids, i, j)
                 label = dist < min_dist ? j : label
                 min_dist = dist < min_dist ? dist : min_dist
                 lb[j, i] = dist
@@ -161,7 +165,8 @@ function chunk_initialize(::Elkan, containers, centroids, X, weights, r, idx)
     end
 end
 
-function update_containers(::Elkan, containers, centroids, n_threads)
+
+function update_containers(::Elkan, containers, centroids, n_threads, metric)
     # unpack containers for easier manipulations
     centroids_dist = containers.centroids_dist
     T = eltype(centroids)
@@ -170,7 +175,7 @@ function update_containers(::Elkan, containers, centroids, n_threads)
     @inbounds for j in axes(centroids_dist, 2)
         min_dist = T(Inf)
         for i in j + 1:k
-            d = distance(centroids, centroids, i, j)
+            d = distance(metric, centroids, centroids, i, j)
             centroids_dist[i, j] = d
             centroids_dist[j, i] = d
             min_dist = min_dist < d ? min_dist : d
@@ -189,7 +194,8 @@ function update_containers(::Elkan, containers, centroids, n_threads)
     return centroids_dist
 end
 
-function chunk_update_centroids(::Elkan, containers, centroids, X, weights, r, idx)
+
+function chunk_update_centroids(::Elkan, containers, centroids, X, weights, metric, r, idx)
     # unpack
     ub = containers.ub
     lb = containers.lb
@@ -214,14 +220,14 @@ function chunk_update_centroids(::Elkan, containers, centroids, X, weights, r, i
 
             # one calculation per iteration is enough
             if stale[i]
-                min_dist = distance(X, centroids, i, label)
+                min_dist = distance(metric, X, centroids, i, label)
                 lb[label, i] = min_dist
                 ub[i] = min_dist
                 stale[i] = false
             end
 
             if (min_dist > lb[j, i]) | (min_dist > centroids_dist[j, label])
-                dist = distance(X, centroids, i, j)
+                dist = distance(metric, X, centroids, i, j)
                 lb[j, i] = dist
                 if dist < min_dist
                     min_dist = dist
@@ -242,16 +248,18 @@ function chunk_update_centroids(::Elkan, containers, centroids, X, weights, r, i
     end
 end
 
-function calculate_centroids_movement(alg::Elkan, containers, centroids)
+
+function calculate_centroids_movement(alg::Elkan, containers, centroids, metric)
     p = containers.p
     centroids_new = containers.centroids_new[end]
 
     for i in axes(centroids, 2)
-        p[i] = distance(centroids, centroids_new, i, i)
+        p[i] = distance(metric, centroids, centroids_new, i, i)
     end
 end
 
-function chunk_update_bounds(alg, containers, centroids, r, idx)
+
+function chunk_update_bounds(alg, containers, centroids, metric::Euclidean, r, idx)
     p = containers.p
     lb = containers.lb
     ub = containers.ub
@@ -267,3 +275,22 @@ function chunk_update_bounds(alg, containers, centroids, r, idx)
         ub[i] += p[labels[i]] + T(2)*sqrt(abs(ub[i]*p[labels[i]]))
     end
 end
+
+
+function chunk_update_bounds(alg, containers, centroids, metric::Metric, r, idx)
+    p = containers.p
+    lb = containers.lb
+    ub = containers.ub
+    stale = containers.stale
+    labels = containers.labels
+    T = eltype(centroids)
+
+    @inbounds for i in r
+        for j in axes(centroids, 2)
+            lb[j, i] -= p[j]
+        end
+        stale[i] = true
+        ub[i] += p[labels[i]]
+    end
+
+end