-- SPDX-License-Identifier: MPL-2.0

{- | Immutable point-cloud spatial indexes. IDs are original zero-based row
indices, duplicate points are retained, range boxes are closed, and nearest
results sort by (distance,ID). Up to 65536 points in dimensions 1..16 are
accepted (quadtrees require 2D, octrees 3D). Hash cells are real hash buckets;
KD, ball, quadrant, packed R-tree, and BVH nodes retain actual hierarchical
bounds. Query pruning is exact, with a caller leaf budget for nearest search.
-}
module HMath.Geometry.Spatial (
    SpatialKind (..),
    PointCloud,
    pointCloud,
    cloudPoints,
    SpatialIndex,
    buildSpatialIndex,
    rangeQuery,
    radiusQuery,
    NeighborQuery (..),
    nearestQuery,
) where

import qualified Data.IntMap.Strict as IntMap
import Data.List (sortOn, transpose)
import qualified Data.Map.Strict as Map
import HMath.Error (MathError (..))
import HMath.Geometry.Internal

-- | Explicit index construction algorithm; hash width must be positive finite.
data SpatialKind = SpatialHash Double | KDTree | BallTree | Quadtree | Octree | RTree | BVH
    deriving (SpatialKind -> SpatialKind -> Bool
(SpatialKind -> SpatialKind -> Bool)
-> (SpatialKind -> SpatialKind -> Bool) -> Eq SpatialKind
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: SpatialKind -> SpatialKind -> Bool
== :: SpatialKind -> SpatialKind -> Bool
$c/= :: SpatialKind -> SpatialKind -> Bool
/= :: SpatialKind -> SpatialKind -> Bool
Eq, Int -> SpatialKind -> ShowS
[SpatialKind] -> ShowS
SpatialKind -> String
(Int -> SpatialKind -> ShowS)
-> (SpatialKind -> String)
-> ([SpatialKind] -> ShowS)
-> Show SpatialKind
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> SpatialKind -> ShowS
showsPrec :: Int -> SpatialKind -> ShowS
$cshow :: SpatialKind -> String
show :: SpatialKind -> String
$cshowList :: [SpatialKind] -> ShowS
showList :: [SpatialKind] -> ShowS
Show)

-- | Validated homogeneous finite cloud; a nonempty cloud defines its dimension.
newtype PointCloud = Cloud [PointND] deriving (PointCloud -> PointCloud -> Bool
(PointCloud -> PointCloud -> Bool)
-> (PointCloud -> PointCloud -> Bool) -> Eq PointCloud
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: PointCloud -> PointCloud -> Bool
== :: PointCloud -> PointCloud -> Bool
$c/= :: PointCloud -> PointCloud -> Bool
/= :: PointCloud -> PointCloud -> Bool
Eq, Int -> PointCloud -> ShowS
[PointCloud] -> ShowS
PointCloud -> String
(Int -> PointCloud -> ShowS)
-> (PointCloud -> String)
-> ([PointCloud] -> ShowS)
-> Show PointCloud
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> PointCloud -> ShowS
showsPrec :: Int -> PointCloud -> ShowS
$cshow :: PointCloud -> String
show :: PointCloud -> String
$cshowList :: [PointCloud] -> ShowS
showList :: [PointCloud] -> ShowS
Show)

data Node = Leaf AABB [(Int, PointND)] | Branch AABB [Node] | Ball AABB [Double] Double [Node]
    deriving (Node -> Node -> Bool
(Node -> Node -> Bool) -> (Node -> Node -> Bool) -> Eq Node
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Node -> Node -> Bool
== :: Node -> Node -> Bool
$c/= :: Node -> Node -> Bool
/= :: Node -> Node -> Bool
Eq, Int -> Node -> ShowS
[Node] -> ShowS
Node -> String
(Int -> Node -> ShowS)
-> (Node -> String) -> ([Node] -> ShowS) -> Show Node
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Node -> ShowS
showsPrec :: Int -> Node -> ShowS
$cshow :: Node -> String
show :: Node -> String
$cshowList :: [Node] -> ShowS
showList :: [Node] -> ShowS
Show)

{- | Opaque checked spatial-search tree/hash index. Construction and query functions
document dimension/work limits and exactness flags.
-}
data SpatialIndex = Tree Int Node | Hash Int Double (Map.Map [Integer] [(Int, PointND)])
    deriving (SpatialIndex -> SpatialIndex -> Bool
(SpatialIndex -> SpatialIndex -> Bool)
-> (SpatialIndex -> SpatialIndex -> Bool) -> Eq SpatialIndex
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: SpatialIndex -> SpatialIndex -> Bool
== :: SpatialIndex -> SpatialIndex -> Bool
$c/= :: SpatialIndex -> SpatialIndex -> Bool
/= :: SpatialIndex -> SpatialIndex -> Bool
Eq, Int -> SpatialIndex -> ShowS
[SpatialIndex] -> ShowS
SpatialIndex -> String
(Int -> SpatialIndex -> ShowS)
-> (SpatialIndex -> String)
-> ([SpatialIndex] -> ShowS)
-> Show SpatialIndex
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> SpatialIndex -> ShowS
showsPrec :: Int -> SpatialIndex -> ShowS
$cshow :: SpatialIndex -> String
show :: SpatialIndex -> String
$cshowList :: [SpatialIndex] -> ShowS
showList :: [SpatialIndex] -> ShowS
Show)

-- | Construct a nonempty cloud of at most 65536 points, with consistent dimensions.
pointCloud :: [PointND] -> Either MathError PointCloud
pointCloud :: [PointND] -> Either MathError PointCloud
pointCloud [PointND]
ps = do
    Bool -> String -> Either MathError ()
require (Bool -> Bool
not ([PointND] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [PointND]
ps) Bool -> Bool -> Bool
&& [PointND] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [PointND]
ps Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
65536) String
"point cloud requires 1..65536 points"
    (PointND -> Either MathError ())
-> [PointND] -> Either MathError ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (PointND -> PointND -> Either MathError ()
checkPair (PointND -> [PointND] -> PointND
forall a. a -> [a] -> a
firstOr ([Double] -> PointND
PointND []) [PointND]
ps)) [PointND]
ps
    PointCloud -> Either MathError PointCloud
forall a b. b -> Either a b
Right ([PointND] -> PointCloud
Cloud [PointND]
ps)

-- | Original cloud observations.
cloudPoints :: PointCloud -> [PointND]
cloudPoints :: PointCloud -> [PointND]
cloudPoints (Cloud [PointND]
ps) = [PointND]
ps

bounds :: [(Int, PointND)] -> AABB
bounds :: [(Int, PointND)] -> AABB
bounds [(Int, PointND)]
pts =
    let columns :: [[Double]]
columns = [[Double]] -> [[Double]]
forall a. [[a]] -> [[a]]
transpose (((Int, PointND) -> [Double]) -> [(Int, PointND)] -> [[Double]]
forall a b. (a -> b) -> [a] -> [b]
map (PointND -> [Double]
coordsND (PointND -> [Double])
-> ((Int, PointND) -> PointND) -> (Int, PointND) -> [Double]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, PointND) -> PointND
forall a b. (a, b) -> b
snd) [(Int, PointND)]
pts)
     in PointND -> PointND -> AABB
AABB ([Double] -> PointND
PointND (([Double] -> Double) -> [[Double]] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map [Double] -> Double
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [[Double]]
columns)) ([Double] -> PointND
PointND (([Double] -> Double) -> [[Double]] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map [Double] -> Double
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [[Double]]
columns))
nodeBounds :: Node -> AABB
nodeBounds :: Node -> AABB
nodeBounds (Leaf AABB
b [(Int, PointND)]
_) = AABB
b
nodeBounds (Branch AABB
b [Node]
_) = AABB
b
nodeBounds (Ball AABB
b [Double]
_ Double
_ [Node]
_) = AABB
b

{- | Build the named structure. RTree uses sort-tile-recursive leaf packing
(fanout 8); BVH splits longest-axis medians; quadrant trees split cells at
their midplanes with a depth limit of 64 for coincident/near-coincident data.
-}
buildSpatialIndex :: SpatialKind -> PointCloud -> Either MathError SpatialIndex
buildSpatialIndex :: SpatialKind -> PointCloud -> Either MathError SpatialIndex
buildSpatialIndex SpatialKind
kind (Cloud [PointND]
ps) = do
    let dim :: Int
dim = [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (PointND -> [Double]
coordsND (PointND -> [PointND] -> PointND
forall a. a -> [a] -> a
firstOr ([Double] -> PointND
PointND []) [PointND]
ps)); pts :: [(Int, PointND)]
pts = [Int] -> [PointND] -> [(Int, PointND)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [PointND]
ps
    case SpatialKind
kind of
        SpatialHash Double
width -> do
            Bool -> String -> Either MathError ()
require
                (Double -> Bool
finite Double
width Bool -> Bool -> Bool
&& Double
width Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0 Bool -> Bool -> Bool
&& (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Double -> Bool
finite (Double -> Bool) -> (Double -> Double) -> Double -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
width)) ((PointND -> [Double]) -> [PointND] -> [Double]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap PointND -> [Double]
coordsND [PointND]
ps))
                String
"invalid spatial hash width/cell coordinate"
            let buckets :: Map [Integer] [(Int, PointND)]
buckets = ([(Int, PointND)] -> [(Int, PointND)] -> [(Int, PointND)])
-> [([Integer], [(Int, PointND)])]
-> Map [Integer] [(Int, PointND)]
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith [(Int, PointND)] -> [(Int, PointND)] -> [(Int, PointND)]
forall a. [a] -> [a] -> [a]
(++) [((Double -> Integer) -> [Double] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map (Double -> Integer
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor (Double -> Integer) -> (Double -> Double) -> Double -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
width)) (PointND -> [Double]
coordsND PointND
p), [(Int
i, PointND
p)]) | (Int
i, PointND
p) <- [(Int, PointND)]
pts]
            SpatialIndex -> Either MathError SpatialIndex
forall a b. b -> Either a b
Right (Int -> Double -> Map [Integer] [(Int, PointND)] -> SpatialIndex
Hash Int
dim Double
width Map [Integer] [(Int, PointND)]
buckets)
        SpatialKind
_ -> do
            Bool -> String -> Either MathError ()
require (SpatialKind
kind SpatialKind -> SpatialKind -> Bool
forall a. Eq a => a -> a -> Bool
/= SpatialKind
Quadtree Bool -> Bool -> Bool
|| Int
dim Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2) String
"quadtree requires 2D"
            Bool -> String -> Either MathError ()
require (SpatialKind
kind SpatialKind -> SpatialKind -> Bool
forall a. Eq a => a -> a -> Bool
/= SpatialKind
Octree Bool -> Bool -> Bool
|| Int
dim Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
3) String
"octree requires 3D"
            SpatialIndex -> Either MathError SpatialIndex
forall a b. b -> Either a b
Right (Int -> Node -> SpatialIndex
Tree Int
dim (Int -> [(Int, PointND)] -> Node
build Int
0 [(Int, PointND)]
pts))
  where
    build :: Int -> [(Int, PointND)] -> Node
    build :: Int -> [(Int, PointND)] -> Node
build Int
depth [(Int, PointND)]
pts
        | [(Int, PointND)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Int, PointND)]
pts Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
8 = AABB -> [(Int, PointND)] -> Node
Leaf ([(Int, PointND)] -> AABB
bounds [(Int, PointND)]
pts) [(Int, PointND)]
pts
        | Int
depth Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
64 = AABB -> [(Int, PointND)] -> Node
Leaf ([(Int, PointND)] -> AABB
bounds [(Int, PointND)]
pts) [(Int, PointND)]
pts
        | Bool
otherwise =
            let box :: AABB
box@(AABB PointND
lo PointND
hi) = [(Int, PointND)] -> AABB
bounds [(Int, PointND)]
pts
                spans :: [Double]
spans = [Double] -> [Double] -> [Double]
sub (PointND -> [Double]
coordsND PointND
hi) (PointND -> [Double]
coordsND PointND
lo)
                axis :: Int
axis = (Double, Int) -> Int
forall a b. (a, b) -> b
snd ([(Double, Int)] -> (Double, Int)
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum ([Double] -> [Int] -> [(Double, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Double]
spans [Int
0 :: Int ..]))
                ordered :: [(Int, PointND)]
ordered = ((Int, PointND) -> Double) -> [(Int, PointND)] -> [(Int, PointND)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (([Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis) ([Double] -> Double)
-> ((Int, PointND) -> [Double]) -> (Int, PointND) -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PointND -> [Double]
coordsND (PointND -> [Double])
-> ((Int, PointND) -> PointND) -> (Int, PointND) -> [Double]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, PointND) -> PointND
forall a b. (a, b) -> b
snd) [(Int, PointND)]
pts
                mid :: Int
mid = [(Int, PointND)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Int, PointND)]
pts Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
                halves :: [[(Int, PointND)]]
halves = [Int -> [(Int, PointND)] -> [(Int, PointND)]
forall a. Int -> [a] -> [a]
take Int
mid [(Int, PointND)]
ordered, Int -> [(Int, PointND)] -> [(Int, PointND)]
forall a. Int -> [a] -> [a]
drop Int
mid [(Int, PointND)]
ordered]
                recurse :: [[(Int, PointND)]] -> [Node]
recurse [[(Int, PointND)]]
groups = ([(Int, PointND)] -> Node) -> [[(Int, PointND)]] -> [Node]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> [(Int, PointND)] -> Node
build (Int
depth Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)) (([(Int, PointND)] -> Bool)
-> [[(Int, PointND)]] -> [[(Int, PointND)]]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool)
-> ([(Int, PointND)] -> Bool) -> [(Int, PointND)] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Int, PointND)] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null) [[(Int, PointND)]]
groups)
             in case SpatialKind
kind of
                    SpatialKind
BallTree ->
                        let columns :: [[Double]]
columns = [[Double]] -> [[Double]]
forall a. [[a]] -> [[a]]
transpose (((Int, PointND) -> [Double]) -> [(Int, PointND)] -> [[Double]]
forall a b. (a -> b) -> [a] -> [b]
map (PointND -> [Double]
coordsND (PointND -> [Double])
-> ((Int, PointND) -> PointND) -> (Int, PointND) -> [Double]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, PointND) -> PointND
forall a b. (a, b) -> b
snd) [(Int, PointND)]
pts)
                            center :: [Double]
center = ([Double] -> Double) -> [[Double]] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (\[Double]
xs -> [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
xs Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs)) [[Double]]
columns
                            radius :: Double
radius = [Double] -> Double
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [[Double] -> Double
norm ([Double] -> [Double] -> [Double]
sub (PointND -> [Double]
coordsND PointND
p) [Double]
center) | (Int
_, PointND
p) <- [(Int, PointND)]
pts]
                            pivot :: Int
pivot = (Double, Int) -> Int
forall a b. (a, b) -> b
snd ([(Double, Int)] -> (Double, Int)
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [([Double] -> Double
norm ([Double] -> [Double] -> [Double]
sub (PointND -> [Double]
coordsND PointND
p) [Double]
center), Int
i) | (Int
i, PointND
p) <- [(Int, PointND)]
pts])
                            chosen :: [Double]
chosen = PointND -> [Double]
coordsND ((Int, PointND) -> PointND
forall a b. (a, b) -> b
snd ((Int, PointND) -> [(Int, PointND)] -> (Int, PointND)
forall a. a -> [a] -> a
firstOr (-Int
1, [Double] -> PointND
PointND []) (((Int, PointND) -> Bool) -> [(Int, PointND)] -> [(Int, PointND)]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
pivot) (Int -> Bool) -> ((Int, PointND) -> Int) -> (Int, PointND) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, PointND) -> Int
forall a b. (a, b) -> a
fst) [(Int, PointND)]
pts)))
                            sorted :: [(Int, PointND)]
sorted = ((Int, PointND) -> Double) -> [(Int, PointND)] -> [(Int, PointND)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (\(Int
_, PointND
p) -> [Double] -> Double
norm ([Double] -> [Double] -> [Double]
sub (PointND -> [Double]
coordsND PointND
p) [Double]
chosen)) [(Int, PointND)]
pts
                         in AABB -> [Double] -> Double -> [Node] -> Node
Ball AABB
box [Double]
center Double
radius ([[(Int, PointND)]] -> [Node]
recurse [Int -> [(Int, PointND)] -> [(Int, PointND)]
forall a. Int -> [a] -> [a]
take Int
mid [(Int, PointND)]
sorted, Int -> [(Int, PointND)] -> [(Int, PointND)]
forall a. Int -> [a] -> [a]
drop Int
mid [(Int, PointND)]
sorted])
                    SpatialKind
Quadtree -> AABB -> PointND -> PointND -> [(Int, PointND)] -> Int -> Node
quadrants AABB
box PointND
lo PointND
hi [(Int, PointND)]
pts Int
depth
                    SpatialKind
Octree -> AABB -> PointND -> PointND -> [(Int, PointND)] -> Int -> Node
quadrants AABB
box PointND
lo PointND
hi [(Int, PointND)]
pts Int
depth
                    SpatialKind
RTree ->
                        let count :: Int
count = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
ceiling (Double -> Double
forall a. Floating a => a -> a
sqrt (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([(Int, PointND)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Int, PointND)]
pts) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
8) :: Double) :: Int
                            sliceSize :: Int
sliceSize = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
8 (Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
ceiling (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([(Int, PointND)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Int, PointND)]
pts) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
count :: Double))
                            axis2 :: Int
axis2 = if [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (PointND -> [Double]
coordsND PointND
lo) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1 then Int
1 else Int
0
                            tiles :: [[(Int, PointND)]]
tiles =
                                ([(Int, PointND)] -> [[(Int, PointND)]])
-> [[(Int, PointND)]] -> [[(Int, PointND)]]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap
                                    (Int -> [(Int, PointND)] -> [[(Int, PointND)]]
forall {a}. Int -> [a] -> [[a]]
chunks Int
8 ([(Int, PointND)] -> [[(Int, PointND)]])
-> ([(Int, PointND)] -> [(Int, PointND)])
-> [(Int, PointND)]
-> [[(Int, PointND)]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Int, PointND) -> Double) -> [(Int, PointND)] -> [(Int, PointND)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (([Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis2) ([Double] -> Double)
-> ((Int, PointND) -> [Double]) -> (Int, PointND) -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PointND -> [Double]
coordsND (PointND -> [Double])
-> ((Int, PointND) -> PointND) -> (Int, PointND) -> [Double]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, PointND) -> PointND
forall a b. (a, b) -> b
snd))
                                    (Int -> [(Int, PointND)] -> [[(Int, PointND)]]
forall {a}. Int -> [a] -> [[a]]
chunks Int
sliceSize (((Int, PointND) -> Double) -> [(Int, PointND)] -> [(Int, PointND)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (Double -> [Double] -> Double
forall a. a -> [a] -> a
firstOr Double
0 ([Double] -> Double)
-> ((Int, PointND) -> [Double]) -> (Int, PointND) -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PointND -> [Double]
coordsND (PointND -> [Double])
-> ((Int, PointND) -> PointND) -> (Int, PointND) -> [Double]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, PointND) -> PointND
forall a b. (a, b) -> b
snd) [(Int, PointND)]
pts))
                         in [Node] -> Node
pack [AABB -> [(Int, PointND)] -> Node
Leaf ([(Int, PointND)] -> AABB
bounds [(Int, PointND)]
tile) [(Int, PointND)]
tile | [(Int, PointND)]
tile <- [[(Int, PointND)]]
tiles]
                    SpatialKind
_ -> AABB -> [Node] -> Node
Branch AABB
box ([[(Int, PointND)]] -> [Node]
recurse [[(Int, PointND)]]
halves)
    quadrants :: AABB -> PointND -> PointND -> [(Int, PointND)] -> Int -> Node
quadrants AABB
box PointND
lo PointND
hi [(Int, PointND)]
pts Int
depth =
        let mids :: [Double]
mids = (Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2) ([Double] -> [Double] -> [Double]
plus (PointND -> [Double]
coordsND PointND
lo) (PointND -> [Double]
coordsND PointND
hi))
            groups :: [[(Int, PointND)]]
groups = Map [Bool] [(Int, PointND)] -> [[(Int, PointND)]]
forall k a. Map k a -> [a]
Map.elems (([(Int, PointND)] -> [(Int, PointND)] -> [(Int, PointND)])
-> [([Bool], [(Int, PointND)])] -> Map [Bool] [(Int, PointND)]
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith [(Int, PointND)] -> [(Int, PointND)] -> [(Int, PointND)]
forall a. [a] -> [a] -> [a]
(++) [((Double -> Double -> Bool) -> [Double] -> [Double] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
(>) (PointND -> [Double]
coordsND PointND
p) [Double]
mids, [(Int
i, PointND
p)]) | (Int
i, PointND
p) <- [(Int, PointND)]
pts])
         in if [[(Int, PointND)]] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[(Int, PointND)]]
groups Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1 then AABB -> [(Int, PointND)] -> Node
Leaf AABB
box [(Int, PointND)]
pts else AABB -> [Node] -> Node
Branch AABB
box (([(Int, PointND)] -> Node) -> [[(Int, PointND)]] -> [Node]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> [(Int, PointND)] -> Node
build (Int
depth Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)) [[(Int, PointND)]]
groups)
    pack :: [Node] -> Node
pack [] = AABB -> [(Int, PointND)] -> Node
Leaf (PointND -> PointND -> AABB
AABB ([Double] -> PointND
PointND []) ([Double] -> PointND
PointND [])) []
    pack [Node
node] = Node
node
    pack [Node]
nodes
        | [Node] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Node]
nodes Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
8 = AABB -> [Node] -> Node
Branch ([Node] -> AABB
unionBounds [Node]
nodes) [Node]
nodes
        | Bool
otherwise = [Node] -> Node
pack [AABB -> [Node] -> Node
Branch ([Node] -> AABB
unionBounds [Node]
group) [Node]
group | [Node]
group <- Int -> [Node] -> [[Node]]
forall {a}. Int -> [a] -> [[a]]
chunks Int
8 [Node]
nodes]
    unionBounds :: [Node] -> AABB
unionBounds [Node]
nodes =
        let boxes :: [AABB]
boxes = (Node -> AABB) -> [Node] -> [AABB]
forall a b. (a -> b) -> [a] -> [b]
map Node -> AABB
nodeBounds [Node]
nodes
            lows :: [[Double]]
lows = [PointND -> [Double]
coordsND PointND
low | AABB PointND
low PointND
_ <- [AABB]
boxes]
            highs :: [[Double]]
highs = [PointND -> [Double]
coordsND PointND
high | AABB PointND
_ PointND
high <- [AABB]
boxes]
         in PointND -> PointND -> AABB
AABB ([Double] -> PointND
PointND (([Double] -> Double) -> [[Double]] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map [Double] -> Double
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum ([[Double]] -> [[Double]]
forall a. [[a]] -> [[a]]
transpose [[Double]]
lows))) ([Double] -> PointND
PointND (([Double] -> Double) -> [[Double]] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map [Double] -> Double
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum ([[Double]] -> [[Double]]
forall a. [[a]] -> [[a]]
transpose [[Double]]
highs)))
    chunks :: Int -> [a] -> [[a]]
chunks Int
_ [] = []
    chunks Int
k [a]
xs = Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
k [a]
xs [a] -> [[a]] -> [[a]]
forall a. a -> [a] -> [a]
: Int -> [a] -> [[a]]
chunks Int
k (Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop Int
k [a]
xs)

checkRange :: Int -> AABB -> Either MathError ()
checkRange :: Int -> AABB -> Either MathError ()
checkRange Int
dim (AABB PointND
lo PointND
hi) =
    PointND -> PointND -> Either MathError ()
checkPair PointND
lo PointND
hi
        Either MathError () -> Either MathError () -> Either MathError ()
forall a b.
Either MathError a -> Either MathError b -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Bool -> String -> Either MathError ()
require
            ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (PointND -> [Double]
coordsND PointND
lo) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
dim Bool -> Bool -> Bool
&& [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ((Double -> Double -> Bool) -> [Double] -> [Double] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
(<=) (PointND -> [Double]
coordsND PointND
lo) (PointND -> [Double]
coordsND PointND
hi)))
            String
"invalid range box/dimension"
overlaps :: AABB -> AABB -> Bool
overlaps :: AABB -> AABB -> Bool
overlaps (AABB PointND
a PointND
b) (AABB PointND
c PointND
d) = [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ((Double -> Double -> Bool) -> [Double] -> [Double] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
(<=) (PointND -> [Double]
coordsND PointND
a) (PointND -> [Double]
coordsND PointND
d)) Bool -> Bool -> Bool
&& [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ((Double -> Double -> Bool) -> [Double] -> [Double] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
(<=) (PointND -> [Double]
coordsND PointND
c) (PointND -> [Double]
coordsND PointND
b))
inBox :: AABB -> PointND -> Bool
inBox :: AABB -> PointND -> Bool
inBox (AABB PointND
a PointND
b) PointND
p = [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ((Double -> Double -> Double -> Bool)
-> [Double] -> [Double] -> [Double] -> [Bool]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 (\Double
x Double
lo Double
hi -> Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
lo Bool -> Bool -> Bool
&& Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
hi) (PointND -> [Double]
coordsND PointND
p) (PointND -> [Double]
coordsND PointND
a) (PointND -> [Double]
coordsND PointND
b))

{- | Closed range search, sorted original IDs. Hash queries inspect occupied
bucket bounds, avoiding an unbounded Cartesian enumeration of empty cells.
-}
rangeQuery :: SpatialIndex -> AABB -> Either MathError [Int]
rangeQuery :: SpatialIndex -> AABB -> Either MathError [Int]
rangeQuery SpatialIndex
index AABB
range = do
    let dim :: Int
dim = case SpatialIndex
index of Tree Int
d Node
_ -> Int
d; Hash Int
d Double
_ Map [Integer] [(Int, PointND)]
_ -> Int
d
    Int -> AABB -> Either MathError ()
checkRange Int
dim AABB
range
    let answer :: [(Int, PointND)]
answer = case SpatialIndex
index of
            Tree Int
_ Node
root -> Node -> [(Int, PointND)]
walk Node
root
            Hash Int
_ Double
_ Map [Integer] [(Int, PointND)]
buckets -> ([(Int, PointND)] -> [(Int, PointND)])
-> [[(Int, PointND)]] -> [(Int, PointND)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (((Int, PointND) -> Bool) -> [(Int, PointND)] -> [(Int, PointND)]
forall a. (a -> Bool) -> [a] -> [a]
filter (AABB -> PointND -> Bool
inBox AABB
range (PointND -> Bool)
-> ((Int, PointND) -> PointND) -> (Int, PointND) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, PointND) -> PointND
forall a b. (a, b) -> b
snd)) (Map [Integer] [(Int, PointND)] -> [[(Int, PointND)]]
forall k a. Map k a -> [a]
Map.elems Map [Integer] [(Int, PointND)]
buckets)
    [Int] -> Either MathError [Int]
forall a b. b -> Either a b
Right ((Int -> Int) -> [Int] -> [Int]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn Int -> Int
forall a. a -> a
id (((Int, PointND) -> Int) -> [(Int, PointND)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, PointND) -> Int
forall a b. (a, b) -> a
fst [(Int, PointND)]
answer))
  where
    walk :: Node -> [(Int, PointND)]
walk Node
node | Bool -> Bool
not (AABB -> AABB -> Bool
overlaps AABB
range (Node -> AABB
nodeBounds Node
node)) = []
    walk (Leaf AABB
_ [(Int, PointND)]
pts) = ((Int, PointND) -> Bool) -> [(Int, PointND)] -> [(Int, PointND)]
forall a. (a -> Bool) -> [a] -> [a]
filter (AABB -> PointND -> Bool
inBox AABB
range (PointND -> Bool)
-> ((Int, PointND) -> PointND) -> (Int, PointND) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, PointND) -> PointND
forall a b. (a, b) -> b
snd) [(Int, PointND)]
pts
    walk (Branch AABB
_ [Node]
children) = (Node -> [(Int, PointND)]) -> [Node] -> [(Int, PointND)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Node -> [(Int, PointND)]
walk [Node]
children
    walk (Ball AABB
_ [Double]
_ Double
_ [Node]
children) = (Node -> [(Int, PointND)]) -> [Node] -> [(Int, PointND)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Node -> [(Int, PointND)]
walk [Node]
children

boxDistance :: PointND -> AABB -> Double
boxDistance :: PointND -> AABB -> Double
boxDistance PointND
p (AABB PointND
a PointND
b) =
    [Double] -> Double
norm ((Double -> Double -> Double -> Double)
-> [Double] -> [Double] -> [Double] -> [Double]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 (\Double
x Double
lo Double
hi -> Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
0 (Double -> Double -> Double
forall a. Ord a => a -> a -> a
max (Double
lo Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
x) (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
hi))) (PointND -> [Double]
coordsND PointND
p) (PointND -> [Double]
coordsND PointND
a) (PointND -> [Double]
coordsND PointND
b))

{- | A budgeted nearest-neighbor answer; False means unexplored leaves may
improve the answer. Returned distances are exact Double Euclidean distances.
-}
data NeighborQuery = NeighborQuery
    { NeighborQuery -> [(Int, Double)]
queryNeighbors :: [(Int, Double)]
    -- ^ Nearest candidate IDs paired with Euclidean distances, in increasing distance order.
    , NeighborQuery -> Int
queryLeaves :: Int
    -- ^ Number of visited spatial-index leaves charged to the query budget.
    , NeighborQuery -> Bool
queryExhaustive :: Bool
    -- ^ True only when unexplored nodes cannot improve the returned neighbors; False marks a budget-limited approximation.
    }
    deriving (NeighborQuery -> NeighborQuery -> Bool
(NeighborQuery -> NeighborQuery -> Bool)
-> (NeighborQuery -> NeighborQuery -> Bool) -> Eq NeighborQuery
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: NeighborQuery -> NeighborQuery -> Bool
== :: NeighborQuery -> NeighborQuery -> Bool
$c/= :: NeighborQuery -> NeighborQuery -> Bool
/= :: NeighborQuery -> NeighborQuery -> Bool
Eq, Int -> NeighborQuery -> ShowS
[NeighborQuery] -> ShowS
NeighborQuery -> String
(Int -> NeighborQuery -> ShowS)
-> (NeighborQuery -> String)
-> ([NeighborQuery] -> ShowS)
-> Show NeighborQuery
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> NeighborQuery -> ShowS
showsPrec :: Int -> NeighborQuery -> ShowS
$cshow :: NeighborQuery -> String
show :: NeighborQuery -> String
$cshowList :: [NeighborQuery] -> ShowS
showList :: [NeighborQuery] -> ShowS
Show)

{- | Given neighbor count k, positive leaf budget, index, then query point, search for
nearest neighbors. Require matching finite coordinate dimension and 1\<=k\<=65536. The
result reports neighbors, visited leaves and whether search was exhaustive; an
exhausted budget is not an exact-neighbor guarantee.
-}
nearestQuery :: Int -> Int -> SpatialIndex -> PointND -> Either MathError NeighborQuery
nearestQuery :: Int
-> Int -> SpatialIndex -> PointND -> Either MathError NeighborQuery
nearestQuery Int
k Int
budget SpatialIndex
index PointND
p = do
    PointND -> Either MathError ()
checkND PointND
p
    let (Int
dim, [Node]
roots) = case SpatialIndex
index of
            Tree Int
d Node
root -> (Int
d, [Node
root])
            Hash Int
d Double
_ Map [Integer] [(Int, PointND)]
buckets -> (Int
d, [AABB -> [(Int, PointND)] -> Node
Leaf ([(Int, PointND)] -> AABB
bounds [(Int, PointND)]
pts) [(Int, PointND)]
pts | [(Int, PointND)]
pts <- Map [Integer] [(Int, PointND)] -> [[(Int, PointND)]]
forall k a. Map k a -> [a]
Map.elems Map [Integer] [(Int, PointND)]
buckets])
    Bool -> String -> Either MathError ()
require
        ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (PointND -> [Double]
coordsND PointND
p) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
dim Bool -> Bool -> Bool
&& Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
65536 Bool -> Bool -> Bool
&& Int
budget Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0)
        String
"invalid nearest dimension/count/budget"
    let search :: [Node] -> [(Int, Double)] -> Int -> NeighborQuery
search [] [(Int, Double)]
best Int
used = [(Int, Double)] -> Int -> Bool -> NeighborQuery
NeighborQuery [(Int, Double)]
best Int
used Bool
True
        search [Node]
pending [(Int, Double)]
best Int
used
            | Int
used Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
budget = [(Int, Double)] -> Int -> Bool -> NeighborQuery
NeighborQuery [(Int, Double)]
best Int
used Bool
False
            | Bool
otherwise =
                let ordered :: [Node]
ordered = (Node -> Double) -> [Node] -> [Node]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (PointND -> AABB -> Double
boxDistance PointND
p (AABB -> Double) -> (Node -> AABB) -> Node -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Node -> AABB
nodeBounds) [Node]
pending
                    node :: Node
node = Node -> [Node] -> Node
forall a. a -> [a] -> a
firstOr (AABB -> [(Int, PointND)] -> Node
Leaf (PointND -> PointND -> AABB
AABB ([Double] -> PointND
PointND []) ([Double] -> PointND
PointND [])) []) [Node]
ordered
                    rest :: [Node]
rest = Int -> [Node] -> [Node]
forall a. Int -> [a] -> [a]
drop Int
1 [Node]
ordered
                    bound :: Double
bound = if [(Int, Double)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Int, Double)]
best Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
k then Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0 else (Int, Double) -> Double
forall a b. (a, b) -> b
snd ([(Int, Double)] -> (Int, Double)
forall a. HasCallStack => [a] -> a
last [(Int, Double)]
best)
                 in if Node -> Double
lower Node
node Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
bound
                        then [Node] -> [(Int, Double)] -> Int -> NeighborQuery
search [Node]
rest [(Int, Double)]
best Int
used
                        else case Node
node of
                            Leaf AABB
_ [(Int, PointND)]
pts ->
                                let candidates :: [(Int, Double)]
candidates =
                                        Int -> [(Int, Double)] -> [(Int, Double)]
forall a. Int -> [a] -> [a]
take
                                            Int
k
                                            (((Int, Double) -> (Double, Int))
-> [(Int, Double)] -> [(Int, Double)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (\(Int
i, Double
d) -> (Double
d, Int
i)) ([(Int, Double)]
best [(Int, Double)] -> [(Int, Double)] -> [(Int, Double)]
forall a. [a] -> [a] -> [a]
++ [(Int
i, [Double] -> Double
norm ([Double] -> [Double] -> [Double]
sub (PointND -> [Double]
coordsND PointND
p) (PointND -> [Double]
coordsND PointND
q))) | (Int
i, PointND
q) <- [(Int, PointND)]
pts]))
                                 in [Node] -> [(Int, Double)] -> Int -> NeighborQuery
search [Node]
rest [(Int, Double)]
candidates (Int
used Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                            Branch AABB
_ [Node]
children -> [Node] -> [(Int, Double)] -> Int -> NeighborQuery
search ([Node]
children [Node] -> [Node] -> [Node]
forall a. [a] -> [a] -> [a]
++ [Node]
rest) [(Int, Double)]
best Int
used
                            Ball AABB
_ [Double]
_ Double
_ [Node]
children -> [Node] -> [(Int, Double)] -> Int -> NeighborQuery
search ([Node]
children [Node] -> [Node] -> [Node]
forall a. [a] -> [a] -> [a]
++ [Node]
rest) [(Int, Double)]
best Int
used
        lower :: Node -> Double
lower (Ball AABB
box [Double]
center Double
radius [Node]
_) = Double -> Double -> Double
forall a. Ord a => a -> a -> a
max (PointND -> AABB -> Double
boxDistance PointND
p AABB
box) ([Double] -> Double
norm ([Double] -> [Double] -> [Double]
sub (PointND -> [Double]
coordsND PointND
p) [Double]
center) Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
radius)
        lower Node
node = PointND -> AABB -> Double
boxDistance PointND
p (Node -> AABB
nodeBounds Node
node)
    NeighborQuery -> Either MathError NeighborQuery
forall a b. b -> Either a b
Right ([Node] -> [(Int, Double)] -> Int -> NeighborQuery
search [Node]
roots [] Int
0)

-- | Exact closed radius query; uses the index range query then Euclidean filtering.
radiusQuery :: SpatialIndex -> PointND -> Double -> Either MathError [Int]
radiusQuery :: SpatialIndex -> PointND -> Double -> Either MathError [Int]
radiusQuery SpatialIndex
index PointND
p Double
radius = do
    PointND -> Either MathError ()
checkND PointND
p
    Bool -> String -> Either MathError ()
require (Double -> Bool
coord Double
radius Bool -> Bool -> Bool
&& Double
radius Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
0) String
"invalid query radius"
    let lo :: PointND
lo = [Double] -> PointND
PointND ((Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (\Double
x -> Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
radius) (PointND -> [Double]
coordsND PointND
p)); hi :: PointND
hi = [Double] -> PointND
PointND ((Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
radius) (PointND -> [Double]
coordsND PointND
p))
    ids <- SpatialIndex -> AABB -> Either MathError [Int]
rangeQuery SpatialIndex
index (PointND -> PointND -> AABB
AABB PointND
lo PointND
hi)
    let points = case SpatialIndex
index of Hash Int
_ Double
_ Map [Integer] [(Int, PointND)]
buckets -> [[(Int, PointND)]] -> [(Int, PointND)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (Map [Integer] [(Int, PointND)] -> [[(Int, PointND)]]
forall k a. Map k a -> [a]
Map.elems Map [Integer] [(Int, PointND)]
buckets); Tree Int
_ Node
root -> Node -> [(Int, PointND)]
allPoints Node
root
        indexed = [(Int, PointND)] -> IntMap PointND
forall a. [(Int, a)] -> IntMap a
IntMap.fromList [(Int, PointND)]
points
    Right
        [i | i <- ids, Just q <- [IntMap.lookup i indexed], norm (sub (coordsND p) (coordsND q)) <= radius]
  where
    allPoints :: Node -> [(Int, PointND)]
allPoints (Leaf AABB
_ [(Int, PointND)]
pts) = [(Int, PointND)]
pts
    allPoints (Branch AABB
_ [Node]
children) = (Node -> [(Int, PointND)]) -> [Node] -> [(Int, PointND)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Node -> [(Int, PointND)]
allPoints [Node]
children
    allPoints (Ball AABB
_ [Double]
_ Double
_ [Node]
children) = (Node -> [(Int, PointND)]) -> [Node] -> [(Int, PointND)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Node -> [(Int, PointND)]
allPoints [Node]
children