Use KDTree beyond nmax in KMedoids when possible

Project.toml

@@ -1,7 +1,7 @@
 name = "TableTransforms"
 uuid = "0d432bfd-3ee1-4ac1-886a-39f05cc69a3e"
-authors = ["Júlio Hoffimann <julio.hoffimann@gmail.com> and contributors"]
 version = "1.38.9"
+authors = ["Júlio Hoffimann <julio.hoffimann@gmail.com> and contributors"]
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
@@ -13,6 +13,7 @@ Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 InverseFunctions = "3587e190-3f89-42d0-90ee-14403ec27112"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+NearestNeighbors = "b8a86587-4115-5ab1-83bc-aa920d37bbce"
 NelderMead = "2f6b4ddb-b4ff-44c0-b59b-2ab99302f970"
 PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -34,6 +35,7 @@ Distributed = "1.10"
 Distributions = "0.25"
 InverseFunctions = "0.1"
 LinearAlgebra = "1.10"
+NearestNeighbors = "0.4"
 NelderMead = "0.4"
 PrettyTables = "3.0"
 Random = "1.10"

src/TableTransforms.jl

@@ -26,6 +26,7 @@ using InverseFunctions: NoInverse, inverse as invfun
 using StatsBase: AbstractWeights, Weights, sample
 using Distributed: CachingPool, pmap, workers
 using SpectralIndices: indices, bands, compute
+using NearestNeighbors: KDTree, nn
 using NelderMead: optimise
 
 import Distributions: quantile, cdf

src/transforms/kmedoids.jl

@@ -129,14 +129,9 @@ end
 
 function _updatelabels!(td, table, medoids, labels, dists)
   for (k, mₖ) in enumerate(medoids)
-    inds = 1:_nrows(table)
-
-    X = Tables.subset(table, inds, viewhint=true)
     μ = Tables.subset(table, [mₖ], viewhint=true)
-
-    δ = pairwise(td, X, μ)
-
-    @inbounds for i in inds
+    δ = pairwise(td, table, μ) |> vec
+    @inbounds for i in eachindex(δ)
       if δ[i] < dists[i]
         dists[i] = δ[i]
         labels[i] = k
@@ -152,33 +147,50 @@ function _updatemedoids!(td, table, medoids, labels)
     # if cluster is empty, then no medoid to update
     isempty(inds) && continue
 
-    X = Tables.subset(table, inds, viewhint=true)
-
-    j = _medoid(td, X)
+    # compute new medoid for cluster k
+    clust = Tables.subset(table, inds, viewhint=true)
+    _, j = findmin(sum, eachcol(pairwise(td, clust)))
 
     @inbounds medoids[k] = inds[j]
   end
 end
 
-function _medoid(td, table)
-  Δ = pairwise(td, table)
-  _, j = findmin(sum, eachcol(Δ))
-  j
-end
-
 function _interp(labels, inds, table, td)
   nobs = _nrows(table)
 
   ilabels = fill(0, nobs)
   ilabels[inds] .= labels
 
   X = Tables.subset(table, inds, viewhint=true)
+  s = _searcher(X, td)
+
   for i in setdiff(1:nobs, inds)
     x = Tables.subset(table, [i], viewhint=true)
-    δ = pairwise(td, X, x)
-    _, j = findmin(δ)
+    j = _search(s, x)
     ilabels[i] = labels[j]
   end
 
   ilabels
 end
+
+function _searcher(X, td)
+  # check if all variables are continuous
+  cols = Tables.columns(X)
+  vars = Tables.columnnames(cols)
+  allcont = all(vars) do var
+    x = Tables.getcolumn(cols, var)
+    elscitype(x) <: Continuous
+  end
+
+  # use KDTree if all variables are continuous
+  if allcont
+    data = Tables.matrix(X)
+    KDTree(transpose(data))
+  else
+    X, td
+  end
+end
+
+_search(s::KDTree, x) = nn(s, Tables.matrix(x) |> vec) |> first
+
+_search((X, td), x) = pairwise(td, X, x) |> vec |> argmin

test/transforms/sample.jl

@@ -64,5 +64,5 @@
   t = (; x, y, c)
 
   T = Sample(10_000)
-  @test @elapsed(apply(T, t)) < 0.5
+  @test @elapsed(apply(T, t)) < 1.0
 end