Merge pull request #5 from dillondaudert/local_connect

Implement local_connectivity
dillondaudert · Jan 28, 2019 · 40a5e56 · 40a5e56
2 parents 4f44708 + 9d0a7c5
commit 40a5e56
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 7 deletions.
diff --git a/src/umap_.jl b/src/umap_.jl
@@ -7,11 +7,13 @@ struct UMAP_{S}
     embedding::AbstractMatrix{S}
 
     function UMAP_(graph::AbstractMatrix{S}, embedding::AbstractMatrix{S}) where {S<:AbstractFloat}
-        issymmetric(graph) || throw(MethodError("UMAP_ constructor expected graph to be a symmetric matrix"))
+        issymmetric(graph) || isapprox(graph, graph') || error("UMAP_ constructor expected graph to be a symmetric matrix")
         new{S}(graph, embedding)
     end
 end
 
+const SMOOTH_K_TOLERANCE = 1e-5
+
 
 """
     umap(X::AbstractMatrix[, n_components=2]; <kwargs>) -> embedding
@@ -108,18 +110,33 @@ and the nearest neighbor (nn_dists) from each point.
 """
 function smooth_knn_dists(knn_dists::AbstractMatrix{S}, 
                           k::Integer, 
-                          local_connectivity::Integer;
+                          local_connectivity::Real;
                           niter::Integer=64,
                           bandwidth::AbstractFloat=1.,
                           ktol = 1e-5) where {S <: Real}
-    # TODO: implement local_connectivity
-    @inline minimum_nonzero(dists) = minimum(dists[dists .> 0.])
-    ρs = S[minimum_nonzero(knn_dists[:, i]) for i in 1:size(knn_dists, 2)]
+
+    nonzero_dists(dists) = @view dists[dists .> 0.]
+    ρs = zeros(S, size(knn_dists, 2)) 
     σs = Array{S}(undef, size(knn_dists, 2))
-
     for i in 1:size(knn_dists, 2)
+        nz_dists = nonzero_dists(knn_dists[:, i])
+        if length(nz_dists) >= local_connectivity
+            index = floor(Int, local_connectivity)
+            interpolation = local_connectivity - index
+            if index > 0
+                ρs[i] = nz_dists[index]
+                if interpolation > SMOOTH_K_TOLERANCE
+                    ρs[i] += interpolation * (nz_dists[index+1] - nz_dists[index]) 
+                end
+            else
+                ρs[i] = interpolation * nz_dists[1] 
+            end
+        elseif length(nz_dists) > 0
+            ρs[i] = maximum(nz_dists) 
+        end
         @inbounds σs[i] = smooth_knn_dist(knn_dists[:, i], ρs[i], k, local_connectivity, bandwidth, niter, ktol)
     end
+
     return ρs, σs
 end
 

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -1,5 +1,5 @@
 using Test
-using Distances: Euclidean
+using Distances: Euclidean, CosineDist
 using SparseArrays
 using LinearAlgebra
 using UMAP

diff --git a/test/umap_tests.jl b/test/umap_tests.jl
@@ -33,7 +33,13 @@
         @test all(0. .<= umap_graph .<= 1.)
         data = rand(Float32, 20, 500)
         umap_graph = fuzzy_simplicial_set(data, k, Euclidean(), 1, 1.f0)
+        @test issymmetric(umap_graph)
         @test eltype(umap_graph) == Float32
+
+        data = 2 .* rand(20, 1000) .- 1
+        umap_graph = fuzzy_simplicial_set(data, k, CosineDist(), 1, 1.)
+        @test issymmetric(umap_graph)
+        @test all(0. .<= umap_graph .<= 1.)
     end
 
     @testset "smooth_knn_dists" begin
@@ -58,6 +64,15 @@
         @test rhos == [1., 2., 3.]
         diffs = [psum(knn_dists[:,i], rhos[i], sigmas[i]) for i in 1:3] .- log2(6)
         @test all(diffs .< 1e-5)
+
+        knn_dists = [0. 0. 0.;
+                     0. 1. 2.;
+                     0. 2. 3.]
+        rhos, sigmas = smooth_knn_dists(knn_dists, 2, 1)
+        @test rhos == [0., 1., 2.]
+
+        rhos, sigmas = smooth_knn_dists(knn_dists, 2, 1.5)
+        @test rhos == [0., 1.5, 2.5]
     end
 
     @testset "compute_membership_strengths" begin