Search methods: handle different units (#1064)

* Search methods: handle different units * Minor adjustments * Minor adjustments --------- Co-authored-by: Júlio Hoffimann <[email protected]>
JuliaGeometry · Sep 16, 2024 · 52d8010 · 52d8010
1 parent 23a2317
commit 52d8010
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 8 deletions.
diff --git a/src/neighborsearch/ball.jl b/src/neighborsearch/ball.jl
@@ -27,9 +27,13 @@ BallSearch(geoms, ball) = BallSearch(GeometrySet(geoms), ball)
 
 function search(pₒ::Point, method::BallSearch; mask=nothing)
   tree = method.tree
-  dmax = radius(method.ball)
 
-  inds = inrange(tree, ustrip.(to(pₒ)), ustrip(dmax))
+  # adjust units of query point and radius
+  u = unit(lentype(method.domain))
+  r = ustrip(u, radius(method.ball))
+  x = ustrip.(u, to(pₒ))
+
+  inds = inrange(tree, x, r)
 
   if isnothing(mask)
     inds

diff --git a/src/neighborsearch/kball.jl b/src/neighborsearch/kball.jl
@@ -30,15 +30,18 @@ KBallSearch(geoms, k, ball) = KBallSearch(GeometrySet(geoms), k, ball)
 maxneighbors(method::KBallSearch) = method.k
 
 function searchdists!(neighbors, distances, pₒ::Point, method::KBallSearch; mask=nothing)
-  u = unit(lentype(pₒ))
   tree = method.tree
-  dmax = radius(method.ball)
   k = method.k
 
-  inds, dists = knn(tree, ustrip.(to(pₒ)), k, true)
+  # adjust units of query point and radius
+  u = unit(lentype(method.domain))
+  r = ustrip(u, radius(method.ball))
+  x = ustrip.(u, to(pₒ))
+
+  inds, dists = knn(tree, x, k, true)
 
   # keep neighbors inside ball
-  keep = dists .≤ ustrip(dmax)
+  keep = dists .≤ r
 
   # possibly mask some of the neighbors
   isnothing(mask) || (keep .*= mask[inds])

diff --git a/src/neighborsearch/knearest.jl b/src/neighborsearch/knearest.jl
@@ -28,11 +28,14 @@ KNearestSearch(geoms, k; metric=Euclidean()) = KNearestSearch(GeometrySet(geoms)
 maxneighbors(method::KNearestSearch) = method.k
 
 function searchdists!(neighbors, distances, pₒ::Point, method::KNearestSearch; mask=nothing)
-  u = unit(lentype(pₒ))
   tree = method.tree
   k = method.k
 
-  inds, dists = knn(tree, ustrip.(to(pₒ)), k, true)
+  # adjust units of query point
+  u = unit(lentype(method.domain))
+  x = ustrip.(u, to(pₒ))
+
+  inds, dists = knn(tree, x, k, true)
 
   if isnothing(mask)
     nneigh = k

diff --git a/test/neighborsearch.jl b/test/neighborsearch.jl
@@ -21,6 +21,13 @@
   n = search(cart(9, 9), s)
   @test Set(n) == Set([89, 90, 99, 100])
 
+  # different units
+  s = BallSearch(𝒟, MetricBall(T(10) * u"dm"))
+  n = search(Point(T(900) * u"cm", T(900) * u"cm"), s)
+  @test Set(n) == Set([100, 99, 90])
+  n = search(Point(T(9000) * u"mm", T(9000) * u"mm"), s)
+  @test Set(n) == Set([100, 99, 90])
+
   # non MinkowskiMetric example
   𝒟 = CartesianGrid((360, 180), T.((0.0, -90.0)), T.((1.0, 1.0)))
   s = BallSearch(𝒟, MetricBall(T(150), Haversine(T(6371))))
@@ -56,6 +63,13 @@ end
   @test nn == 3
   @test Set(n[1:nn]) == Set([100, 99, 90])
 
+  # different units
+  s = KNearestSearch(𝒟, 3)
+  n = search(Point(T(900) * u"cm", T(900) * u"cm"), s)
+  @test Set(n) == Set([100, 99, 90])
+  n = search(Point(T(9000) * u"mm", T(9000) * u"mm"), s)
+  @test Set(n) == Set([100, 99, 90])
+
   # construct from vector of geometries
   s = KNearestSearch(randpoint2(100), 3)
   @test s isa KNearestSearch
@@ -105,6 +119,13 @@ end
   @test length(n) == 4
   @test length(d) == 4
 
+  # different units
+  s = KBallSearch(𝒟, 10, MetricBall(T(10) * u"dm"))
+  n = search(Point(T(500) * u"cm", T(500) * u"cm"), s)
+  @test Set(n) == Set([56, 66, 55, 57, 46])
+  n = search(Point(T(5000) * u"mm", T(5000) * u"mm"), s)
+  @test Set(n) == Set([56, 66, 55, 57, 46])
+
   # construct from vector of geometries
   s = KBallSearch(randpoint2(100), 10, MetricBall(T(1)))
   @test s isa KBallSearch