updated FLANN plug-in type

Author: wildart

misc/nearestneighbors.jl

@@ -13,16 +13,19 @@ struct KDTree <: ManifoldLearning.AbstractNearestNeighbors
     tree::NearestNeighbors.KDTree
 end
 show(io::IO, NN::KDTree) = print(io, "KDTree")
+size(NN::KDTree) = (length(NN.fitted.data[1]), length(NN.fitted.data))
+
 fit(::Type{KDTree}, X::AbstractMatrix{T}) where {T<:Real} =
     KDTree(X, NearestNeighbors.KDTree(X))
-size(NN::KDTree) = (length(NN.fitted.data[1]), length(NN.fitted.data))
+
 function knn(NN::KDTree, X::AbstractVecOrMat{T}, k::Integer;
              self::Bool=false, weights::Bool=true, kwargs...) where {T<:Real}
     m, n = size(X)
     @assert n > k "Number of observations must be more then $(k)"
     A, D = NearestNeighbors.knn(NN.tree, X, k, true)
     return A, D
 end
+
 function inradius(NN::KDTree, X::AbstractVecOrMat{T}, r::Real;
                   weights::Bool=false, kwargs...) where {T<:Real}
     m, n = size(X)
@@ -42,20 +45,50 @@ function inradius(NN::KDTree, X::AbstractVecOrMat{T}, r::Real;
     return A, W
 end
 
-#=
 # Wrapper around FLANN functionality
 using FLANN: FLANN
 struct FLANNTree{T <: Real} <: ManifoldLearning.AbstractNearestNeighbors
+    d::Int
     index::FLANN.FLANNIndex{T}
 end
 show(io::IO, NN::FLANNTree) = print(io, "FLANNTree")
+size(NN::FLANNTree) = (NN.d, length(NN.index))
+
 function fit(::Type{FLANNTree}, X::AbstractMatrix{T}) where {T<:Real}
     params = FLANN.FLANNParameters()
     idx = FLANN.flann(X, params)
-    FLANNTree(idx)
+    FLANNTree(size(X,1), idx)
+end
+
+function knn(NN::FLANNTree, X::AbstractVecOrMat{T}, k::Integer;
+             self::Bool=false, weights::Bool=false, kwargs...) where {T<:Real}
+    m, n = size(X)
+    E, D = FLANN.knn(NN.index, X, k+1)
+    idxs = (1:k).+(!self)
+
+    A = Vector{Vector{Int}}(undef, n)
+    W = Vector{Vector{T}}(undef, (weights ? n : 0))
+    for (i,(es, ds)) in enumerate(zip(eachcol(E), eachcol(D)))
+        A[i] = es[idxs]
+        if weights
+            W[i] = sqrt.(ds[idxs])
+        end
+    end
+    return A, W
 end
-function knn(NN::FLANNTree, X::AbstractVecOrMat{T}, k::Int; self=false) where {T<:Real}
-    E, D = FLANN.knn(NN.index, X, NN.k+1)
-    sqrt.(@view D[2:end, :]), @view E[2:end, :]
+
+function inradius(NN::FLANNTree, X::AbstractVecOrMat{T}, r::Real;
+                  weights::Bool=false, kwargs...) where {T<:Real}
+    m, n = size(X)
+    A = Vector{Vector{Int}}(undef, n)
+    W = Vector{Vector{T}}(undef, (weights ? n : 0))
+    for (i, x) in enumerate(eachcol(X))
+        E, D = FLANN.inrange(NN.index, x, r)
+        A[i] = E
+        if weights
+            W[i] = D
+        end
+    end
+    return A, W
 end
-=#
+