From bbdb0285066862791bd77ca2b92fc333af35034f Mon Sep 17 00:00:00 2001 From: Kristoffer Date: Tue, 23 Jul 2024 13:20:07 +0200 Subject: [PATCH] format with runic --- benchmark/benchmarkdatafreetree.jl | 8 +- benchmark/benchmarks.jl | 8 +- src/NearestNeighbors.jl | 41 +++--- src/ball_tree.jl | 153 +++++++++++++---------- src/brute_tree.jl | 71 +++++++---- src/datafreetree.jl | 6 +- src/evaluation.jl | 14 ++- src/hyperrectangles.jl | 4 +- src/hyperspheres.jl | 84 ++++++++----- src/inrange.jl | 26 ++-- src/kd_tree.jl | 156 +++++++++++++---------- src/knn.jl | 18 +-- src/tree_data.jl | 8 +- src/tree_ops.jl | 36 +++--- src/utilities.jl | 20 +-- test/runtests.jl | 16 +-- test/test_datafreetree.jl | 10 +- test/test_inrange.jl | 192 +++++++++++++++-------------- test/test_knn.jl | 28 +++-- test/test_monkey.jl | 4 +- 20 files changed, 514 insertions(+), 389 deletions(-) diff --git a/benchmark/benchmarkdatafreetree.jl b/benchmark/benchmarkdatafreetree.jl index 89e70bd..118f1bf 100644 --- a/benchmark/benchmarkdatafreetree.jl +++ b/benchmark/benchmarkdatafreetree.jl @@ -5,7 +5,7 @@ using Mmap runtimes = [] runtimesreordered = [] -function create_tree(n, reorder=false) +function create_tree(n, reorder = false) filename = tempname() d = 10 data = Mmap.mmap(filename, Matrix{Float32}, (d, n)) @@ -22,7 +22,7 @@ end function knnbench(tree, data, n, N) ind = rand(1:n, N) - knn(tree, data[:,ind], 3)[2] + knn(tree, data[:, ind], 3)[2] end function bench() @@ -34,10 +34,10 @@ function bench() tr, datar, filenamer = create_tree(n, true) bm = @benchmark knnbench(t, data, n, 1000) - push!(runtimes, mean(bm.samples.elapsed_times) / 1e9) + push!(runtimes, mean(bm.samples.elapsed_times) / 1.0e9) bmr = @benchmark knnbench(tr, datar, n, 1000) - push!(runtimesreordered, mean(bmr.samples.elapsed_times) / 1e9) + push!(runtimesreordered, mean(bmr.samples.elapsed_times) / 1.0e9) rm(filename) rm(filenamer) diff --git a/benchmark/benchmarks.jl b/benchmark/benchmarks.jl index b987905..485ac46 100644 --- a/benchmark/benchmarks.jl +++ b/benchmark/benchmarks.jl @@ -16,8 +16,10 @@ for n_points in (EXTENSIVE_BENCHMARK ? (10^3, 10^5) : 10^5) data = rand(StableRNG(123), dim, n_points) for leafsize in (EXTENSIVE_BENCHMARK ? (1, 10) : 10) for reorder in (true, false) - for (tree_type, SUITE_name) in ((KDTree, "kd tree"), - (BallTree, "ball tree")) + for (tree_type, SUITE_name) in ( + (KDTree, "kd tree"), + (BallTree, "ball tree"), + ) tree = tree_type(data; leafsize = leafsize, reorder = reorder) SUITE["build tree"]["$(tree_type) $dim × $n_points, ls = $leafsize"] = @benchmarkable $(tree_type)($data; leafsize = $leafsize, reorder = $reorder) for input_size in (1, 1000) @@ -27,7 +29,7 @@ for n_points in (EXTENSIVE_BENCHMARK ? (10^3, 10^5) : 10^5) end perc = 0.01 V = π^(dim / 2) / gamma(dim / 2 + 1) * (1 / 2)^dim - r = (V * perc * gamma(dim / 2 + 1))^(1/dim) + r = (V * perc * gamma(dim / 2 + 1))^(1 / dim) r_formatted = @sprintf("%3.2e", r) SUITE["inrange"]["$(tree_type) $dim × $n_points, ls = $leafsize, input_size = $input_size, r = $r_formatted"] = @benchmarkable inrange($tree, $input_data, $r) end diff --git a/src/NearestNeighbors.jl b/src/NearestNeighbors.jl index 28d998d..328964f 100644 --- a/src/NearestNeighbors.jl +++ b/src/NearestNeighbors.jl @@ -11,37 +11,46 @@ export knn, knn!, nn, inrange, inrange!,inrangecount # TODOs? , allpairs, distma export injectdata export Euclidean, - Cityblock, - Minkowski, - Chebyshev, - Hamming, - WeightedEuclidean, - WeightedCityblock, - WeightedMinkowski + Cityblock, + Minkowski, + Chebyshev, + Hamming, + WeightedEuclidean, + WeightedCityblock, + WeightedMinkowski -abstract type NNTree{V <: AbstractVector,P <: PreMetric} end +abstract type NNTree{V <: AbstractVector, P <: PreMetric} end -const NonweightedMinowskiMetric = Union{Euclidean,Chebyshev,Cityblock,Minkowski} -const WeightedMinowskiMetric = Union{WeightedEuclidean,WeightedCityblock,WeightedMinkowski} +const NonweightedMinowskiMetric = Union{Euclidean, Chebyshev, Cityblock, Minkowski} +const WeightedMinowskiMetric = Union{WeightedEuclidean, WeightedCityblock, WeightedMinkowski} const MinkowskiMetric = Union{NonweightedMinowskiMetric, WeightedMinowskiMetric} function check_input(::NNTree{V1}, ::AbstractVector{V2}) where {V1, V2 <: AbstractVector} if length(V1) != length(V2) - throw(ArgumentError( - "dimension of input points:$(length(V2)) and tree data:$(length(V1)) must agree")) + throw( + ArgumentError( + "dimension of input points:$(length(V2)) and tree data:$(length(V1)) must agree", + ), + ) end end function check_input(::NNTree{V1}, point::AbstractVector{T}) where {V1, T <: Number} if length(V1) != length(point) - throw(ArgumentError( - "dimension of input points:$(length(point)) and tree data:$(length(V1)) must agree")) + throw( + ArgumentError( + "dimension of input points:$(length(point)) and tree data:$(length(V1)) must agree", + ), + ) end end function check_input(::NNTree{V1}, m::AbstractMatrix) where {V1} if length(V1) != size(m, 1) - throw(ArgumentError( - "dimension of input points:$(size(m, 1)) and tree data:$(length(V1)) must agree")) + throw( + ArgumentError( + "dimension of input points:$(size(m, 1)) and tree data:$(length(V1)) must agree", + ), + ) end end diff --git a/src/ball_tree.jl b/src/ball_tree.jl index 1be8cc3..6006455 100644 --- a/src/ball_tree.jl +++ b/src/ball_tree.jl @@ -3,9 +3,9 @@ # which radius are determined from the given metric. # The tree uses the triangle inequality to prune the search space # when finding the neighbors to a point, -struct BallTree{V <: AbstractVector,N,T,M <: Metric} <: NNTree{V,M} +struct BallTree{V <: AbstractVector, N, T, M <: Metric} <: NNTree{V, M} data::Vector{V} - hyper_spheres::Vector{HyperSphere{N,T}} # Each hyper sphere bounds its children + hyper_spheres::Vector{HyperSphere{N, T}} # Each hyper sphere bounds its children indices::Vector{Int} # Translates from tree index -> point index metric::M # Metric used for tree tree_data::TreeData # Some constants needed @@ -18,12 +18,14 @@ end Creates a `BallTree` from the data using the given `metric` and `leafsize`. """ -function BallTree(data::AbstractVector{V}, - metric::Metric = Euclidean(); - leafsize::Int = 25, - reorder::Bool = true, - storedata::Bool = true, - reorderbuffer::Vector{V} = Vector{V}()) where {V <: AbstractArray} +function BallTree( + data::AbstractVector{V}, + metric::Metric = Euclidean(); + leafsize::Int = 25, + reorder::Bool = true, + storedata::Bool = true, + reorderbuffer::Vector{V} = Vector{V}(), + ) where {V <: AbstractArray} reorder = !isempty(reorderbuffer) || (storedata ? reorder : false) tree_data = TreeData(data, leafsize) @@ -32,7 +34,7 @@ function BallTree(data::AbstractVector{V}, indices = collect(1:n_p) # Bottom up creation of hyper spheres so need spheres even for leafs) - hyper_spheres = Vector{HyperSphere{length(V),eltype(V)}}(undef, tree_data.n_internal_nodes + tree_data.n_leafs) + hyper_spheres = Vector{HyperSphere{length(V), eltype(V)}}(undef, tree_data.n_internal_nodes + tree_data.n_leafs) indices_reordered = Vector{Int}() data_reordered = Vector{V}() @@ -49,53 +51,64 @@ function BallTree(data::AbstractVector{V}, if metric isa Distances.UnionMetrics p = parameters(metric) if p !== nothing && length(p) != length(V) - throw(ArgumentError( - "dimension of input points:$(length(V)) and metric parameter:$(length(p)) must agree")) + throw( + ArgumentError( + "dimension of input points:$(length(V)) and metric parameter:$(length(p)) must agree", + ), + ) end end if n_p > 0 # Call the recursive BallTree builder - build_BallTree(1, data, data_reordered, hyper_spheres, metric, indices, indices_reordered, - 1:length(data), tree_data, reorder) + build_BallTree( + 1, data, data_reordered, hyper_spheres, metric, indices, indices_reordered, + 1:length(data), tree_data, reorder, + ) end if reorder - data = data_reordered - indices = indices_reordered + data = data_reordered + indices = indices_reordered end BallTree(storedata ? data : similar(data, 0), hyper_spheres, indices, metric, tree_data, reorder) end -function BallTree(data::AbstractVecOrMat{T}, - metric::Metric = Euclidean(); - leafsize::Int = 25, - storedata::Bool = true, - reorder::Bool = true, - reorderbuffer::Matrix{T} = Matrix{T}(undef, 0, 0)) where {T <: AbstractFloat} +function BallTree( + data::AbstractVecOrMat{T}, + metric::Metric = Euclidean(); + leafsize::Int = 25, + storedata::Bool = true, + reorder::Bool = true, + reorderbuffer::Matrix{T} = Matrix{T}(undef, 0, 0), + ) where {T <: AbstractFloat} dim = size(data, 1) points = copy_svec(T, data, Val(dim)) if isempty(reorderbuffer) - reorderbuffer_points = Vector{SVector{dim,T}}() + reorderbuffer_points = Vector{SVector{dim, T}}() else reorderbuffer_points = copy_svec(T, reorderbuffer, Val(dim)) end - BallTree(points, metric; leafsize, storedata, reorder, - reorderbuffer = reorderbuffer_points) + BallTree( + points, metric; leafsize, storedata, reorder, + reorderbuffer = reorderbuffer_points, + ) end # Recursive function to build the tree. -function build_BallTree(index::Int, - data::AbstractVector{V}, - data_reordered::Vector{V}, - hyper_spheres::Vector{HyperSphere{N,T}}, - metric::Metric, - indices::Vector{Int}, - indices_reordered::Vector{Int}, - range::UnitRange{Int}, - tree_data::TreeData, - reorder::Bool) where {V <: AbstractVector, N, T} +function build_BallTree( + index::Int, + data::AbstractVector{V}, + data_reordered::Vector{V}, + hyper_spheres::Vector{HyperSphere{N, T}}, + metric::Metric, + indices::Vector{Int}, + indices_reordered::Vector{Int}, + range::UnitRange{Int}, + tree_data::TreeData, + reorder::Bool, + ) where {V <: AbstractVector, N, T} n_points = length(range) # Points left if n_points <= tree_data.leafsize @@ -118,35 +131,45 @@ function build_BallTree(index::Int, # to compare select_spec!(indices, mid_idx, first(range), last(range), data, split_dim) - build_BallTree(getleft(index), data, data_reordered, hyper_spheres, metric, - indices, indices_reordered, first(range):mid_idx - 1, - tree_data, reorder) + build_BallTree( + getleft(index), data, data_reordered, hyper_spheres, metric, + indices, indices_reordered, first(range):(mid_idx - 1), + tree_data, reorder, + ) - build_BallTree(getright(index), data, data_reordered, hyper_spheres, metric, - indices, indices_reordered, mid_idx:last(range), - tree_data, reorder) + build_BallTree( + getright(index), data, data_reordered, hyper_spheres, metric, + indices, indices_reordered, mid_idx:last(range), + tree_data, reorder, + ) # Finally create bounding hyper sphere from the two children's hyper spheres - hyper_spheres[index] = create_bsphere(metric, hyper_spheres[getleft(index)], - hyper_spheres[getright(index)]) + hyper_spheres[index] = create_bsphere( + metric, hyper_spheres[getleft(index)], + hyper_spheres[getright(index)], + ) end -function _knn(tree::BallTree, - point::AbstractVector, - best_idxs::AbstractVector{<:Integer}, - best_dists::AbstractVector, - skip::F) where {F} +function _knn( + tree::BallTree, + point::AbstractVector, + best_idxs::AbstractVector{<:Integer}, + best_dists::AbstractVector, + skip::F, + ) where {F} knn_kernel!(tree, 1, point, best_idxs, best_dists, skip) return end -function knn_kernel!(tree::BallTree{V}, - index::Int, - point::AbstractArray, - best_idxs::AbstractVector{<:Integer}, - best_dists::AbstractVector, - skip::F) where {V, F} +function knn_kernel!( + tree::BallTree{V}, + index::Int, + point::AbstractArray, + best_idxs::AbstractVector{<:Integer}, + best_dists::AbstractVector, + skip::F, + ) where {V, F} if isleaf(tree.tree_data.n_internal_nodes, index) add_points_knn!(best_dists, best_idxs, tree, index, point, true, skip) return @@ -174,19 +197,23 @@ function knn_kernel!(tree::BallTree{V}, return end -function _inrange(tree::BallTree{V}, - point::AbstractVector, - radius::Number, - idx_in_ball::Union{Nothing, Vector{<:Integer}}) where {V} +function _inrange( + tree::BallTree{V}, + point::AbstractVector, + radius::Number, + idx_in_ball::Union{Nothing, Vector{<:Integer}}, + ) where {V} ball = HyperSphere(convert(V, point), convert(eltype(V), radius)) # The "query ball" return inrange_kernel!(tree, 1, point, ball, idx_in_ball) # Call the recursive range finder end -function inrange_kernel!(tree::BallTree, - index::Int, - point::AbstractVector, - query_ball::HyperSphere, - idx_in_ball::Union{Nothing, Vector{<:Integer}}) +function inrange_kernel!( + tree::BallTree, + index::Int, + point::AbstractVector, + query_ball::HyperSphere, + idx_in_ball::Union{Nothing, Vector{<:Integer}}, + ) if index > length(tree.hyper_spheres) return 0 @@ -215,7 +242,7 @@ function inrange_kernel!(tree::BallTree, count += addall(tree, index, idx_in_ball) else # Recursively call the left and right sub tree. - count += inrange_kernel!(tree, getleft(index), point, query_ball, idx_in_ball) + count += inrange_kernel!(tree, getleft(index), point, query_ball, idx_in_ball) count += inrange_kernel!(tree, getright(index), point, query_ball, idx_in_ball) end return count diff --git a/src/brute_tree.jl b/src/brute_tree.jl index ed5ce9a..43b4c16 100644 --- a/src/brute_tree.jl +++ b/src/brute_tree.jl @@ -1,4 +1,4 @@ -struct BruteTree{V <: AbstractVector,M <: PreMetric} <: NNTree{V,M} +struct BruteTree{V <: AbstractVector, M <: PreMetric} <: NNTree{V, M} data::Vector{V} metric::M reordered::Bool @@ -9,41 +9,54 @@ end Creates a `BruteTree` from the data using the given `metric`. """ -function BruteTree(data::AbstractVector{V}, metric::PreMetric = Euclidean(); - reorder::Bool=false, leafsize::Int=0, storedata::Bool=true) where {V <: AbstractVector} +function BruteTree( + data::AbstractVector{V}, metric::PreMetric = Euclidean(); + reorder::Bool = false, leafsize::Int = 0, storedata::Bool = true, + ) where {V <: AbstractVector} if metric isa Distances.UnionMetrics p = parameters(metric) if p !== nothing && length(p) != length(V) - throw(ArgumentError( - "dimension of input points:$(length(V)) and metric parameter:$(length(p)) must agree")) + throw( + ArgumentError( + "dimension of input points:$(length(V)) and metric parameter:$(length(p)) must agree", + ), + ) end end BruteTree(storedata ? Vector(data) : Vector{V}(), metric, reorder) end -function BruteTree(data::AbstractVecOrMat{T}, metric::PreMetric = Euclidean(); - reorder::Bool=false, leafsize::Int=0, storedata::Bool=true) where {T} +function BruteTree( + data::AbstractVecOrMat{T}, metric::PreMetric = Euclidean(); + reorder::Bool = false, leafsize::Int = 0, storedata::Bool = true, + ) where {T} dim = size(data, 1) - BruteTree(copy_svec(T, data, Val(dim)), - metric, reorder = reorder, leafsize = leafsize, storedata = storedata) + BruteTree( + copy_svec(T, data, Val(dim)), + metric, reorder = reorder, leafsize = leafsize, storedata = storedata, + ) end -function _knn(tree::BruteTree{V}, - point::AbstractVector, - best_idxs::AbstractVector{<:Integer}, - best_dists::AbstractVector, - skip::F) where {V, F} +function _knn( + tree::BruteTree{V}, + point::AbstractVector, + best_idxs::AbstractVector{<:Integer}, + best_dists::AbstractVector, + skip::F, + ) where {V, F} knn_kernel!(tree, point, best_idxs, best_dists, skip) return end -function knn_kernel!(tree::BruteTree{V}, - point::AbstractVector, - best_idxs::AbstractVector{<:Integer}, - best_dists::AbstractVector, - skip::F) where {V, F} +function knn_kernel!( + tree::BruteTree{V}, + point::AbstractVector, + best_idxs::AbstractVector{<:Integer}, + best_dists::AbstractVector, + skip::F, + ) where {V, F} for i in 1:length(tree.data) if skip(i) continue @@ -58,18 +71,22 @@ function knn_kernel!(tree::BruteTree{V}, end end -function _inrange(tree::BruteTree, - point::AbstractVector, - radius::Number, - idx_in_ball::Union{Nothing, Vector{<:Integer}}) +function _inrange( + tree::BruteTree, + point::AbstractVector, + radius::Number, + idx_in_ball::Union{Nothing, Vector{<:Integer}}, + ) return inrange_kernel!(tree, point, radius, idx_in_ball) end -function inrange_kernel!(tree::BruteTree, - point::AbstractVector, - r::Number, - idx_in_ball::Union{Nothing, Vector{<:Integer}}) +function inrange_kernel!( + tree::BruteTree, + point::AbstractVector, + r::Number, + idx_in_ball::Union{Nothing, Vector{<:Integer}}, + ) count = 0 for i in 1:length(tree.data) d = evaluate(tree.metric, tree.data[i], point) diff --git a/src/datafreetree.jl b/src/datafreetree.jl index b0b3977..0ea0297 100644 --- a/src/datafreetree.jl +++ b/src/datafreetree.jl @@ -1,7 +1,7 @@ # A DataFreeTree wraps a descendant of NNTree # which does not contain a copy of the data struct DataFreeTree{N <: NNTree} - size::Tuple{Int,Int} + size::Tuple{Int, Int} hash::UInt64 tree::N end @@ -45,13 +45,13 @@ function injectdata(datafreetree::DataFreeTree, data::AbstractMatrix{T}) where { if isbitstype(T) new_data = copy_svec(T, data, Val(dim)) else - new_data = SVector{dim,T}[SVector{dim,T}(data[:, i]) for i in 1:npoints] + new_data = SVector{dim, T}[SVector{dim, T}(data[:, i]) for i in 1:npoints] end new_hash = hash(data) injectdata(datafreetree, new_data, new_hash) end -function injectdata(datafreetree::DataFreeTree, data::AbstractVector{V}, new_hash::UInt64=0) where {V <: AbstractVector} +function injectdata(datafreetree::DataFreeTree, data::AbstractVector{V}, new_hash::UInt64 = 0) where {V <: AbstractVector} if new_hash == 0 new_hash = hash(data) end diff --git a/src/evaluation.jl b/src/evaluation.jl index 807adbf..375caaf 100644 --- a/src/evaluation.jl +++ b/src/evaluation.jl @@ -5,10 +5,12 @@ @inline eval_diff(::NonweightedMinowskiMetric, a, b, dim) = a - b @inline eval_diff(::Chebyshev, ::Any, b, dim) = b -@inline eval_diff(m::WeightedMinowskiMetric, a, b, dim) = m.weights[dim] * (a-b) +@inline eval_diff(m::WeightedMinowskiMetric, a, b, dim) = m.weights[dim] * (a - b) -function evaluate_maybe_end(d::Distances.UnionMetrics, a::AbstractVector, - b::AbstractVector, do_end::Bool) +function evaluate_maybe_end( + d::Distances.UnionMetrics, a::AbstractVector, + b::AbstractVector, do_end::Bool, + ) p = Distances.parameters(d) s = eval_start(d, a, b) if p === nothing @@ -32,7 +34,9 @@ function evaluate_maybe_end(d::Distances.UnionMetrics, a::AbstractVector, end end -function evaluate_maybe_end(d::Distances.PreMetric, a::AbstractVector, - b::AbstractVector, ::Bool) +function evaluate_maybe_end( + d::Distances.PreMetric, a::AbstractVector, + b::AbstractVector, ::Bool, + ) evaluate(d, a, b) end diff --git a/src/hyperrectangles.jl b/src/hyperrectangles.jl index 261fac1..bd3b2a4 100644 --- a/src/hyperrectangles.jl +++ b/src/hyperrectangles.jl @@ -4,8 +4,8 @@ struct HyperRectangle{V <: AbstractVector} end function compute_bbox(data::AbstractVector{V}) where {V <: AbstractVector} - mins = mapreduce(identity, (a, b) -> min.(a, b), data; init=fill(Inf,V)) - maxes = mapreduce(identity, (a, b) -> max.(a, b), data; init=fill(-Inf,V)) + mins = mapreduce(identity, (a, b) -> min.(a, b), data; init = fill(Inf, V)) + maxes = mapreduce(identity, (a, b) -> max.(a, b), data; init = fill(-Inf, V)) return HyperRectangle(mins, maxes) end diff --git a/src/hyperspheres.jl b/src/hyperspheres.jl index f913552..169ab44 100644 --- a/src/hyperspheres.jl +++ b/src/hyperspheres.jl @@ -1,47 +1,57 @@ -const NormMetric = Union{Euclidean,Chebyshev,Cityblock,Minkowski,WeightedEuclidean,WeightedCityblock,WeightedMinkowski,Mahalanobis} +const NormMetric = Union{Euclidean, Chebyshev, Cityblock, Minkowski, WeightedEuclidean, WeightedCityblock, WeightedMinkowski, Mahalanobis} -struct HyperSphere{N,T <: AbstractFloat} - center::SVector{N,T} +struct HyperSphere{N, T <: AbstractFloat} + center::SVector{N, T} r::T end -HyperSphere(center::SVector{N,T1}, r) where {N, T1} = HyperSphere(center, convert(T1, r)) +HyperSphere(center::SVector{N, T1}, r) where {N, T1} = HyperSphere(center, convert(T1, r)) HyperSphere(center::AbstractVector, r) = HyperSphere(SVector{length(center)}(center), r) -@inline function intersects(m::Metric, - s1::HyperSphere{N}, - s2::HyperSphere{N}) where {N} +@inline function intersects( + m::Metric, + s1::HyperSphere{N}, + s2::HyperSphere{N}, + ) where {N} d = evaluate(m, s1.center, s2.center) return d, d <= s1.r + s2.r end -@inline function intersects(m::MinkowskiMetric, - s1::HyperSphere{N}, - s2::HyperSphere{N}) where {N} +@inline function intersects( + m::MinkowskiMetric, + s1::HyperSphere{N}, + s2::HyperSphere{N}, + ) where {N} d = evaluate_maybe_end(m, s1.center, s2.center, false) return d, d <= eval_pow(m, s1.r + s2.r) end -@inline function encloses(m::Metric, - s1::HyperSphere{N}, - s2::HyperSphere{N}) where {N} +@inline function encloses( + m::Metric, + s1::HyperSphere{N}, + s2::HyperSphere{N}, + ) where {N} evaluate(m, s1.center, s2.center) + s1.r <= s2.r end -@inline function encloses_fast(d, m::Metric, - s1::HyperSphere{N}, - s2::HyperSphere{N}) where {N} - if s1.r > s2.r +@inline function encloses_fast( + d, m::Metric, + s1::HyperSphere{N}, + s2::HyperSphere{N}, + ) where {N} + if s1.r > s2.r return false else return d + s1.r <= s2.r end end -@inline function encloses_fast(d, m::MinkowskiMetric, - s1::HyperSphere{N}, - s2::HyperSphere{N}) where {N} +@inline function encloses_fast( + d, m::MinkowskiMetric, + s1::HyperSphere{N}, + s2::HyperSphere{N}, + ) where {N} if s1.r > s2.r return false else @@ -49,21 +59,25 @@ end end end -@inline function interpolate(::NormMetric, - c1::V, - c2::V, - x, - d) where {V <: AbstractVector} +@inline function interpolate( + ::NormMetric, + c1::V, + c2::V, + x, + d, + ) where {V <: AbstractVector} alpha = x / d center = (1 - alpha) * c1 + alpha * c2 return center, true end -@inline function interpolate(::Metric, - c1::V, - ::V, - ::Any, - ::Any) where {V <: AbstractVector} +@inline function interpolate( + ::Metric, + c1::V, + ::V, + ::Any, + ::Any, + ) where {V <: AbstractVector} return c1, false end @@ -77,9 +91,11 @@ function create_bsphere(data::AbstractVector{V}, metric::Metric, indices::Vector end # Creates a bounding sphere from two other spheres -function create_bsphere(m::Metric, - s1::HyperSphere{N,T}, - s2::HyperSphere{N,T}) where {N, T <: AbstractFloat} +function create_bsphere( + m::Metric, + s1::HyperSphere{N, T}, + s2::HyperSphere{N, T}, + ) where {N, T <: AbstractFloat} if encloses(m, s1, s2) return HyperSphere(s2.center, s2.r) elseif encloses(m, s2, s1) @@ -98,7 +114,7 @@ function create_bsphere(m::Metric, rad = max(s1.r + evaluate(m, s1.center, center), s2.r + evaluate(m, s2.center, center)) end - return HyperSphere(SVector{N,T}(center), rad) + return HyperSphere(SVector{N, T}(center), rad) end function distance_to_sphere(metric::Metric, point, sphere::HyperSphere) diff --git a/src/inrange.jl b/src/inrange.jl index d271639..d7e9e6c 100644 --- a/src/inrange.jl +++ b/src/inrange.jl @@ -8,10 +8,12 @@ Find all the points in the tree which is closer than `radius` to `points`. If See also: `inrange!`, `inrangecount`. """ -function inrange(tree::NNTree, - points::AbstractVector{T}, - radius::Number, - sortres=false) where {T <: AbstractVector} +function inrange( + tree::NNTree, + points::AbstractVector{T}, + radius::Number, + sortres = false, + ) where {T <: AbstractVector} check_input(tree, points) check_radius(radius) @@ -44,7 +46,7 @@ Useful if one want to avoid allocations or specify the element type of the outpu See also: `inrange`, `inrangecount`. """ -function inrange!(idxs::AbstractVector, tree::NNTree{V}, point::AbstractVector{T}, radius::Number, sortres=false) where {V, T <: Number} +function inrange!(idxs::AbstractVector, tree::NNTree{V}, point::AbstractVector{T}, radius::Number, sortres = false) where {V, T <: Number} check_input(tree, point) check_radius(radius) length(idxs) == 0 || throw(ArgumentError("idxs must be empty")) @@ -52,11 +54,11 @@ function inrange!(idxs::AbstractVector, tree::NNTree{V}, point::AbstractVector{T return idxs end -function inrange(tree::NNTree{V}, point::AbstractVector{T}, radius::Number, sortres=false) where {V, T <: Number} +function inrange(tree::NNTree{V}, point::AbstractVector{T}, radius::Number, sortres = false) where {V, T <: Number} return inrange!(Int[], tree, point, radius, sortres) end -function inrange(tree::NNTree{V}, points::AbstractMatrix{T}, radius::Number, sortres=false) where {V, T <: Number} +function inrange(tree::NNTree{V}, points::AbstractMatrix{T}, radius::Number, sortres = false) where {V, T <: Number} dim = size(points, 1) inrange_matrix(tree, points, radius, Val(dim), sortres) end @@ -69,7 +71,7 @@ function inrange_matrix(tree::NNTree{V}, points::AbstractMatrix{T}, radius::Numb idxs = [Vector{Int}() for _ in 1:n_points] for i in 1:n_points - point = SVector{dim,T}(ntuple(j -> points[j, i], Val(dim))) + point = SVector{dim, T}(ntuple(j -> points[j, i], Val(dim))) inrange_point!(tree, point, radius, sortres, idxs[i]) end return idxs @@ -86,9 +88,11 @@ function inrangecount(tree::NNTree{V}, point::AbstractVector{T}, radius::Number) return inrange_point!(tree, point, radius, false, nothing) end -function inrangecount(tree::NNTree, +function inrangecount( + tree::NNTree, points::AbstractVector{T}, - radius::Number) where {T <: AbstractVector} + radius::Number, + ) where {T <: AbstractVector} check_input(tree, points) check_radius(radius) return inrange_point!.(Ref(tree), points, radius, false, nothing) @@ -100,7 +104,7 @@ function inrangecount(tree::NNTree{V}, point::AbstractMatrix{T}, radius::Number) if isbitstype(T) new_data = copy_svec(T, point, Val(dim)) else - new_data = SVector{dim,T}[SVector{dim,T}(point[:, i]) for i in 1:npoints] + new_data = SVector{dim, T}[SVector{dim, T}(point[:, i]) for i in 1:npoints] end return inrangecount(tree, new_data, radius) end diff --git a/src/kd_tree.jl b/src/kd_tree.jl index 5518d7d..dbcd251 100644 --- a/src/kd_tree.jl +++ b/src/kd_tree.jl @@ -1,4 +1,4 @@ -struct KDTree{V <: AbstractVector, M <: MinkowskiMetric, T, TH} <: NNTree{V,M} +struct KDTree{V <: AbstractVector, M <: MinkowskiMetric, T, TH} <: NNTree{V, M} data::Vector{V} hyper_rec::HyperRectangle{TH} indices::Vector{Int} @@ -16,12 +16,14 @@ end Creates a `KDTree` from the data using the given `metric` and `leafsize`. The `metric` must be a `MinkowskiMetric`. """ -function KDTree(data::AbstractVector{V}, - metric::M = Euclidean(); - leafsize::Int = 25, - storedata::Bool = true, - reorder::Bool = true, - reorderbuffer::Vector{V} = Vector{V}()) where {V <: AbstractArray, M <: MinkowskiMetric} +function KDTree( + data::AbstractVector{V}, + metric::M = Euclidean(); + leafsize::Int = 25, + storedata::Bool = true, + reorder::Bool = true, + reorderbuffer::Vector{V} = Vector{V}(), + ) where {V <: AbstractArray, M <: MinkowskiMetric} reorder = !isempty(reorderbuffer) || (storedata ? reorder : false) tree_data = TreeData(data, leafsize) @@ -47,8 +49,11 @@ function KDTree(data::AbstractVector{V}, if metric isa Distances.UnionMetrics p = parameters(metric) if p !== nothing && length(p) != length(V) - throw(ArgumentError( - "dimension of input points:$(length(V)) and metric parameter:$(length(p)) must agree")) + throw( + ArgumentError( + "dimension of input points:$(length(V)) and metric parameter:$(length(p)) must agree", + ), + ) end end @@ -56,8 +61,10 @@ function KDTree(data::AbstractVector{V}, hyper_rec = compute_bbox(data) # Call the recursive KDTree builder - build_KDTree(1, data, data_reordered, hyper_rec, split_vals, split_dims, indices, indices_reordered, - 1:length(data), tree_data, reorder) + build_KDTree( + 1, data, data_reordered, hyper_rec, split_vals, split_dims, indices, indices_reordered, + 1:length(data), tree_data, reorder, + ) if reorder data = data_reordered indices = indices_reordered @@ -66,42 +73,51 @@ function KDTree(data::AbstractVector{V}, if metric isa Distances.UnionMetrics p = parameters(metric) if p !== nothing && length(p) != length(V) - throw(ArgumentError( - "dimension of input points:$(length(V)) and metric parameter:$(length(p)) must agree")) + throw( + ArgumentError( + "dimension of input points:$(length(V)) and metric parameter:$(length(p)) must agree", + ), + ) end end KDTree(storedata ? data : similar(data, 0), hyper_rec, indices, metric, split_vals, split_dims, tree_data, reorder) end - function KDTree(data::AbstractVecOrMat{T}, - metric::M = Euclidean(); - leafsize::Int = 25, - storedata::Bool = true, - reorder::Bool = true, - reorderbuffer::Matrix{T} = Matrix{T}(undef, 0, 0)) where {T <: AbstractFloat, M <: MinkowskiMetric} +function KDTree( + data::AbstractVecOrMat{T}, + metric::M = Euclidean(); + leafsize::Int = 25, + storedata::Bool = true, + reorder::Bool = true, + reorderbuffer::Matrix{T} = Matrix{T}(undef, 0, 0), + ) where {T <: AbstractFloat, M <: MinkowskiMetric} dim = size(data, 1) points = copy_svec(T, data, Val(dim)) if isempty(reorderbuffer) - reorderbuffer_points = Vector{SVector{dim,T}}() + reorderbuffer_points = Vector{SVector{dim, T}}() else reorderbuffer_points = copy_svec(T, reorderbuffer, Val(dim)) end - KDTree(points, metric; leafsize, storedata, reorder, - reorderbuffer = reorderbuffer_points) + KDTree( + points, metric; leafsize, storedata, reorder, + reorderbuffer = reorderbuffer_points, + ) end -function build_KDTree(index::Int, - data::AbstractVector{V}, - data_reordered::Vector{V}, - hyper_rec::HyperRectangle, - split_vals::Vector{T}, - split_dims::Vector{UInt16}, - indices::Vector{Int}, - indices_reordered::Vector{Int}, - range, - tree_data::TreeData, - reorder::Bool) where {V <: AbstractVector, T} +function build_KDTree( + index::Int, + data::AbstractVector{V}, + data_reordered::Vector{V}, + hyper_rec::HyperRectangle, + split_vals::Vector{T}, + split_dims::Vector{UInt16}, + indices::Vector{Int}, + indices_reordered::Vector{Int}, + range, + tree_data::TreeData, + reorder::Bool, + ) where {V <: AbstractVector, T} n_p = length(range) # Points left if n_p <= tree_data.leafsize if reorder @@ -133,22 +149,28 @@ function build_KDTree(index::Int, # Call the left sub tree with an updated hyper rectangle new_maxes = @inbounds setindex(hyper_rec.maxes, split_val, split_dim) hyper_rec_left = HyperRectangle(hyper_rec.mins, new_maxes) - build_KDTree(getleft(index), data, data_reordered, hyper_rec_left, split_vals, split_dims, - indices, indices_reordered, first(range):mid_idx - 1, tree_data, reorder) + build_KDTree( + getleft(index), data, data_reordered, hyper_rec_left, split_vals, split_dims, + indices, indices_reordered, first(range):(mid_idx - 1), tree_data, reorder, + ) # Call the right sub tree with an updated hyper rectangle new_mins = @inbounds setindex(hyper_rec.mins, split_val, split_dim) hyper_rec_right = HyperRectangle(new_mins, hyper_rec.maxes) - build_KDTree(getright(index), data, data_reordered, hyper_rec_right, split_vals, split_dims, - indices, indices_reordered, mid_idx:last(range), tree_data, reorder) + build_KDTree( + getright(index), data, data_reordered, hyper_rec_right, split_vals, split_dims, + indices, indices_reordered, mid_idx:last(range), tree_data, reorder, + ) end -function _knn(tree::KDTree, - point::AbstractVector, - best_idxs::AbstractVector{<:Integer}, - best_dists::AbstractVector, - skip::F) where {F} +function _knn( + tree::KDTree, + point::AbstractVector, + best_idxs::AbstractVector{<:Integer}, + best_dists::AbstractVector, + skip::F, + ) where {F} init_min = get_min_distance_no_end(tree.metric, tree.hyper_rec, point) knn_kernel!(tree, 1, point, best_idxs, best_dists, init_min, tree.hyper_rec, skip) @simd for i in eachindex(best_dists) @@ -156,14 +178,16 @@ function _knn(tree::KDTree, end end -function knn_kernel!(tree::KDTree{V}, - index::Int, - point::AbstractVector, - best_idxs::AbstractVector{<:Integer}, - best_dists::AbstractVector, - min_dist, - hyper_rec::HyperRectangle, - skip::F) where {V, F} +function knn_kernel!( + tree::KDTree{V}, + index::Int, + point::AbstractVector, + best_idxs::AbstractVector{<:Integer}, + best_dists::AbstractVector, + min_dist, + hyper_rec::HyperRectangle, + skip::F, + ) where {V, F} # At a leaf node. Go through all points in node and add those in range if isleaf(tree.tree_data.n_internal_nodes, index) add_points_knn!(best_dists, best_idxs, tree, index, point, false, skip) @@ -204,23 +228,29 @@ function knn_kernel!(tree::KDTree{V}, return end -function _inrange(tree::KDTree, - point::AbstractVector, - radius::Number, - idx_in_ball::Union{Nothing, Vector{<:Integer}} = Int[]) +function _inrange( + tree::KDTree, + point::AbstractVector, + radius::Number, + idx_in_ball::Union{Nothing, Vector{<:Integer}} = Int[], + ) init_min = get_min_distance_no_end(tree.metric, tree.hyper_rec, point) - return inrange_kernel!(tree, 1, point, eval_pow(tree.metric, radius), idx_in_ball, - tree.hyper_rec, init_min) + return inrange_kernel!( + tree, 1, point, eval_pow(tree.metric, radius), idx_in_ball, + tree.hyper_rec, init_min, + ) end # Explicitly check the distance between leaf node and point while traversing -function inrange_kernel!(tree::KDTree, - index::Int, - point::AbstractVector, - r::Number, - idx_in_ball::Union{Nothing, Vector{<:Integer}}, - hyper_rec::HyperRectangle, - min_dist) +function inrange_kernel!( + tree::KDTree, + index::Int, + point::AbstractVector, + r::Number, + idx_in_ball::Union{Nothing, Vector{<:Integer}}, + hyper_rec::HyperRectangle, + min_dist, + ) # Point is outside hyper rectangle, skip the whole sub tree if min_dist > r return 0 diff --git a/src/knn.jl b/src/knn.jl index 775f7eb..5c12450 100644 --- a/src/knn.jl +++ b/src/knn.jl @@ -14,7 +14,7 @@ index. See also: `knn!`, `nn`. """ -function knn(tree::NNTree{V}, points::AbstractVector{T}, k::Int, sortres=false, skip::F=always_false) where {V, T <: AbstractVector, F<:Function} +function knn(tree::NNTree{V}, points::AbstractVector{T}, k::Int, sortres = false, skip::F = always_false) where {V, T <: AbstractVector, F <: Function} check_input(tree, points) check_k(tree, k) n_points = length(points) @@ -52,7 +52,7 @@ Useful if one want to avoid allocations or specify the element type of the outpu See also: `knn`, `nn`. """ -function knn!(idxs::AbstractVector{<:Integer}, dists::AbstractVector, tree::NNTree{V}, point::AbstractVector{T}, k::Int, sortres=false, skip::F=always_false) where {V, T <: Number, F<:Function} +function knn!(idxs::AbstractVector{<:Integer}, dists::AbstractVector, tree::NNTree{V}, point::AbstractVector{T}, k::Int, sortres = false, skip::F = always_false) where {V, T <: Number, F <: Function} check_k(tree, k) length(idxs) == k || throw(ArgumentError("idxs must be of length k")) length(dists) == k || throw(ArgumentError("dists must be of length k")) @@ -60,19 +60,19 @@ function knn!(idxs::AbstractVector{<:Integer}, dists::AbstractVector, tree::NNTr return idxs, dists end -function knn(tree::NNTree{V}, point::AbstractVector{T}, k::Int, sortres=false, skip::F=always_false) where {V, T <: Number, F<:Function} +function knn(tree::NNTree{V}, point::AbstractVector{T}, k::Int, sortres = false, skip::F = always_false) where {V, T <: Number, F <: Function} idx = Vector{Int}(undef, k) dist = Vector{get_T(eltype(V))}(undef, k) return knn!(idx, dist, tree, point, k, sortres, skip) end -function knn(tree::NNTree{V}, points::AbstractMatrix{T}, k::Int, sortres=false, skip::F=always_false) where {V, T <: Number, F<:Function} +function knn(tree::NNTree{V}, points::AbstractMatrix{T}, k::Int, sortres = false, skip::F = always_false) where {V, T <: Number, F <: Function} dim = size(points, 1) knn_matrix(tree, points, k, Val(dim), sortres, skip) end # Function barrier -function knn_matrix(tree::NNTree{V}, points::AbstractMatrix{T}, k::Int, ::Val{dim}, sortres=false, skip::F=always_false) where {V, T <: Number, F<:Function, dim} +function knn_matrix(tree::NNTree{V}, points::AbstractMatrix{T}, k::Int, ::Val{dim}, sortres = false, skip::F = always_false) where {V, T <: Number, F <: Function, dim} # TODO: DRY with knn for AbstractVector check_input(tree, points) check_k(tree, k) @@ -81,7 +81,7 @@ function knn_matrix(tree::NNTree{V}, points::AbstractMatrix{T}, k::Int, ::Val{di idxs = [Vector{Int}(undef, k) for _ in 1:n_points] for i in 1:n_points - point = SVector{dim,T}(ntuple(j -> points[j, i], Val(dim))) + point = SVector{dim, T}(ntuple(j -> points[j, i], Val(dim))) knn_point!(tree, point, sortres, dists[i], idxs[i], skip) end return idxs, dists @@ -95,9 +95,9 @@ Performs a lookup of the single nearest neigbours to the `points` from the data. See also: `knn`. """ -nn(tree::NNTree{V}, points::AbstractVector{T}, skip::F=always_false) where {V, T <: Number, F <: Function} = _nn(tree, points, skip) .|> only -nn(tree::NNTree{V}, points::AbstractVector{T}, skip::F=always_false) where {V, T <: AbstractVector, F <: Function} = _nn(tree, points, skip) |> _onlyeach -nn(tree::NNTree{V}, points::AbstractMatrix{T}, skip::F=always_false) where {V, T <: Number, F <: Function} = _nn(tree, points, skip) |> _onlyeach +nn(tree::NNTree{V}, points::AbstractVector{T}, skip::F = always_false) where {V, T <: Number, F <: Function} = _nn(tree, points, skip) .|> only +nn(tree::NNTree{V}, points::AbstractVector{T}, skip::F = always_false) where {V, T <: AbstractVector, F <: Function} = _nn(tree, points, skip) |> _onlyeach +nn(tree::NNTree{V}, points::AbstractMatrix{T}, skip::F = always_false) where {V, T <: Number, F <: Function} = _nn(tree, points, skip) |> _onlyeach _nn(tree, points, skip) = knn(tree, points, 1, false, skip) diff --git a/src/tree_data.jl b/src/tree_data.jl index f8b692c..7a6c276 100644 --- a/src/tree_data.jl +++ b/src/tree_data.jl @@ -10,7 +10,7 @@ struct TreeData end -function TreeData(data::AbstractVector{V}, leafsize) where V +function TreeData(data::AbstractVector{V}, leafsize) where {V} n_p = length(data) # If number of points is zero @@ -35,6 +35,8 @@ function TreeData(data::AbstractVector{V}, leafsize) where V k2 = -cross_node * leafsize + 1 last_full_node = n_leafs + n_internal_nodes - TreeData(last_node_size, leafsize, n_leafs, - n_internal_nodes, cross_node, k1, k2, last_full_node) + TreeData( + last_node_size, leafsize, n_leafs, + n_internal_nodes, cross_node, k1, k2, last_full_node, + ) end diff --git a/src/tree_ops.jl b/src/tree_ops.jl index 39338cf..fc02980 100644 --- a/src/tree_ops.jl +++ b/src/tree_ops.jl @@ -9,7 +9,7 @@ function show(io::IO, tree::NNTree{V}) where {V} println(io, " Number of points: ", length(tree.data)) println(io, " Dimensions: ", length(V)) println(io, " Metric: ", tree.metric) - print(io, " Reordered: ", tree.reordered) + print(io, " Reordered: ", tree.reordered) end # We split the tree such that one of the sub trees has exactly 2^p points @@ -28,22 +28,20 @@ function find_split(low, leafsize, n_p) # The conditionals here fulfill the desired splitting procedure but # can probably be written in a nicer way - - # Can fill less than two nodes -> leafsize to left node. if n_p <= 2 * leafsize + # Can fill less than two nodes -> leafsize to left node. mid_idx = leafsize - # The last leaf node will be in the right sub tree -> fill the left - # sub tree with elseif rest > 2^(k - 1) # Last node over the "half line" in the row + # The last leaf node will be in the right sub tree -> fill the left + # sub tree with mid_idx = 2^k * leafsize - # Perfectly filling both sub trees -> half to left and right sub tree elseif rest == 0 + # Perfectly filling both sub trees -> half to left and right sub tree mid_idx = 2^(k - 1) * leafsize - - # Else we fill the right sub tree -> send the rest to the left sub tree else + # Else we fill the right sub tree -> send the rest to the left sub tree mid_idx = n_p - 2^(k - 1) * leafsize end return mid_idx + low @@ -72,13 +70,15 @@ end @inline function get_leaf_range(td::TreeData, index) p_index = point_index(index, td) n_p = n_ps(index, td) - return p_index:p_index + n_p - 1 + return p_index:(p_index + n_p - 1) end # Store all the points in a leaf node continuously in memory in data_reordered to improve cache locality. # Also stores the mapping to get the index into the original data from the reordered data. -function reorder_data!(data_reordered::Vector{V}, data::AbstractVector{V}, index::Int, - indices::Vector{Int}, indices_reordered::Vector{Int}, tree_data::TreeData) where {V} +function reorder_data!( + data_reordered::Vector{V}, data::AbstractVector{V}, index::Int, + indices::Vector{Int}, indices_reordered::Vector{Int}, tree_data::TreeData, + ) where {V} for i in get_leaf_range(tree_data, index) idx = indices[i] @@ -90,9 +90,11 @@ end # Checks the distance function and add those points that are among the k best. # Uses a heap for fast insertion. -@inline function add_points_knn!(best_dists::AbstractVector, best_idxs::AbstractVector{<:Integer}, - tree::NNTree, index::Int, point::AbstractVector, - do_end::Bool, skip::F) where {F} +@inline function add_points_knn!( + best_dists::AbstractVector, best_idxs::AbstractVector{<:Integer}, + tree::NNTree, index::Int, point::AbstractVector, + do_end::Bool, skip::F, + ) where {F} for z in get_leaf_range(tree.tree_data, index) idx = tree.reordered ? z : tree.indices[z] dist_d = evaluate_maybe_end(tree.metric, tree.data[idx], point, do_end) @@ -114,8 +116,10 @@ end # stop computing the distance function as soon as we reach the desired radius. # This will probably prevent SIMD and other optimizations so some care is needed # to evaluate if it is worth it. -@inline function add_points_inrange!(idx_in_ball::Union{Nothing, AbstractVector{<:Integer}}, tree::NNTree, - index::Int, point::AbstractVector, r::Number) +@inline function add_points_inrange!( + idx_in_ball::Union{Nothing, AbstractVector{<:Integer}}, tree::NNTree, + index::Int, point::AbstractVector, r::Number, + ) count = 0 for z in get_leaf_range(tree.tree_data, index) idx = tree.reordered ? z : tree.indices[z] diff --git a/src/utilities.jl b/src/utilities.jl index 7a7f30a..5d7798b 100644 --- a/src/utilities.jl +++ b/src/utilities.jl @@ -24,8 +24,10 @@ end # Taken from https://github.com/JuliaLang/julia/blob/v0.3.5/base/sort.jl # and modified to compare against a matrix -@inline function select_spec!(v::Vector{Int}, k::Int, lo::Int, - hi::Int, data::AbstractVector, dim::Int) +@inline function select_spec!( + v::Vector{Int}, k::Int, lo::Int, + hi::Int, data::AbstractVector, dim::Int, + ) lo <= k <= hi || error("select index $k is out of range $lo:$hi") @inbounds while lo < hi if hi - lo == 1 @@ -65,11 +67,13 @@ end end # Binary max-heap percolate down. -@inline function percolate_down!(xs::AbstractArray, - xis::AbstractArray, - dist::Number, - index::Integer, - len::Integer=length(xs)) +@inline function percolate_down!( + xs::AbstractArray, + xis::AbstractArray, + dist::Number, + index::Integer, + len::Integer = length(xs), + ) i = 1 @inbounds while (l = getleft(i)) <= len r = getright(i) @@ -94,4 +98,4 @@ end # Instead of ReinterpretArray wrapper, copy an array, interpreting it as a vector of SVectors copy_svec(::Type{T}, data, ::Val{dim}) where {T, dim} = - [SVector{dim,T}(ntuple(i -> data[n+i], Val(dim))) for n in 0:dim:(length(data)-1)] + [SVector{dim, T}(ntuple(i -> data[n + i], Val(dim))) for n in 0:dim:(length(data) - 1)] diff --git a/test/runtests.jl b/test/runtests.jl index 584a271..215e64e 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -7,12 +7,14 @@ using LinearAlgebra using Distances: Distances, Metric, evaluate, PeriodicEuclidean struct CustomMetric1 <: Metric end Distances.evaluate(::CustomMetric1, a::AbstractVector, b::AbstractVector) = maximum(abs.(a .- b)) -function NearestNeighbors.interpolate(::CustomMetric1, - a::V, - b::V, - x, - d, - ab) where {V <: AbstractVector} +function NearestNeighbors.interpolate( + ::CustomMetric1, + a::V, + b::V, + x, + d, + ab, + ) where {V <: AbstractVector} idx = (abs.(b .- a) .>= d - x) c = copy(Array(a)) c[idx] = (1 - x / d) * a[idx] + (x / d) * b[idx] @@ -46,7 +48,7 @@ end pred = PeriodicEuclidean([Inf, 2.5]) l = [0.0 0.0; 0.0 2.5] S = BallTree(l, pred) - @test inrange(S,[0.0,0.0], 1e-2, true) == [1, 2] + @test inrange(S, [0.0, 0.0], 1.0e-2, true) == [1, 2] end using NearestNeighbors: HyperRectangle, get_min_distance_no_end, get_max_distance_no_end diff --git a/test/test_datafreetree.jl b/test/test_datafreetree.jl index f6070af..94e8aa3 100644 --- a/test/test_datafreetree.jl +++ b/test/test_datafreetree.jl @@ -14,16 +14,16 @@ using Mmap data[:] = rand(Float32, d, n) t = injectdata(DataFreeTree(typ, data), data) tr = typ(data) - for i = 1:n - @test knn(t, data[:,i], 3) == knn(tr, data[:,i], 3) + for i in 1:n + @test knn(t, data[:, i], 3) == knn(tr, data[:, i], 3) end finalize(data) end end end - data = rand(2,100) - data2 = rand(2,100) - data3 = rand(3,100) + data = rand(2, 100) + data2 = rand(2, 100) + data3 = rand(3, 100) test(data, data2, data3) test(view(data, :, :), view(data2, :, :), view(data3, :, :)) end diff --git a/test/test_inrange.jl b/test/test_inrange.jl index 83fc56d..65f6b14 100644 --- a/test/test_inrange.jl +++ b/test/test_inrange.jl @@ -1,95 +1,97 @@ -# Does not test leafsize -@testset "inrange" begin - @testset "metric" for metric in [Euclidean()] - @testset "tree type" for TreeType in trees_with_brute - function test(data) - tree = TreeType(data, metric; leafsize=2) - dosort = true - - idxs = inrange(tree, [1.1, 1.1, 1.1], 0.2, dosort) - @test idxs == [8] # Only corner 8 at least 0.2 distance away from [1.1, 1.1, 1.1] - counts = inrangecount(tree, [1.1, 1.1, 1.1], 0.2) - @test counts == 1 - - idxs = inrange(tree, [0.0, 0.0, 0.5], 0.6, dosort) - @test idxs == [1, 2] # Corner 1 and 2 at least 0.6 distance away from [0.0, 0.0, 0.5] - counts = inrangecount(tree, [0.0, 0.0, 0.5], 0.6) - @test counts == 2 - - idxs = inrange(tree, [0, 0, 0], 0.6, dosort) - @test idxs == [1] - counts = inrangecount(tree, [0, 0, 0], 0.6) - @test counts == 1 - - X = [0.0 0.0; 0.0 0.0; 0.5 0.0] - idxs1 = inrange(tree, X, 0.6, dosort) - idxs2 = inrange(tree, view(X,:,1:2), 0.6, dosort) - @test idxs1 == idxs2 - @test idxs1[1] == [1,2] - @test idxs1[2] == [1] - counts1 = inrangecount(tree, X, 0.6) - counts2 = inrangecount(tree, view(X,:,1:2), 0.6) - @test counts1 == counts2 - @test counts1 == [2, 1] - - idxs = inrange(tree, [SVector{3,Float64}(0.0, 0.0, 0.5), SVector{3,Float64}(0.0, 0.0, 0.0)], 0.6, dosort) - @test idxs[1] == [1,2] - @test idxs[2] == [1] - counts = inrangecount(tree, [SVector{3,Float64}(0.0, 0.0, 0.5), SVector{3,Float64}(0.0, 0.0, 0.0)], 0.6) - @test counts == [2, 1] - - idxs = inrange(tree, [0.33333333333, 0.33333333333, 0.33333333333], 1, dosort) - @test idxs == [1, 2, 3, 5] - counts = inrangecount(tree, [0.33333333333, 0.33333333333, 0.33333333333], 1) - @test counts == 4 - - idxs = inrange(tree, [0.5, 0.5, 0.5], 0.2, dosort) - @test idxs == [] - counts = inrangecount(tree, [0.5, 0.5, 0.5], 0.2) - @test counts == 0 - - idxs = inrange(tree, [0.5, 0.5, 0.5], 1.0, dosort) - @test idxs == [1, 2, 3, 4, 5, 6, 7, 8] - counts = inrangecount(tree, [0.5, 0.5, 0.5], 1.0) - @test counts == 8 - - @test_throws ArgumentError inrange(tree, rand(3), -0.1) - @test_throws ArgumentError inrange(tree, rand(5), 1.0) - - empty_tree = TreeType(rand(3,0), metric) - idxs = inrange(empty_tree, [0.5, 0.5, 0.5], 1.0) - @test idxs == [] - counts = inrangecount(empty_tree, [0.5, 0.5, 0.5], 1.0) - @test counts == 0 - - one_point_tree = TreeType([0.5, 0.5, 0.5], metric) - idxs = inrange(one_point_tree, data, 1.0) - @test idxs == repeat([[1]], size(data, 2)) - counts = inrangecount(one_point_tree, data, 1.0) - @test counts == repeat([1], size(data, 2)) - end - data = [0.0 0.0 0.0 0.0 1.0 1.0 1.0 1.0; - 0.0 0.0 1.0 1.0 0.0 0.0 1.0 1.0; - 0.0 1.0 0.0 1.0 0.0 1.0 0.0 1.0] # 8 node cube - test(data) - test(view(data, :, :)) - end - end -end - -@testset "view" begin - points = rand(SVector{3, Float64}, 100) - kdtree = KDTree(points) - idxs = inrange(kdtree, view(points, 1:10), 0.1) - @test idxs isa Vector{Vector{Int}} -end - -@testset "mutating" begin - for T in (KDTree, BallTree, BruteTree) - data = T(rand(3, 100)) - idxs = Vector{Int32}(undef, 0) - inrange!(idxs, data, [0.5, 0.5, 0.5], 1.0) - idxs2 = inrange(data, [0.5, 0.5, 0.5], 3) - @test idxs == idxs2 - end -end +# Does not test leafsize +@testset "inrange" begin + @testset "metric" for metric in [Euclidean()] + @testset "tree type" for TreeType in trees_with_brute + function test(data) + tree = TreeType(data, metric; leafsize = 2) + dosort = true + + idxs = inrange(tree, [1.1, 1.1, 1.1], 0.2, dosort) + @test idxs == [8] # Only corner 8 at least 0.2 distance away from [1.1, 1.1, 1.1] + counts = inrangecount(tree, [1.1, 1.1, 1.1], 0.2) + @test counts == 1 + + idxs = inrange(tree, [0.0, 0.0, 0.5], 0.6, dosort) + @test idxs == [1, 2] # Corner 1 and 2 at least 0.6 distance away from [0.0, 0.0, 0.5] + counts = inrangecount(tree, [0.0, 0.0, 0.5], 0.6) + @test counts == 2 + + idxs = inrange(tree, [0, 0, 0], 0.6, dosort) + @test idxs == [1] + counts = inrangecount(tree, [0, 0, 0], 0.6) + @test counts == 1 + + X = [0.0 0.0; 0.0 0.0; 0.5 0.0] + idxs1 = inrange(tree, X, 0.6, dosort) + idxs2 = inrange(tree, view(X, :, 1:2), 0.6, dosort) + @test idxs1 == idxs2 + @test idxs1[1] == [1, 2] + @test idxs1[2] == [1] + counts1 = inrangecount(tree, X, 0.6) + counts2 = inrangecount(tree, view(X, :, 1:2), 0.6) + @test counts1 == counts2 + @test counts1 == [2, 1] + + idxs = inrange(tree, [SVector{3, Float64}(0.0, 0.0, 0.5), SVector{3, Float64}(0.0, 0.0, 0.0)], 0.6, dosort) + @test idxs[1] == [1, 2] + @test idxs[2] == [1] + counts = inrangecount(tree, [SVector{3, Float64}(0.0, 0.0, 0.5), SVector{3, Float64}(0.0, 0.0, 0.0)], 0.6) + @test counts == [2, 1] + + idxs = inrange(tree, [0.33333333333, 0.33333333333, 0.33333333333], 1, dosort) + @test idxs == [1, 2, 3, 5] + counts = inrangecount(tree, [0.33333333333, 0.33333333333, 0.33333333333], 1) + @test counts == 4 + + idxs = inrange(tree, [0.5, 0.5, 0.5], 0.2, dosort) + @test idxs == [] + counts = inrangecount(tree, [0.5, 0.5, 0.5], 0.2) + @test counts == 0 + + idxs = inrange(tree, [0.5, 0.5, 0.5], 1.0, dosort) + @test idxs == [1, 2, 3, 4, 5, 6, 7, 8] + counts = inrangecount(tree, [0.5, 0.5, 0.5], 1.0) + @test counts == 8 + + @test_throws ArgumentError inrange(tree, rand(3), -0.1) + @test_throws ArgumentError inrange(tree, rand(5), 1.0) + + empty_tree = TreeType(rand(3, 0), metric) + idxs = inrange(empty_tree, [0.5, 0.5, 0.5], 1.0) + @test idxs == [] + counts = inrangecount(empty_tree, [0.5, 0.5, 0.5], 1.0) + @test counts == 0 + + one_point_tree = TreeType([0.5, 0.5, 0.5], metric) + idxs = inrange(one_point_tree, data, 1.0) + @test idxs == repeat([[1]], size(data, 2)) + counts = inrangecount(one_point_tree, data, 1.0) + @test counts == repeat([1], size(data, 2)) + end + data = [ + 0.0 0.0 0.0 0.0 1.0 1.0 1.0 1.0; + 0.0 0.0 1.0 1.0 0.0 0.0 1.0 1.0; + 0.0 1.0 0.0 1.0 0.0 1.0 0.0 1.0 + ] # 8 node cube + test(data) + test(view(data, :, :)) + end + end +end + +@testset "view" begin + points = rand(SVector{3, Float64}, 100) + kdtree = KDTree(points) + idxs = inrange(kdtree, view(points, 1:10), 0.1) + @test idxs isa Vector{Vector{Int}} +end + +@testset "mutating" begin + for T in (KDTree, BallTree, BruteTree) + data = T(rand(3, 100)) + idxs = Vector{Int32}(undef, 0) + inrange!(idxs, data, [0.5, 0.5, 0.5], 1.0) + idxs2 = inrange(data, [0.5, 0.5, 0.5], 3) + @test idxs == idxs2 + end +end diff --git a/test/test_knn.jl b/test/test_knn.jl index 99d96a5..6c29dbe 100644 --- a/test/test_knn.jl +++ b/test/test_knn.jl @@ -6,7 +6,7 @@ import Distances.evaluate @testset "metric" for metric in [metrics; WeightedEuclidean(ones(2))] @testset "tree type" for TreeType in trees_with_brute function test(data) - tree = TreeType(data, metric; leafsize=2) + tree = TreeType(data, metric; leafsize = 2) idxs, dists = knn(tree, [0.8, 0.8], 1) @test idxs[1] == 8 # Should be closest to top right corner @@ -21,7 +21,7 @@ import Distances.evaluate X = [0.8 0.1; 0.8 0.8] idxs1, dists1 = knn(tree, X, 1, true) - idxs2, dists2 = knn(tree, view(X,:,1:2), 1, true) + idxs2, dists2 = knn(tree, view(X, :, 1:2), 1, true) @test idxs1 == idxs2 @test dists1 == dists2 @test idxs1[1][1] == 8 @@ -31,22 +31,22 @@ import Distances.evaluate @test idxs[1] == 8 @test idxs[2] == 3 - idxs, dists = knn(tree, [SVector{2, Float64}(0.8,0.8), SVector{2, Float64}(0.1,0.8)], 1, true) + idxs, dists = knn(tree, [SVector{2, Float64}(0.8, 0.8), SVector{2, Float64}(0.1, 0.8)], 1, true) @test idxs[1][1] == 8 @test idxs[2][1] == 3 - idxs, dists = nn(tree, [SVector{2, Float64}(0.8,0.8), SVector{2, Float64}(0.1,0.8)]) + idxs, dists = nn(tree, [SVector{2, Float64}(0.8, 0.8), SVector{2, Float64}(0.1, 0.8)]) @test idxs[1] == 8 @test idxs[2] == 3 - idxs, dists = knn(tree, [1//10, 8//10], 3, true) + idxs, dists = knn(tree, [1 // 10, 8 // 10], 3, true) @test idxs == [3, 2, 5] @test_throws ArgumentError knn(tree, [0.1, 0.8], -1) # k < 0 @test_throws ArgumentError knn(tree, [0.1, 0.8], 10) # k > n_points @test_throws ArgumentError knn(tree, [0.1], 10) # n_dim != trees dim - empty_tree = TreeType(rand(2,0), metric; leafsize=2) + empty_tree = TreeType(rand(2, 0), metric; leafsize = 2) idxs, dists = knn(empty_tree, [0.5, 0.5], 0, true) @test idxs == Int[] @test_throws ArgumentError knn(empty_tree, [0.1, 0.8], -1) # k < 0 @@ -59,8 +59,10 @@ import Distances.evaluate @test_throws ArgumentError knn(one_point_tree, [0.1, 0.8], 2) # k > n_points end # 8 node rectangle - data = [0.0 0.0 0.0 0.5 0.5 1.0 1.0 1.0; - 0.0 0.5 1.0 0.0 1.0 0.0 0.5 1.0] + data = [ + 0.0 0.0 0.0 0.5 0.5 1.0 1.0 1.0; + 0.0 0.5 1.0 0.0 1.0 0.0 0.5 1.0 + ] test(data) test(view(data, :, :)) end @@ -85,7 +87,7 @@ end test(view(data, :, :)) end - data = [[0.13380863416387367, 0.7845254987714512],[0.1563342025559629, 0.7956456895676077],[0.23320094627474594, 0.9055515160266435]] + data = [[0.13380863416387367, 0.7845254987714512], [0.1563342025559629, 0.7956456895676077], [0.23320094627474594, 0.9055515160266435]] tree = KDTree(hcat(map(p -> [p[1], p[2]], data)...)) nearest, distance = knn(tree, [0.15, 0.8], 3, true, x -> x == 2) @test nearest == [1, 3] @@ -93,12 +95,12 @@ end end @testset "weighted" begin - m = WeightedEuclidean([1e-5, 1]) - data = [ + m = WeightedEuclidean([1.0e-5, 1]) + data = Float64[ 0 0 1 - 0 1 0. + 0 1 0 ] - tree = KDTree(data, m, leafsize=1) + tree = KDTree(data, m, leafsize = 1) p = [1, 0.9] @test nn(tree, p)[1] == 2 end diff --git a/test/test_monkey.jl b/test/test_monkey.jl index 4d40570..e4ce45a 100644 --- a/test/test_monkey.jl +++ b/test/test_monkey.jl @@ -18,10 +18,10 @@ import NearestNeighbors.MinkowskiMetric dim_data = rand(1:4) size_data = rand(1000:1300) data = rand(T, dim_data, size_data) - for j = 1:5 + for j in 1:5 tree = TreeType(data, metric; leafsize = rand(1:15)) n = rand(1:size_data) - idx, dist = knn(tree, data[:,n], rand(1:30), true) + idx, dist = knn(tree, data[:, n], rand(1:30), true) @test issorted(dist) == true @test n == idx[1] end