|
| 1 | +""" |
| 2 | + Coreset() |
| 3 | +
|
| 4 | +Coreset algorithm implementation, based on "Lucic, Mario & Bachem, |
| 5 | +Olivier & Krause, Andreas. (2015). Strong Coresets for Hard and Soft Bregman |
| 6 | +Clustering with Applications to Exponential Family Mixtures. " |
| 7 | +
|
| 8 | +It can be used directly in `kmeans` function |
| 9 | +
|
| 10 | +```julia |
| 11 | +X = rand(30, 100_000) # 100_000 random points in 30 dimensions |
| 12 | +
|
| 13 | +kmeans(Coreset(), X, 3) # 3 clusters, Coreset algorithm |
| 14 | +``` |
| 15 | +""" |
| 16 | +struct Coreset{T <: AbstractKMeansAlg} <: AbstractKMeansAlg |
| 17 | + m::Int |
| 18 | + alg::T |
| 19 | +end |
| 20 | + |
| 21 | +Coreset() = Coreset(100, Lloyd()) |
| 22 | + |
| 23 | +function kmeans!(alg::Coreset, containers, X, k, weights; |
| 24 | + n_threads = Threads.nthreads(), |
| 25 | + k_init = "k-means++", max_iters = 300, |
| 26 | + tol = eltype(design_matrix)(1e-6), verbose = false, init = nothing) |
| 27 | + nrow, ncol = size(X) |
| 28 | + centroids = isnothing(init) ? smart_init(X, k, n_threads, init=k_init).centroids : deepcopy(init) |
| 29 | + |
| 30 | + T = eltype(X) |
| 31 | + # Steps 2-4 of the paper's algorithm 3 |
| 32 | + # We distribute points over the centers and calculate weights of each cluster |
| 33 | + @parallelize n_threads ncol chunk_fit(alg, containers, centroids, X, weights) |
| 34 | + |
| 35 | + # after this step, containers.centroids_new |
| 36 | + collect_containers(alg, containers, n_threads) |
| 37 | + |
| 38 | + # step 7 of the algorithm 3 |
| 39 | + @parallelize n_threads ncol chunk_update_sensitivity(alg, containers) |
| 40 | + |
| 41 | + # sample from containers.s |
| 42 | + coreset_ids = wsample(1:ncol, containers.s, alg.m) |
| 43 | + coreset = X[:, coreset_ids] |
| 44 | + # create new weights as 1/s[i] |
| 45 | + coreset_weights = one(T) ./ @view containers.s[coreset_ids] |
| 46 | + |
| 47 | + # run usual kmeans for new set with new weights. |
| 48 | + res = kmeans(alg.alg, coreset, k, coreset_weights, tol = tol, max_iters = max_iters, |
| 49 | + verbose = verbose, init = centroids, n_threads = n_threads) |
| 50 | + |
| 51 | + @parallelize n_threads ncol chunk_apply(alg, containers, res.centers, X, weights) |
| 52 | + |
| 53 | + totalcost = sum(containers.totalcost) |
| 54 | + |
| 55 | + return KmeansResult(res.centers, containers.labels, T[], Int[], T[], totalcost, res.iterations, res.converged) |
| 56 | +end |
| 57 | + |
| 58 | +function create_containers(alg::Coreset, X, k, nrow, ncol, n_threads) |
| 59 | + T = eltype(X) |
| 60 | + |
| 61 | + centroids_cnt = Vector{Vector{T}}(undef, n_threads) |
| 62 | + centroids_dist = Vector{Vector{T}}(undef, n_threads) |
| 63 | + |
| 64 | + # sensitivity |
| 65 | + |
| 66 | + for i in 1:n_threads |
| 67 | + centroids_cnt[i] = zeros(T, k) |
| 68 | + centroids_dist[i] = zeros(T, k) |
| 69 | + end |
| 70 | + |
| 71 | + labels = Vector{Int}(undef, ncol) |
| 72 | + s = Vector{T}(undef, ncol) |
| 73 | + |
| 74 | + # J is the same as $c_\phi$ in the paper. |
| 75 | + J = Vector{T}(undef, n_threads) |
| 76 | + |
| 77 | + alpha = 16 * (log(k) + 2) |
| 78 | + |
| 79 | + centroids_const = Vector{T}(undef, k) |
| 80 | + |
| 81 | + # total_sum_calculation |
| 82 | + totalcost = Vector{T}(undef, n_threads) |
| 83 | + |
| 84 | + return ( |
| 85 | + centroids_cnt = centroids_cnt, |
| 86 | + centroids_dist = centroids_dist, |
| 87 | + s = s, |
| 88 | + labels = labels, |
| 89 | + totalcost = totalcost, |
| 90 | + J = J, |
| 91 | + centroids_const = centroids_const, |
| 92 | + alpha = alpha |
| 93 | + ) |
| 94 | +end |
| 95 | + |
| 96 | +function chunk_fit(alg::Coreset, containers, centroids, X, weights, r, idx) |
| 97 | + centroids_cnt = containers.centroids_cnt[idx] |
| 98 | + centroids_dist = containers.centroids_dist[idx] |
| 99 | + labels = containers.labels |
| 100 | + s = containers.s |
| 101 | + T = eltype(X) |
| 102 | + |
| 103 | + J = zero(T) |
| 104 | + for i in r |
| 105 | + dist = distance(X, centroids, i, 1) |
| 106 | + label = 1 |
| 107 | + for j in 2:size(centroids, 2) |
| 108 | + new_dist = distance(X, centroids, i, j) |
| 109 | + |
| 110 | + # calculation of the closest center (steps 2-3 of the paper's algorithm 3) |
| 111 | + label = new_dist < dist ? j : label |
| 112 | + dist = new_dist < dist ? new_dist : dist |
| 113 | + end |
| 114 | + labels[i] = label |
| 115 | + |
| 116 | + # calculation of the $c_\phi$ (step 4) |
| 117 | + # Note: $d_A(x', B) = min_{b \in B} d_A(x', b)$ |
| 118 | + # Not exactly sure about whole `weights` thing, needs further investigation |
| 119 | + # for Nothing `weights` (default) it'll work as intendent |
| 120 | + centroids_cnt[label] += isnothing(weights) ? one(T) : weights[i] |
| 121 | + centroids_dist[label] += isnothing(weights) ? dist : weights[i] * dist |
| 122 | + J += dist |
| 123 | + |
| 124 | + # for now we write dist to sensitivity, update it later |
| 125 | + s[i] = dist |
| 126 | + end |
| 127 | + |
| 128 | + containers.J[idx] = J |
| 129 | +end |
| 130 | + |
| 131 | +function collect_containers(::Coreset, containers, n_threads) |
| 132 | + # Here we transform formula of the step 6 |
| 133 | + # By multiplying both sides of equation on $c_\phi / \alpha$ we obtain |
| 134 | + # $s(x) <- d_A(x, B) + 2 \sum d_A(x, B) / |B_i| + 4 c_\phi |\Chi| / (|B_i| * \alpha)$ |
| 135 | + # Taking into account that $c_\phi = 1/|\Chi| \sum d_A(x', B) = J / |\Chi|$ we get |
| 136 | + # $s(x) <- d_A(x, B) + 2 \sum d_A(x, B) / |B_i| + 4 J / \alpha * 1/ |B_i|$ |
| 137 | + |
| 138 | + alpha = containers.alpha |
| 139 | + centroids_const = containers.centroids_const |
| 140 | + |
| 141 | + centroids_cnt = containers.centroids_cnt[1] |
| 142 | + centroids_dist = containers.centroids_dist[1] |
| 143 | + J = containers.J[1] |
| 144 | + |
| 145 | + @inbounds for i in 2:n_threads |
| 146 | + centroids_cnt .+= containers.centroids_cnt[i] |
| 147 | + centroids_dist .+= containers.centroids_dist[i] |
| 148 | + J += containers.J[i] |
| 149 | + end |
| 150 | + |
| 151 | + J = 4 * J / alpha |
| 152 | + |
| 153 | + for i in 1:length(centroids_dist) |
| 154 | + centroids_const[i] = 2 * centroids_dist[i] / centroids_cnt[i] + |
| 155 | + J / centroids_cnt[i] |
| 156 | + end |
| 157 | +end |
| 158 | + |
| 159 | +function chunk_update_sensitivity(alg::Coreset, containers, r, idx) |
| 160 | + labels = containers.labels |
| 161 | + centroids_const = containers.centroids_const |
| 162 | + s = containers.s |
| 163 | + |
| 164 | + @inbounds for i in r |
| 165 | + s[i] += centroids_const[labels[i]] |
| 166 | + end |
| 167 | +end |
| 168 | + |
| 169 | +function chunk_apply(alg::Coreset, containers, centroids, X, weights, r, idx) |
| 170 | + centroids_cnt = containers.centroids_cnt[idx] |
| 171 | + centroids_dist = containers.centroids_dist[idx] |
| 172 | + labels = containers.labels |
| 173 | + T = eltype(X) |
| 174 | + |
| 175 | + J = zero(T) |
| 176 | + for i in r |
| 177 | + dist = distance(X, centroids, i, 1) |
| 178 | + label = 1 |
| 179 | + for j in 2:size(centroids, 2) |
| 180 | + new_dist = distance(X, centroids, i, j) |
| 181 | + |
| 182 | + # calculation of the closest center (steps 2-3 of the paper's algorithm 3) |
| 183 | + label = new_dist < dist ? j : label |
| 184 | + dist = new_dist < dist ? new_dist : dist |
| 185 | + end |
| 186 | + labels[i] = label |
| 187 | + J += isnothing(weights) ? dist : weights[i] * dist |
| 188 | + end |
| 189 | + |
| 190 | + containers.totalcost[idx] = J |
| 191 | +end |
0 commit comments