kdt (empty) → 0.1.0
raw patch · 10 files changed
+1568/−0 lines, 10 filesdep +MonadRandomdep +QuickCheckdep +basesetup-changed
Dependencies added: MonadRandom, QuickCheck, base, criterion, deepseq, deepseq-generics, kdt, mersenne-random-pure64, pqueue
Files
- LICENSE +20/−0
- Setup.hs +2/−0
- app-src/Benchmarks/KDTBenchmark.hs +140/−0
- app-src/Tests/KdTreeTest.hs +10/−0
- kdt.cabal +69/−0
- lib-src/Data/KdMap/Dynamic.hs +336/−0
- lib-src/Data/KdMap/Static.hs +523/−0
- lib-src/Data/KdTree/Dynamic.hs +158/−0
- lib-src/Data/KdTree/Static.hs +283/−0
- lib-src/Data/Point2d.hs +27/−0
+ LICENSE view
@@ -0,0 +1,20 @@+Copyright (c) 2014 Luis G. Torres++Permission is hereby granted, free of charge, to any person obtaining+a copy of this software and associated documentation files (the+"Software"), to deal in the Software without restriction, including+without limitation the rights to use, copy, modify, merge, publish,+distribute, sublicense, and/or sell copies of the Software, and to+permit persons to whom the Software is furnished to do so, subject to+the following conditions:++The above copyright notice and this permission notice shall be included+in all copies or substantial portions of the Software.++THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,+EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.+IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY+CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,+TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE+SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+ Setup.hs view
@@ -0,0 +1,2 @@+import Distribution.Simple+main = defaultMain
+ app-src/Benchmarks/KDTBenchmark.hs view
@@ -0,0 +1,140 @@+import Data.Point2d+import Data.KdTree.Static as KDT+import Data.KdTree.Dynamic as DKDT++import Control.Monad+import qualified Control.Monad.Random as CMR+import Criterion.Main+import Data.List+import qualified Data.PQueue.Prio.Max as Q+import System.Random.Mersenne.Pure64++zeroOnePointSampler :: CMR.Rand PureMT Point2d+zeroOnePointSampler =+ liftM2 Point2d+ (CMR.getRandomR (0.0, 1.0))+ (CMR.getRandomR (0.0, 1.0))++-- Input: List of pairs of points, where first of each pair is the+-- point to add to the dynamic KdTree, and the second is the point to+-- query for nearest neighbor+interleaveBuildQuery :: [(Point2d, Point2d)] -> [Point2d]+interleaveBuildQuery =+ let f :: (DKDT.KdTree Double Point2d, [Point2d]) ->+ (Point2d, Point2d) ->+ (DKDT.KdTree Double Point2d, [Point2d])+ f (kdt, accList) (treePt, queryPt) =+ let newKdt = DKDT.insert kdt treePt+ nearest = DKDT.nearestNeighbor newKdt queryPt+ in (newKdt, nearest : accList)+ start = (DKDT.emptyKdTreeWithDistFn pointAsList2d distSqr2d, [])+ in snd . foldl' f start++-- nn implemented with optimized linear scan+nearestLinear :: [Point2d] -> Point2d -> Point2d+nearestLinear [] _ = error "nearestLinear called on an empty list!"+nearestLinear (ph : pt) query = fst $ foldl' f (ph, distSqr2d query ph) pt+ where {-# INLINE f #-}+ f b@(_, dBest) x+ | d < dBest = (x, d)+ | otherwise = b+ where d = distSqr2d query x++pointsInRadiusLinear :: [Point2d] -> Double -> Point2d -> [Point2d]+pointsInRadiusLinear ps radius query =+ filter ((<= radius * radius) . distSqr2d query) ps++-- knn implemented with priority queue+kNearestNeighborsLinear :: [Point2d] -> Int -> Point2d -> [Point2d]+kNearestNeighborsLinear ps k query = reverse $ map snd $ Q.toList $ foldl' f Q.empty ps+ where f q p = let insertBounded queue dist x+ | Q.size queue < k = Q.insert dist x queue+ | otherwise = if dist < fst (Q.findMax queue)+ then Q.insert dist x $ Q.deleteMax queue+ else queue+ in insertBounded q (distSqr2d query p) p++rangeLinear :: Point2d -> Point2d -> [Point2d] -> [Point2d]+rangeLinear lowers uppers xs =+ let lowersAsList = pointAsList2d lowers+ uppersAsList = pointAsList2d uppers+ valInRange l x u = l <= x && x <= u+ pointInRange p =+ and $ zipWith3 valInRange+ lowersAsList (pointAsList2d p) uppersAsList+ in filter pointInRange xs++pointToBounds :: Point2d -> Double -> (Point2d, Point2d)+pointToBounds (Point2d x y) w =+ (Point2d (x - w) (y - w), Point2d (x + w) (y + w))++rangeOfPointLinear :: [Point2d] -> Double -> Point2d -> [Point2d]+rangeOfPointLinear xs w q =+ let (lowers, uppers) = pointToBounds q w+ in rangeLinear lowers uppers xs++rangeOfPointKdt :: KDT.KdTree Double Point2d -> Double -> Point2d -> [Point2d]+rangeOfPointKdt kdt w q =+ let (lowers, uppers) = pointToBounds q w+ in KDT.pointsInRange kdt lowers uppers++linearInterleaveBuildQuery :: [(Point2d, Point2d)] -> [Point2d]+linearInterleaveBuildQuery =+ let f :: ([Point2d], [Point2d]) -> (Point2d, Point2d) -> ([Point2d], [Point2d])+ f (ps, accList) (structPt, queryPt) =+ let ps' = structPt : ps+ nearest = nearestLinear ps' queryPt+ in (ps', nearest : accList)+ in snd . foldl' f ([], [])++main :: IO ()+main =+ let seed = 1+ treePoints = CMR.evalRand (sequence $ repeat zeroOnePointSampler) $ pureMT seed+ kdtN n = KDT.buildKdTreeWithDistFn pointAsList2d distSqr2d $ take n treePoints+ queryPoints = CMR.evalRand (sequence $ repeat zeroOnePointSampler) $ pureMT (seed + 1)+ buildKdtBench n = bench (show n) $ nf kdtN n+ nnKdtBench nq np =+ bench ("np-" ++ show np ++ "-nq-" ++ show nq) $+ nf (map (KDT.nearestNeighbor (kdtN np))) (take nq queryPoints)+ inRadKdtBench nq r np =+ bench ("np-" ++ show np ++ "-nq-" ++ show nq ++ "-r-" ++ show r) $+ nf (map (KDT.pointsInRadius (kdtN np) r)) (take nq queryPoints)+ knnKdtBench nq k np =+ bench ("np-" ++ show np ++ "-nq-" ++ show nq ++ "-k-" ++ show k) $+ nf (map (KDT.kNearestNeighbors (kdtN np) k)) (take nq queryPoints)+ rangeKdtBench nq w np =+ bench ("np-" ++ show np ++ "-nq-" ++ show nq ++ "-w-" ++ show w) $+ nf (map $ rangeOfPointKdt (kdtN np) w) (take nq queryPoints)+ nnLinearBench nq np =+ bench ("np-" ++ show np ++ "-nq-" ++ show nq) $+ nf (map (nearestLinear (take np treePoints))) (take nq queryPoints)+ inRadLinearBench nq r np =+ bench ("np-" ++ show np ++ "-nq-" ++ show nq ++ "-r-" ++ show r) $+ nf (map $ pointsInRadiusLinear (take np treePoints) r) (take nq queryPoints)+ rangeLinearBench nq w np =+ bench ("np-" ++ show np ++ "-nq-" ++ show nq ++ "-w-" ++ show w) $+ nf (map $ rangeOfPointLinear (take np treePoints) w) (take nq queryPoints)+ knnLinearBench nq k np =+ bench ("np-" ++ show np ++ "-nq-" ++ show nq ++ "-k-" ++ show k) $+ nf (map $ kNearestNeighborsLinear (take np treePoints) k) (take nq queryPoints)+ nniDkdtBench n =+ bench ("n-" ++ show n) $+ nf interleaveBuildQuery (zip (take n treePoints) (take n queryPoints))+ numQueries = 100+ pointSetSizes = [100, 1000, 10000, 100000]+ radius = 0.05+ numNeighbors = 10+ rangeWidth = 0.05+ in defaultMain [+ bgroup "linear-nn" $ map (nnLinearBench numQueries) pointSetSizes,+ bgroup "linear-rad" $ map (inRadLinearBench numQueries radius) pointSetSizes,+ bgroup "linear-knn" $ map (knnLinearBench numQueries numNeighbors) pointSetSizes,+ bgroup "linear-range" $ map (rangeLinearBench numQueries rangeWidth) pointSetSizes,+ bgroup "kdt-build" $ map buildKdtBench pointSetSizes,+ bgroup "kdt-nn" $ map (nnKdtBench numQueries) pointSetSizes,+ bgroup "kdt-rad" $ map (inRadKdtBench numQueries radius) pointSetSizes,+ bgroup "kdt-knn" $ map (knnKdtBench numQueries numNeighbors) pointSetSizes,+ bgroup "kdt-range" $ map (rangeKdtBench numQueries rangeWidth) pointSetSizes,+ bgroup "dkdt-nn" $ map nniDkdtBench pointSetSizes+ ]
+ app-src/Tests/KdTreeTest.hs view
@@ -0,0 +1,10 @@+import Data.KdMap.Static as KDM+import Data.KdMap.Dynamic as DKDM++import Control.Monad+import System.Exit++main :: IO ()+main = do+ success <- liftM2 (&&) KDM.runTests DKDM.runTests+ unless success exitFailure
+ kdt.cabal view
@@ -0,0 +1,69 @@+-- Initial kdt.cabal generated by cabal init. For further documentation,+-- see http://haskell.org/cabal/users-guide/++name: kdt+version: 0.1.0+synopsis: Fast and flexible k-d trees for various types of point queries.+description: This package includes static and dynamic versions of k-d trees,+ as well as \"Map\" variants that store data at each point in the+ k-d tree structure. Supports nearest neighbor,+ k nearest neighbors, points within a given radius, and points+ within a given range.+ To learn to use this package, start with the documentation for+ the "Data.KdTree.Static" module.+homepage: https://github.com/giogadi/kdt+license: MIT+license-file: LICENSE+author: Luis G. Torres+maintainer: lgtorres42@gmail.com+copyright: Luis G. Torres, 2014+category: Data+build-type: Simple+-- extra-source-files:+cabal-version: >=1.10+source-repository head+ type: git+ location: https://github.com/giogadi/kdt.git+ branch: master++library+ exposed-modules: Data.KdMap.Static,+ Data.KdTree.Static,+ Data.KdMap.Dynamic,+ Data.KdTree.Dynamic,+ Data.Point2d+ -- other-modules:+ other-extensions: DeriveGeneric, TemplateHaskell+ ghc-options: -Wall -O3+ ghc-prof-options: -Wall -O3 -fprof-auto+ build-depends: base >=4.6 && <4.8,+ deepseq >=1.3 && <1.4,+ QuickCheck >=2.7 && <2.8,+ pqueue >=1.2.1 && <1.3,+ deepseq-generics >=0.1.1.1+ hs-source-dirs: lib-src+ default-language: Haskell2010++Test-Suite KdTreeTest+ type: exitcode-stdio-1.0+ main-is: Tests/KdTreeTest.hs+ hs-source-dirs: app-src+ ghc-options: -Wall -O3+ build-depends: base >=4.6 && <4.8,+ kdt -any+ default-language: Haskell2010++benchmark KDTBenchmark+ type: exitcode-stdio-1.0+ main-is: Benchmarks/KDTBenchmark.hs+ hs-source-dirs: app-src+ ghc-options: -Wall -O3+ ghc-prof-options: -Wall -O3 -fprof-auto+ "-with-rtsopts=-p"+ build-depends: base >=4.6 && <4.8,+ kdt -any,+ MonadRandom >= 0.1.12 && <0.2,+ mersenne-random-pure64 >=0.2.0.4 && <0.3,+ criterion >= 1.0.0.0 && <1.1,+ pqueue >=1.2.1 && <1.3+ default-language: Haskell2010
+ lib-src/Data/KdMap/Dynamic.hs view
@@ -0,0 +1,336 @@+{-# LANGUAGE DeriveGeneric, TemplateHaskell #-}++module Data.KdMap.Dynamic+ ( -- * Usage++ -- $usage++ -- * Reference++ -- ** Types+ PointAsListFn+ , SquaredDistanceFn+ , KdMap+ -- ** Dynamic /k/-d map construction+ , emptyKdMap+ , singleton+ , emptyKdMapWithDistFn+ , singletonWithDistFn+ -- ** Insertion+ , insert+ , batchInsert+ -- ** Query+ , nearestNeighbor+ , pointsInRadius+ , kNearestNeighbors+ , pointsInRange+ , assocs+ , points+ , values+ , null+ , size+ -- ** Folds+ , foldrKdMap+ -- ** Utilities+ , defaultDistSqrFn+ , runTests+ ) where++import Prelude hiding (null)++import Control.Applicative+import Data.Bits+import Data.Foldable+import Data.Function+import Data.List as L hiding (insert, null)+import qualified Data.List (null)+import Data.Traversable++import Control.DeepSeq+import Control.DeepSeq.Generics (genericRnf)+import GHC.Generics+import Test.QuickCheck hiding ((.&.))++import Data.Point2d+import qualified Data.KdMap.Static as KDM+import Data.KdMap.Static (PointAsListFn, SquaredDistanceFn, defaultDistSqrFn)++-- $usage+--+-- The 'KdMap' is a variant of+-- @Data.KdTree.Dynamic.@'Data.KdTree.Dynamic.KdTree' where each point+-- in the tree is associated with some data. It is the dynamic variant+-- of @Data.KdMap.Static.@'Data.KdMap.Static.KdMap'.+--+-- Here's an example of interleaving point-value insertions and point+-- queries using 'KdMap', where points are 3D points and values are+-- 'String's:+--+-- @+-- >>> let dkdm = singleton point3dAsList ((Point3D 0.0 0.0 0.0), \"First\")+--+-- >>> let dkdm' = insert dkdm ((Point3D 1.0 1.0 1.0), \"Second\")+--+-- >>> nearestNeighbor dkdm' (Point3D 0.4 0.4 0.4)+-- (Point3D {x = 0.0, y = 0.0, z = 0.0}, \"First\")+--+-- >>> let dkdm'' = insert dkdm' ((Point3D 0.5 0.5 0.5), \"Third\")+--+-- >>> nearestNeighbor dkdm'' (Point3D 0.4 0.4 0.4)+-- (Point3D {x = 0.5, y = 0.5, z = 0.5}, \"Third\")+-- @++-- | A dynamic /k/-d tree structure that stores points of type @p@+-- with axis values of type @a@. Additionally, each point is+-- associated with a value of type @v@.+data KdMap a p v = KdMap+ { _trees :: [KDM.KdMap a p v]+ , _pointAsList :: PointAsListFn a p+ , _distSqr :: SquaredDistanceFn a p+ , _numNodes :: Int+ } deriving Generic+instance (NFData a, NFData p, NFData v) => NFData (KdMap a p v) where rnf = genericRnf++instance Functor (KdMap a p) where+ fmap f dkdMap = dkdMap { _trees = map (fmap f) $ _trees dkdMap }++-- | Performs a foldr over each point-value pair in the 'KdMap'.+foldrKdMap :: ((p, v) -> b -> b) -> b -> KdMap a p v -> b+foldrKdMap f z dkdMap = L.foldr (flip $ KDM.foldrKdMap f) z $ _trees dkdMap++instance Foldable (KdMap a p) where+ foldr f = foldrKdMap (f . snd)++instance Traversable (KdMap a p) where+ traverse f (KdMap t p d n) =+ KdMap <$> traverse (traverse f) t <*> pure p <*> pure d <*> pure n++-- | Generates an empty 'KdMap' with a user-specified distance function.+emptyKdMapWithDistFn :: PointAsListFn a p -> SquaredDistanceFn a p -> KdMap a p v+emptyKdMapWithDistFn p2l d2 = KdMap [] p2l d2 0++-- | Returns whether the 'KdMap' is empty.+null :: KdMap a p v -> Bool+null (KdMap [] _ _ _) = True+null _ = False++-- | Generates a 'KdMap' with a single point-value pair using a+-- user-specified distance function.+singletonWithDistFn :: Real a => PointAsListFn a p -> SquaredDistanceFn a p -> (p, v) -> KdMap a p v+singletonWithDistFn p2l d2 (k, v) =+ KdMap [KDM.buildKdMapWithDistFn p2l d2 [(k, v)]] p2l d2 1++-- | Generates an empty 'KdMap' with the default distance function.+emptyKdMap :: Real a => PointAsListFn a p -> KdMap a p v+emptyKdMap p2l = emptyKdMapWithDistFn p2l $ defaultDistSqrFn p2l++-- | Generates a 'KdMap' with a single point-value pair using the+-- default distance function.+singleton :: Real a => PointAsListFn a p -> (p, v) -> KdMap a p v+singleton p2l = singletonWithDistFn p2l $ defaultDistSqrFn p2l++-- | Adds a given point-value pair to a 'KdMap'.+--+-- Average time complexity per insert for /n/ inserts: /O(log^2(n))/.+insert :: Real a => KdMap a p v -> p -> v -> KdMap a p v+insert (KdMap trees p2l d2 n) k v =+ let bitList = map ((1 .&.) . (n `shiftR`)) [0..]+ (onesPairs, theRestPairs) = span ((== 1) . fst) $ zip bitList trees+ ((_, ones), (_, theRest)) = (unzip onesPairs, unzip theRestPairs)+ newTree = KDM.buildKdMapWithDistFn p2l d2 $ (k, v) : L.concatMap KDM.assocs ones+ in KdMap (newTree : theRest) p2l d2 $ n + 1++-- | Given a 'KdMap' and a query point, returns the point-value pair in+-- the 'KdMap' with the point nearest to the query.+--+-- Average time complexity: /O(log^2(n))/.+nearestNeighbor :: Real a => KdMap a p v -> p -> (p, v)+nearestNeighbor (KdMap ts _ d2 _) query =+ let nearests = map (`KDM.nearestNeighbor` query) ts+ in if Data.List.null nearests+ then error "Called nearestNeighbor on empty KdMap."+ else L.minimumBy (compare `on` (d2 query . fst)) nearests++insertPair :: Real a => KdMap a p v -> (p, v) -> KdMap a p v+insertPair t = uncurry (insert t)++-- | Given a 'KdMap', a query point, and a number @k@, returns the+-- @k@ point-value pairs with the nearest points to the query.+--+-- Neighbors are returned in order of increasing distance from query+-- point.+--+-- Average time complexity: /log(k) * log^2(n)/ for /k/ nearest+-- neighbors on a structure with /n/ data points.+--+-- Worst case time complexity: /n * log(k)/ for /k/ nearest neighbors+-- on a structure with /n/ data points.+kNearestNeighbors :: Real a => KdMap a p v -> Int -> p -> [(p, v)]+kNearestNeighbors (KdMap trees _ d2 _) k query =+ let neighborSets = map (\t -> KDM.kNearestNeighbors t k query) trees+ in take k $ L.foldr merge [] neighborSets+ where merge [] ys = ys+ merge xs [] = xs+ merge xs@(x:xt) ys@(y:yt)+ | distX <= distY = x : merge xt ys+ | otherwise = y : merge xs yt+ where distX = d2 query $ fst x+ distY = d2 query $ fst y++-- | Given a 'KdMap', a query point, and a radius, returns all+-- point-value pairs in the 'KdTree' with points within the given+-- radius of the query point.+--+-- Points are not returned in any particular order.+--+-- Worst case time complexity: /O(n)/ for /n/ data points.+pointsInRadius :: Real a => KdMap a p v -> a -> p -> [(p, v)]+pointsInRadius (KdMap trees _ _ _) radius query =+ L.concatMap (\t -> KDM.pointsInRadius t radius query) trees++-- | Finds all point-value pairs in a 'KdMap' with points within a+-- given range, where the range is specified as a set of lower and+-- upper bounds.+--+-- Points are not returned in any particular order.+--+-- Worst case time complexity: /O(n)/ for n data points and a range+-- that spans all the points.+pointsInRange :: Real a => KdMap a p v+ -> p -- ^ lower bounds of range+ -> p -- ^ upper bounds of range+ -> [(p, v)] -- ^ point-value pairs within+ -- given range+pointsInRange (KdMap trees _ _ _) lowers uppers =+ L.concatMap (\t -> KDM.pointsInRange t lowers uppers) trees++-- | Returns the number of elements in the 'KdMap'.+--+-- Time complexity: /O(1)/+size :: KdMap a p v -> Int+size (KdMap _ _ _ n) = n++-- | Returns a list of all the point-value pairs in the 'KdMap'.+--+-- Time complexity: /O(n)/ for /n/ data points.+assocs :: KdMap a p v -> [(p, v)]+assocs (KdMap trees _ _ _) = L.concatMap KDM.assocs trees++-- | Returns all points in the 'KdMap'.+--+-- Time complexity: /O(n)/ for /n/ data points.+points :: KdMap a p v -> [p]+points = map fst . assocs++-- | Returns all values in the 'KdMap'.+--+-- Time complexity: /O(n)/ for /n/ data points.+values :: KdMap a p v -> [v]+values = map snd . assocs++-- | Inserts a list of point-value pairs into the 'KdMap'.+batchInsert :: Real a => KdMap a p v -> [(p, v)] -> KdMap a p v+-- TODO: This can be made far more efficient by batch-creating the+-- individual KdMaps before placing them into the KdMap+batchInsert = L.foldl' insertPair++--------------------------------------------------------------------------------+-- Tests+--------------------------------------------------------------------------------++testElements :: [p] -> [(p, Int)]+testElements ps = zip ps [1 ..]++checkLogNTrees :: Real a => PointAsListFn a p -> SquaredDistanceFn a p -> [p] -> Bool+checkLogNTrees p2l d2 ps =+ let lengthIsLogN (KdMap ts _ _ n) = length ts == popCount n+ in L.all lengthIsLogN $ scanl insertPair (emptyKdMapWithDistFn p2l d2) $ testElements ps++prop_logNTrees :: [Point2d] -> Bool+prop_logNTrees = checkLogNTrees pointAsList2d distSqr2d++checkTreeSizesPowerOf2 :: Real a => PointAsListFn a p ->+ SquaredDistanceFn a p ->+ [p] ->+ Bool+checkTreeSizesPowerOf2 p2l d2 ps =+ let sizesPowerOf2 (KdMap ts _ _ _) = L.all (== 1) $ map (popCount . length . KDM.assocs) ts+ in L.all sizesPowerOf2 $ scanl insertPair (emptyKdMapWithDistFn p2l d2) $ testElements ps++prop_treeSizesPowerOf2 :: [Point2d] -> Bool+prop_treeSizesPowerOf2 = checkTreeSizesPowerOf2 pointAsList2d distSqr2d++checkNumElements :: Real a => PointAsListFn a p -> SquaredDistanceFn a p -> [p] -> Bool+checkNumElements p2l d2 ps =+ let numsMatch (num, KdMap ts _ _ n) = n == num && n == L.sum (map (length . KDM.assocs) ts)+ in L.all numsMatch $ zip [0..] $ scanl insertPair (emptyKdMapWithDistFn p2l d2) $ testElements ps++prop_validNumElements :: [Point2d] -> Bool+prop_validNumElements = checkNumElements pointAsList2d distSqr2d++checkNearestEqualToBatch :: (Eq p, Real a) => PointAsListFn a p ->+ SquaredDistanceFn a p ->+ ([p], p) ->+ Bool+checkNearestEqualToBatch p2l d2 (ps, query) =+ let kdt = KDM.buildKdMapWithDistFn p2l d2 $ testElements ps+ kdtAnswer = KDM.nearestNeighbor kdt query+ dkdt = batchInsert (emptyKdMapWithDistFn p2l d2) $ testElements ps+ dkdtAnswer = nearestNeighbor dkdt query+ in dkdtAnswer == kdtAnswer++prop_nearestEqualToBatch :: Point2d -> Property+prop_nearestEqualToBatch query =+ forAll (listOf1 arbitrary) $ \xs ->+ checkNearestEqualToBatch pointAsList2d distSqr2d (xs, query)++checkKNearestEqualToBatch :: (Eq p, Real a) => PointAsListFn a p ->+ SquaredDistanceFn a p ->+ ([p], Int, p) ->+ Bool+checkKNearestEqualToBatch p2l d2 (ps, k, query) =+ let kdt = KDM.buildKdMapWithDistFn p2l d2 $ testElements ps+ kdtAnswer = KDM.kNearestNeighbors kdt k query+ dkdt = batchInsert (emptyKdMapWithDistFn p2l d2) $ testElements ps+ dkdtAnswer = kNearestNeighbors dkdt k query+ in dkdtAnswer == kdtAnswer++prop_kNearestEqualToBatch :: Point2d -> Property+prop_kNearestEqualToBatch query =+ forAll (listOf1 arbitrary) $ \xs ->+ forAll (choose (1, length xs)) $ \k ->+ checkKNearestEqualToBatch pointAsList2d distSqr2d (xs, k, query)++checkInRadiusEqualToBatch :: (Ord p, Real a) => PointAsListFn a p ->+ SquaredDistanceFn a p ->+ ([p], a, p) ->+ Bool+checkInRadiusEqualToBatch p2l d2 (ps, radius, query) =+ let kdt = KDM.buildKdMapWithDistFn p2l d2 $ testElements ps+ kdtAnswer = KDM.pointsInRadius kdt radius query+ dkdt = batchInsert (emptyKdMapWithDistFn p2l d2) $ testElements ps+ dkdtAnswer = pointsInRadius dkdt radius query+ in sort dkdtAnswer == sort kdtAnswer++prop_checkInRadiusEqualToBatch :: Point2d -> Property+prop_checkInRadiusEqualToBatch query =+ forAll (listOf1 arbitrary) $ \xs ->+ forAll (choose (0.0, 1000.0)) $ \radius ->+ checkInRadiusEqualToBatch pointAsList2d distSqr2d (xs, radius, query)++prop_checkInRangeEqualToBatch :: ([Point2d], Point2d, Point2d) -> Bool+prop_checkInRangeEqualToBatch ([], _, _) = True+prop_checkInRangeEqualToBatch (xs, lowers, uppers)+ | L.and $ zipWith (<) (pointAsList2d lowers) (pointAsList2d uppers) =+ let kdt = KDM.buildKdMapWithDistFn pointAsList2d distSqr2d $ testElements xs+ kdtAnswer = KDM.pointsInRange kdt lowers uppers+ dkdt = batchInsert (emptyKdMapWithDistFn pointAsList2d distSqr2d) $ testElements xs+ dkdtAnswer = pointsInRange dkdt lowers uppers+ in sort dkdtAnswer == sort kdtAnswer+ | otherwise = True+++-- Run all tests+return []+runTests :: IO Bool+runTests = $quickCheckAll
+ lib-src/Data/KdMap/Static.hs view
@@ -0,0 +1,523 @@+{-# LANGUAGE DeriveGeneric, TemplateHaskell #-}++module Data.KdMap.Static+ ( -- * Usage++ -- $usage++ -- * Reference++ -- ** Types+ PointAsListFn+ , SquaredDistanceFn+ , KdMap+ -- ** /k/-d map construction+ , buildKdMap+ , buildKdMapWithDistFn+ -- ** Query+ , nearestNeighbor+ , pointsInRadius+ , kNearestNeighbors+ , pointsInRange+ , assocs+ , points+ , values+ , size+ -- ** Folds+ , foldrKdMap+ -- ** Utilities+ , defaultDistSqrFn+ , runTests+ ) where++import Control.DeepSeq+import Control.DeepSeq.Generics (genericRnf)+import GHC.Generics++import Control.Applicative+import Data.Foldable+import Data.Function+import qualified Data.List as L+import Data.Maybe+import Data.Ord+import qualified Data.PQueue.Prio.Max as Q+import Data.Traversable+import Test.QuickCheck++import Data.Point2d++-- $usage+--+-- The 'KdMap' is a variant of 'Data.KdTree.Static.KdTree' where each point in+-- the tree is associated with some data. When talking about 'KdMap's,+-- we'll refer to the points and their associated data as the /points/+-- and /values/ of the 'KdMap', respectively. It might help to think+-- of 'Data.KdTree.Static.KdTree' and 'KdMap' as being analogous to+-- 'Data.Set' and 'Data.Map'.+--+-- Suppose you wanted to perform point queries on a set of 3D points,+-- where each point is associated with a 'String'. Here's how to build+-- a 'KdMap' of the data and perform a nearest neighbor query (if this+-- doesn't make sense, start with the documentation for+-- 'Data.KdTree.Static.KdTree'):+--+-- @+-- >>> let points = [(Point3d 0.0 0.0 0.0), (Point3d 1.0 1.0 1.0)]+--+-- >>> let valueStrings = [\"First\", \"Second\"]+--+-- >>> let pointValuePairs = zip points valueStrings+--+-- >>> let kdm = buildKdMap point3dAsList pointValuePairs+--+-- >>> nearestNeighbor kdm (Point3d 0.1 0.1 0.1)+-- [Point3d {x = 0.0, y = 0.0, z = 0.0}, \"First\"]+-- @++data TreeNode a p v = TreeNode { _treeLeft :: TreeNode a p v+ , _treePoint :: (p, v)+ , _axisValue :: a+ , _treeRight :: TreeNode a p v+ } |+ Empty+ deriving Generic+instance (NFData a, NFData p, NFData v) => NFData (TreeNode a p v) where rnf = genericRnf++mapTreeNode :: (v1 -> v2) -> TreeNode a p v1 -> TreeNode a p v2+mapTreeNode _ Empty = Empty+mapTreeNode f (TreeNode left (k, v) axisValue right) =+ TreeNode (mapTreeNode f left) (k, f v) axisValue (mapTreeNode f right)++-- | Converts a point of type @p@ with axis values of type+-- @a@ into a list of axis values [a].+type PointAsListFn a p = p -> [a]++-- | Returns the squared distance between two points of type+-- @p@ with axis values of type @a@.+type SquaredDistanceFn a p = p -> p -> a++-- | A /k/-d tree structure that stores points of type @p@ with axis+-- values of type @a@. Additionally, each point is associated with a+-- value of type @v@.+data KdMap a p v = KdMap { _pointAsList :: PointAsListFn a p+ , _distSqr :: SquaredDistanceFn a p+ , _rootNode :: TreeNode a p v+ , _size :: Int+ } deriving Generic+instance (NFData a, NFData p, NFData v) => NFData (KdMap a p v) where rnf = genericRnf++instance Functor (KdMap a p) where+ fmap f kdMap = kdMap { _rootNode = mapTreeNode f (_rootNode kdMap) }++foldrTreeNode :: ((p, v) -> b -> b) -> b -> TreeNode a p v -> b+foldrTreeNode _ z Empty = z+foldrTreeNode f z (TreeNode left p _ right) =+ foldrTreeNode f (f p (foldrTreeNode f z right)) left++-- | Performs a foldr over each point-value pair in the 'KdMap'.+foldrKdMap :: ((p, v) -> b -> b) -> b -> KdMap a p v -> b+foldrKdMap f z (KdMap _ _ r _) = foldrTreeNode f z r++instance Foldable (KdMap a p) where+ foldr f = foldrKdMap (f . snd)++traverseTreeNode :: Applicative f => (b -> f c) -> TreeNode a p b -> f (TreeNode a p c)+traverseTreeNode _ Empty = pure Empty+traverseTreeNode f (TreeNode l (p, v) axisValue r) =+ TreeNode <$>+ traverseTreeNode f l <*>+ ((,) p <$> f v) <*> -- would simply be traverse f (p, v), but+ -- base-4.6.* doesn't have a Traversable+ -- instance for tuples.+ pure axisValue <*>+ traverseTreeNode f r++instance Traversable (KdMap a p) where+ traverse f (KdMap p d r n) =+ KdMap <$> pure p <*> pure d <*> traverseTreeNode f r <*> pure n++quickselect :: (b -> b -> Ordering) -> Int -> [b] -> b+quickselect cmp = go+ where go _ [] = error "quickselect must be called on a non-empty list."+ go k (x:xs) | k < l = go k ys+ | k > l = go (k - l - 1) zs+ | otherwise = x+ where (ys, zs) = L.partition ((== LT) . (`cmp` x)) xs+ l = length ys++-- | Builds a 'KdMap' from a list of pairs of points (of type p) and+-- values (of type v), using a user-specified squared distance+-- function.+--+-- Average time complexity: /O(n * log(n))/ for /n/ data points.+--+-- Worst case time complexity: /O(n^2)/ for /n/ data points.+--+-- Worst case space complexity: /O(n)/ for /n/ data points.+--+-- Throws an error if given an empty list of data points.+buildKdMapWithDistFn :: Real a => PointAsListFn a p ->+ SquaredDistanceFn a p ->+ [(p, v)] ->+ KdMap a p v+buildKdMapWithDistFn _ _ [] = error "KdMap must be built with a non-empty list."+buildKdMapWithDistFn pointAsList distSqr dataPoints =+ let axisValsPointsPairs = zip (map (cycle . pointAsList . fst) dataPoints) dataPoints+ in KdMap { _pointAsList = pointAsList+ , _distSqr = distSqr+ , _rootNode = buildTreeInternal axisValsPointsPairs+ , _size = length dataPoints+ }+ where buildTreeInternal [] = Empty+ buildTreeInternal ps =+ let n = length ps+ (medianAxisVal : _, _) =+ quickselect (comparing (head . fst)) (n `div` 2) ps+ f ([], _) _ = error "buildKdMap.f: no empty lists allowed!"+ f (v : vt, p) (lt, maybeMedian, gt)+ | v < medianAxisVal = ((vt, p) : lt, maybeMedian, gt)+ | v > medianAxisVal = (lt, maybeMedian, (vt, p) : gt)+ | otherwise =+ case maybeMedian of+ Nothing -> (lt, Just p, gt)+ Just _ -> ((vt, p) : lt, maybeMedian, gt)+ (leftPoints, maybeMedianPt, rightPoints) = L.foldr f ([], Nothing, []) ps+ in TreeNode+ { _treeLeft = buildTreeInternal leftPoints+ , _treePoint = fromJust maybeMedianPt+ , _axisValue = medianAxisVal+ , _treeRight = buildTreeInternal rightPoints+ }++-- | A default implementation of squared distance given two points and+-- a 'PointAsListFn'.+defaultDistSqrFn :: Num a => PointAsListFn a p -> SquaredDistanceFn a p+defaultDistSqrFn pointAsList k1 k2 =+ L.sum $ map (^ (2 :: Int)) $ zipWith (-) (pointAsList k1) (pointAsList k2)++-- | Builds a 'KdTree' from a list of pairs of points (of type p) and+-- values (of type v) using a default squared distance function+-- 'defaultDistSqrFn'.+--+-- Average complexity: /O(n * log(n))/ for /n/ data points.+--+-- Worst case time complexity: /O(n^2)/ for /n/ data points.+--+-- Worst case space complexity: /O(n)/ for /n/ data points.+--+-- Throws an error if given an empty list of data points.+buildKdMap :: Real a => PointAsListFn a p -> [(p, v)] -> KdMap a p v+buildKdMap pointAsList =+ buildKdMapWithDistFn pointAsList $ defaultDistSqrFn pointAsList++assocsInternal :: TreeNode a p v -> [(p, v)]+assocsInternal t = go t []+ where go Empty = id+ go (TreeNode l p _ r) = go l . (p :) . go r++-- | Returns a list of all the point-value pairs in the 'KdMap'.+--+-- Time complexity: /O(n)/ for /n/ data points.+assocs :: KdMap a p v -> [(p, v)]+assocs (KdMap _ _ t _) = assocsInternal t++-- | Returns all points in the 'KdMap'.+--+-- Time complexity: /O(n)/ for /n/ data points.+points :: KdMap a p v -> [p]+points = map fst . assocs++-- | Returns all values in the 'KdMap'.+--+-- Time complexity: /O(n)/ for /n/ data points.+values :: KdMap a p v -> [v]+values = map snd . assocs++-- | Given a 'KdMap' and a query point, returns the point-value pair+-- in the 'KdMap' with the point nearest to the query.+--+-- Average time complexity: /O(log(n))/ for /n/ data points.+--+-- Worst case time complexity: /O(n)/ for /n/ data points.+nearestNeighbor :: Real a => KdMap a p v -> p -> (p, v)+nearestNeighbor (KdMap _ _ Empty _) _ =+ error "nearestNeighbor: why is there an empty KdMap?"+nearestNeighbor (KdMap pointAsList distSqr t@(TreeNode _ root _ _) _) query =+ -- This is an ugly way to kickstart the function but it's faster+ -- than using a Maybe.+ fst $ go (root, distSqr (fst root) query) (cycle $ pointAsList query) t+ where+ go _ [] _ = error "nearestNeighbor.go: no empty lists allowed!"+ go bestSoFar _ Empty = bestSoFar+ go bestSoFar+ (queryAxisValue : qvs)+ (TreeNode left (nodeK, nodeV) nodeAxisVal right) =+ let better x1@(_, dist1) x2@(_, dist2) = if dist1 < dist2+ then x1+ else x2+ currDist = distSqr query nodeK+ bestAfterNode = better ((nodeK, nodeV), currDist) bestSoFar+ nearestInTree onsideSubtree offsideSubtree =+ let bestAfterOnside = go bestAfterNode qvs onsideSubtree+ checkOffsideSubtree =+ (queryAxisValue - nodeAxisVal)^(2 :: Int) < snd bestAfterOnside+ in if checkOffsideSubtree+ then go bestAfterOnside qvs offsideSubtree+ else bestAfterOnside+ in if queryAxisValue <= nodeAxisVal+ then nearestInTree left right+ else nearestInTree right left++-- | Given a 'KdMap', a query point, and a radius, returns all+-- point-value pairs in the 'KdMap' with points within the given+-- radius of the query point.+--+-- Points are not returned in any particular order.+--+-- Worst case time complexity: /O(n)/ for /n/ data points and a radius+-- that spans all points in the structure.+pointsInRadius :: Real a => KdMap a p v+ -> a -- ^ radius+ -> p -- ^ query point+ -> [(p, v)] -- ^ list of point-value pairs+ -- with points within given+ -- radius of query+pointsInRadius (KdMap pointAsList distSqr t _) radius query =+ go (cycle $ pointAsList query) t []+ where+ go [] _ _ = error "pointsInRadius.go: no empty lists allowed!"+ go _ Empty acc = acc+ go (queryAxisValue : qvs) (TreeNode left (k, v) nodeAxisVal right) acc =+ let onTheLeft = queryAxisValue <= nodeAxisVal+ accAfterOnside = if onTheLeft+ then go qvs left acc+ else go qvs right acc+ accAfterOffside = if abs (queryAxisValue - nodeAxisVal) < radius+ then if onTheLeft+ then go qvs right accAfterOnside+ else go qvs left accAfterOnside+ else accAfterOnside+ accAfterCurrent = if distSqr k query <= radius * radius+ then (k, v) : accAfterOffside+ else accAfterOffside+ in accAfterCurrent++-- | Given a 'KdMap', a query point, and a number @k@, returns the @k@+-- point-value pairs with the nearest points to the query.+--+-- Neighbors are returned in order of increasing distance from query+-- point.+--+-- Average time complexity: /log(k) * log(n)/ for /k/ nearest+-- neighbors on a structure with /n/ data points.+--+-- Worst case time complexity: /n * log(k)/ for /k/ nearest+-- neighbors on a structure with /n/ data points.+kNearestNeighbors :: Real a => KdMap a p v -> Int -> p -> [(p, v)]+kNearestNeighbors (KdMap pointAsList distSqr t _) numNeighbors query =+ reverse $ map snd $ Q.toList $ go (cycle $ pointAsList query) Q.empty t+ where+ -- go :: [Double] -> Q.MaxPQueue Double (p, d) -> TreeNode p d -> KQueue p d+ go [] _ _ = error "kNearestNeighbors.go: no empty lists allowed!"+ go _ q Empty = q+ go (queryAxisValue : qvs) q (TreeNode left (k, v) nodeAxisVal right) =+ let insertBounded queue dist x+ | Q.size queue < numNeighbors = Q.insert dist x queue+ | otherwise = if dist < fst (Q.findMax queue)+ then Q.insert dist x $ Q.deleteMax queue+ else queue+ q' = insertBounded q (distSqr k query) (k, v)+ kNearest queue onsideSubtree offsideSubtree =+ let queue' = go qvs queue onsideSubtree+ checkOffsideTree =+ Q.size queue' < numNeighbors ||+ (queryAxisValue - nodeAxisVal)^(2 :: Int) < fst (Q.findMax queue')+ in if checkOffsideTree+ then go qvs queue' offsideSubtree+ else queue'+ in if queryAxisValue <= nodeAxisVal+ then kNearest q' left right+ else kNearest q' right left++-- | Finds all point-value pairs in a 'KdMap' with points within a+-- given range, where the range is specified as a set of lower and+-- upper bounds.+--+-- Points are not returned in any particular order.+--+-- Worst case time complexity: /O(n)/ for n data points and a range+-- that spans all the points.+--+-- TODO: Maybe use known bounds on entire tree structure to be able to+-- automatically count whole portions of tree as being within given+-- range.+pointsInRange :: Real a => KdMap a p v+ -> p -- ^ lower bounds of range+ -> p -- ^ upper bounds of range+ -> [(p, v)] -- ^ point-value pairs within+ -- given range+pointsInRange (KdMap pointAsList _ t _) lowers uppers =+ go (cycle (pointAsList lowers) `zip` cycle (pointAsList uppers)) t []+ where+ go [] _ _ = error "neighborsInRange.go: no empty lists allowed!"+ go _ Empty acc = acc+ go ((lower, upper) : nextBounds) (TreeNode left p nodeAxisVal right) acc =+ let accAfterLeft = if lower <= nodeAxisVal+ then go nextBounds left acc+ else acc+ accAfterRight = if upper > nodeAxisVal+ then go nextBounds right accAfterLeft+ else accAfterLeft+ valInRange l x u = l <= x && x <= u+ -- maybe "cache" lowers and uppers as lists sooner as hint+ -- to ghc. Also, maybe only need to check previously+ -- unchecked axes?+ currentInRange =+ L.and $ zipWith3 valInRange+ (pointAsList lowers) (pointAsList $ fst p) (pointAsList uppers)+ accAfterCurrent = if currentInRange+ then p : accAfterRight+ else accAfterRight+ in accAfterCurrent++-- | Returns the number of point-value pairs in the 'KdMap'.+--+-- Time complexity: /O(1)/+size :: KdMap a p v -> Int+size (KdMap _ _ _ n) = n++--------------------------------------------------------------------------------+-- Tests+--------------------------------------------------------------------------------++testElements :: [p] -> [(p, Int)]+testElements ps = zip ps [0 ..]++isTreeValid :: Real a => PointAsListFn a p -> Int -> TreeNode a p v -> Bool+isTreeValid _ _ Empty = True+isTreeValid pointAsList axis (TreeNode l (k, _) nodeAxisVal r) =+ let childrenAxisValues = map ((!! axis) . pointAsList . fst) . assocsInternal+ leftSubtreeLess = L.all (<= nodeAxisVal) $ childrenAxisValues l+ rightSubtreeGreater = L.all (> nodeAxisVal) $ childrenAxisValues r+ nextAxis = (axis + 1) `mod` length (pointAsList k)+ in leftSubtreeLess && rightSubtreeGreater &&+ isTreeValid pointAsList nextAxis l && isTreeValid pointAsList nextAxis r++checkValidTree :: Real a => PointAsListFn a p -> [p] -> Bool+checkValidTree pointAsList ps =+ let (KdMap _ _ r _) = buildKdMap pointAsList $ testElements ps+ in isTreeValid pointAsList 0 r++prop_validTree :: Property+prop_validTree = forAll (listOf1 arbitrary) $ checkValidTree pointAsList2d++checkElements :: (Ord p, Real a) => PointAsListFn a p -> [p] -> Bool+checkElements pointAsList ps =+ let kdt = buildKdMap pointAsList $ testElements ps+ in L.sort (assocs kdt) == L.sort (testElements ps)++prop_sameElements :: Property+prop_sameElements = forAll (listOf1 arbitrary) $ checkElements pointAsList2d++checkNumElements :: Real a => PointAsListFn a p -> [p] -> Bool+checkNumElements pointAsList ps =+ let (KdMap _ _ _ n) = buildKdMap pointAsList $ testElements ps+ in n == length ps++prop_validNumElements :: Property+prop_validNumElements = forAll (listOf1 arbitrary) $ checkNumElements pointAsList2d++nearestNeighborLinear :: Real a => PointAsListFn a p -> [(p, v)] -> p -> (p, v)+nearestNeighborLinear pointAsList xs query =+ L.minimumBy (compare `on` (defaultDistSqrFn pointAsList query . fst)) xs++checkNearestEqualToLinear :: (Eq p, Real a) => PointAsListFn a p -> ([p], p) -> Bool+checkNearestEqualToLinear pointAsList (ps, query) =+ let kdt = buildKdMap pointAsList $ testElements ps+ in nearestNeighbor kdt query == nearestNeighborLinear pointAsList (testElements ps) query++prop_nearestEqualToLinear :: Point2d -> Property+prop_nearestEqualToLinear query =+ forAll (listOf1 arbitrary) $ \xs ->+ checkNearestEqualToLinear pointAsList2d (xs, query)++pointsInRadiusLinear :: Real a => PointAsListFn a p -> [(p, v)] -> p -> a -> [(p, v)]+pointsInRadiusLinear pointAsList xs query radius =+ filter ((<= radius * radius) . defaultDistSqrFn pointAsList query . fst) xs++checkInRadiusEqualToLinear :: (Ord p, Real a) => PointAsListFn a p -> a -> ([p], p) -> Bool+checkInRadiusEqualToLinear pointAsList radius (ps, query) =+ let kdt = buildKdMap pointAsList $ testElements ps+ kdtNear = pointsInRadius kdt radius query+ linearNear = pointsInRadiusLinear pointAsList (testElements ps) query radius+ in L.sort kdtNear == L.sort linearNear++prop_inRadiusEqualToLinear :: Point2d -> Property+prop_inRadiusEqualToLinear query =+ forAll (listOf1 arbitrary) $ \xs ->+ forAll (choose (0.0, 1000.0)) $ \radius ->+ checkInRadiusEqualToLinear pointAsList2d radius (xs, query)++kNearestNeighborsLinear :: Real a => PointAsListFn a p -> [(p, v)] -> p -> Int -> [(p, v)]+kNearestNeighborsLinear pointAsList xs query k =+ take k $ L.sortBy (compare `on` (defaultDistSqrFn pointAsList query . fst)) xs++checkKNearestEqualToLinear :: (Ord p, Real a) => PointAsListFn a p -> Int -> ([p], p) -> Bool+checkKNearestEqualToLinear pointAsList k (xs, query) =+ let kdt = buildKdMap pointAsList $ testElements xs+ kdtKNear = kNearestNeighbors kdt k query+ linearKNear = kNearestNeighborsLinear pointAsList (testElements xs) query k+ in kdtKNear == linearKNear++prop_kNearestEqualToLinear :: Point2d -> Property+prop_kNearestEqualToLinear query =+ forAll (listOf1 arbitrary) $ \xs ->+ forAll (choose (1, length xs)) $ \k ->+ checkKNearestEqualToLinear pointAsList2d k (xs, query)++checkKNearestSorted :: (Eq p, Real a) => PointAsListFn a p -> ([p], p) -> Bool+checkKNearestSorted _ ([], _) = True+checkKNearestSorted pointAsList (ps, query) =+ let kdt = buildKdMap pointAsList $ testElements ps+ kNearestDists =+ map (defaultDistSqrFn pointAsList query . fst) $ kNearestNeighbors kdt (length ps) query+ in kNearestDists == L.sort kNearestDists++prop_kNearestSorted :: Point2d -> Property+prop_kNearestSorted query =+ forAll (listOf1 arbitrary) $ \xs ->+ checkKNearestSorted pointAsList2d (xs, query)++rangeLinear :: Real a => PointAsListFn a p -> [(p, v)] -> p -> p -> [(p, v)]+rangeLinear pointAsList xs lowers uppers =+ let valInRange a lower upper = lower <= a && a <= upper+ lowersAsList = pointAsList lowers+ uppersAsList = pointAsList uppers+ pointInRange (p, _) =+ L.and $ zipWith3 valInRange (pointAsList p) lowersAsList uppersAsList+ in filter pointInRange xs++prop_rangeEqualToLinear :: ([Point2d], Point2d, Point2d) -> Bool+prop_rangeEqualToLinear (xs, lowers, uppers)+ | null xs = True+ | L.and $ zipWith (<) (pointAsList2d lowers) (pointAsList2d uppers) =+ let linear = rangeLinear pointAsList2d (testElements xs) lowers uppers+ kdt = buildKdMap pointAsList2d $ testElements xs+ kdtPoints = pointsInRange kdt lowers uppers+ in L.sort linear == L.sort kdtPoints+ | otherwise = True++prop_equalAxisValueSameElems :: Property+prop_equalAxisValueSameElems =+ forAll (listOf1 arbitrary) $ \xs@(Point2d x y : _) ->+ checkElements pointAsList2d $ Point2d x (y + 1) : xs++prop_equalAxisValueEqualToLinear :: Point2d -> Property+prop_equalAxisValueEqualToLinear query =+ forAll (listOf1 arbitrary) $ \xs@(Point2d x y : _) ->+ checkNearestEqualToLinear pointAsList2d (Point2d x (y + 1) : xs, query)++-- Run all tests+return []+runTests :: IO Bool+runTests = $quickCheckAll
+ lib-src/Data/KdTree/Dynamic.hs view
@@ -0,0 +1,158 @@+module Data.KdTree.Dynamic+ ( -- * Usage++ -- $usage++ -- * Reference++ -- ** Types+ PointAsListFn+ , SquaredDistanceFn+ , KdTree+ -- ** Dynamic /k/-d tree construction+ , emptyKdTree+ , singleton+ , emptyKdTreeWithDistFn+ , singletonWithDistFn+ -- ** Insertion+ , insert+ -- ** Query+ , nearestNeighbor+ , pointsInRadius+ , kNearestNeighbors+ , pointsInRange+ , points+ , null+ , size+ -- ** Utilities+ , defaultDistSqrFn+ ) where++import Prelude hiding (null)++import Data.Foldable++import qualified Data.KdMap.Dynamic as DKDM+import Data.KdMap.Dynamic (PointAsListFn, SquaredDistanceFn, defaultDistSqrFn)++-- $usage+--+-- The 'KdTree' is a dynamic variant of+-- @Data.KdTree.Static.@'Data.KdTree.Static.KdTree' that allows for+-- insertion of new points into an existing 'KdTree'. This algorithm+-- was implemented using a+-- <http://repository.cmu.edu/cgi/viewcontent.cgi?article=3453&context=compsci static-to-dynamic transformation>.+--+-- Here's an example of interleaving 3D point insertions and point+-- queries using 'KdTree':+--+-- @+-- >>> let dkdt = singleton point3dAsList (Point3D 0.0 0.0 0.0)+--+-- >>> let dkdt' = insert dkdt (Point3D 1.0 1.0 1.0)+--+-- >>> nearestNeighbor dkdt' (Point3D 0.4 0.4 0.4)+-- Point3D {x = 0.0, y = 0.0, z = 0.0}+--+-- >>> let dkdt'' = insert dkdt' (Point3D 0.5 0.5 0.5)+--+-- >>> nearestNeighbor dkdt'' (Point3D 0.4 0.4 0.4)+-- Point3D {x = 0.5, y = 0.5, z = 0.5}+-- @+--+-- Check out @Data.KdMap.Dynamic.@'Data.KdMap.Dynamic.KdMap' if you want to associate a value+-- with each point in your tree structure.++-- | A dynamic /k/-d tree structure that stores points of type @p@+-- with axis values of type @a@.+newtype KdTree a p = KdTree (DKDM.KdMap a p ())++instance Foldable (KdTree a) where+ foldr f z (KdTree dkdMap) = DKDM.foldrKdMap (f . fst) z dkdMap++-- | Generates an empty 'KdTree' with a user-specified distance function.+emptyKdTreeWithDistFn :: Real a => PointAsListFn a p -> SquaredDistanceFn a p -> KdTree a p+emptyKdTreeWithDistFn p2l d2 = KdTree $ DKDM.emptyKdMapWithDistFn p2l d2++-- | Generates an empty 'KdTree' with the default distance function.+emptyKdTree :: Real a => PointAsListFn a p -> KdTree a p+emptyKdTree p2l = emptyKdTreeWithDistFn p2l $ defaultDistSqrFn p2l++-- | Returns whether the 'KdTree' is empty.+null :: KdTree a p -> Bool+null (KdTree dkdMap) = DKDM.null dkdMap++-- | Generates a 'KdTree' with a single point using a+-- user-specified distance function.+singletonWithDistFn :: Real a => PointAsListFn a p -> SquaredDistanceFn a p -> p -> KdTree a p+singletonWithDistFn p2l d2 p = KdTree $ DKDM.singletonWithDistFn p2l d2 (p, ())++-- | Generates a 'KdTree' with a single point using the default+-- distance function.+singleton :: Real a => PointAsListFn a p -> p -> KdTree a p+singleton p2l = singletonWithDistFn p2l $ defaultDistSqrFn p2l++-- | Adds a given point to a 'KdTree'.+--+-- Average time complexity per insert for /n/ inserts: /O(log^2(n))/.+insert :: Real a => KdTree a p -> p -> KdTree a p+insert (KdTree dkdMap) p = KdTree $ DKDM.insert dkdMap p ()++-- | Given a 'KdTree' and a query point, returns the nearest point+-- in the 'KdTree' to the query point.+--+-- Average time complexity: /O(log^2(n))/.+nearestNeighbor :: Real a => KdTree a p -> p -> p+nearestNeighbor (KdTree dkdMap) = fst . DKDM.nearestNeighbor dkdMap++-- | Given a 'KdTree', a query point, and a number @k@, returns the+-- @k@ nearest points in the 'KdTree' to the query point.+--+-- Neighbors are returned in order of increasing distance from query+-- point.+--+-- Average time complexity: /log(k) * log^2(n)/ for /k/ nearest+-- neighbors on a structure with /n/ data points.+--+-- Worst case time complexity: /n * log(k)/ for /k/ nearest neighbors+-- on a structure with /n/ data points.+kNearestNeighbors :: Real a => KdTree a p -> Int -> p -> [p]+kNearestNeighbors (KdTree dkdMap) k query =+ map fst $ DKDM.kNearestNeighbors dkdMap k query++-- | Given a 'KdTree', a query point, and a radius, returns all+-- points in the 'KdTree' that are within the given radius of the+-- query points.+--+-- Points are not returned in any particular order.+--+-- Worst case time complexity: /O(n)/ for /n/ data points.+pointsInRadius :: Real a => KdTree a p -> a -> p -> [p]+pointsInRadius (KdTree dkdMap) radius query =+ map fst $ DKDM.pointsInRadius dkdMap radius query++-- | Finds all points in a 'KdTree' with points within a given range,+-- where the range is specified as a set of lower and upper bounds.+--+-- Points are not returned in any particular order.+--+-- Worst case time complexity: /O(n)/ for n data points and a range+-- that spans all the points.+pointsInRange :: Real a => KdTree a p+ -> p -- ^ lower bounds of range+ -> p -- ^ upper bounds of range+ -> [p] -- ^ all points within given range+pointsInRange (KdTree dkdMap) lowers uppers =+ map fst $ DKDM.pointsInRange dkdMap lowers uppers++-- | Returns the number of elements in the 'KdTree'.+--+-- Time complexity: /O(1)/+size :: KdTree a p -> Int+size (KdTree dkdMap) = DKDM.size dkdMap++-- | Returns a list of all the points in the 'KdTree'.+--+-- Time complexity: /O(n)/+points :: KdTree a p -> [p]+points (KdTree dkdMap) = DKDM.points dkdMap
+ lib-src/Data/KdTree/Static.hs view
@@ -0,0 +1,283 @@+{-# LANGUAGE DeriveGeneric #-}++module Data.KdTree.Static+ ( -- * Introduction++ -- $intro++ -- * Usage++ -- $usage++ -- * Variants++ -- ** Dynamic /k/-d trees++ -- $dkdtrees++ -- ** /k/-d maps++ -- $kdmaps++ -- * Advanced++ -- ** Custom distance functions++ -- $customdistancefunctions++ -- ** Axis value types++ -- $axisvaluetypes++ -- * Reference++ -- ** Types+ PointAsListFn+ , SquaredDistanceFn+ , KdTree+ -- ** /k/-d tree construction+ , buildKdTree+ , buildKdTreeWithDistFn+ -- ** Query+ , nearestNeighbor+ , pointsInRadius+ , kNearestNeighbors+ , points+ , pointsInRange+ , size+ -- ** Utilities+ , defaultDistSqrFn+ ) where++import Control.DeepSeq+import Control.DeepSeq.Generics (genericRnf)+import GHC.Generics++import Data.Foldable++import qualified Data.KdMap.Static as KDM+import Data.KdMap.Static (PointAsListFn, SquaredDistanceFn, defaultDistSqrFn)++-- $intro+--+-- Let's say you have a large set of 3D points called /data points/,+-- and you'd like to be able to quickly perform /point queries/ on the+-- data points. One example of a point query is the /nearest neighbor/+-- query: given a set of data points @points@ and a query point @p@,+-- which point in @points@ is closest to @p@?+--+-- We can efficiently solve the nearest neighbor query (along with+-- many other types of point queries) if we appropriately organize the+-- data points. One such method of organization is called the /k/-d+-- tree algorithm, which is implemented in this module.++-- $usage+--+-- Let's say you have a list of 3D data points, and each point is of+-- type @Point3d@:+--+-- @+-- data Point3d = Point3d { x :: Double+-- , y :: Double+-- , z :: Double+-- } deriving Show+-- @+--+-- We call a point's individual values /axis values/ (i.e., @x@, @y@,+-- and @z@ in the case of @Point3d@).+--+-- In order to generate a /k/-d tree of @Point3d@'s, we need to define+-- a 'PointAsListFn' that expresses the point's axis values as a list:+--+-- @+-- point3dAsList :: Point3d -> [Double]+-- point3dAsList (Point3d x y z) = [x, y, z]+-- @+--+-- Now we can build a 'KdTree' structure from a list of data points+-- and perform a nearest neighbor query as follows:+--+-- @+-- >>> let dataPoints = [(Point3d 0.0 0.0 0.0), (Point3d 1.0 1.0 1.0)]+--+-- >>> let kdt = 'buildKdTree' point3dAsList dataPoints+--+-- >>> let queryPoint = Point3d 0.1 0.1 0.1+--+-- >>> 'nearestNeighbor' kdt queryPoint+-- Point3d {x = 0.0, y = 0.0, z = 0.0}+-- @++-- $dkdtrees+--+-- The 'KdTree' structure is meant for static sets of data points. If+-- you need to insert points into an existing /k/-d tree, check out+-- @Data.KdTree.Dynamic.@'Data.KdTree.Dynamic.KdTree'.++-- $kdmaps+--+-- If you need to associate additional data with each point in the+-- tree (i.e., points are /keys/ associated with /values/), check out+-- @Data.KdMap.Static.@'Data.KdMap.Static.KdMap' and+-- @Data.KdMap.Dynamic.@'Data.KdMap.Dynamic.KdMap' for static and dynamic+-- variants of this functionality. Please /do not/ try to fake this+-- functionality with a 'KdTree' by augmenting your point type with+-- the extra data; you're gonna have a bad time.++-- $customdistancefunctions+--+-- You may have noticed in the previous use case that we never+-- specified what "nearest" means for our points. By default,+-- 'buildKdTree' uses a Euclidean distance function that is sufficient+-- in most cases. However, point queries are typically faster on a+-- 'KdTree' built with a user-specified custom distance+-- function. Let's generate a 'KdTree' using a custom distance+-- function.+--+-- One idiosyncrasy about 'KdTree' is that custom distance functions+-- are actually specified as /squared distance/ functions+-- ('SquaredDistanceFn'). This means that your custom distance+-- function must return the /square/ of the actual distance between+-- two points. This is for efficiency: regular distance functions+-- often require expensive square root computations, whereas in our+-- case, the squared distance works fine and doesn't require computing+-- any square roots. Here's an example of a squared distance function+-- for @Point3d@:+--+-- @+-- point3dSquaredDistance :: Point3d -> Point3d -> Double+-- point3dSquaredDistance (Point3d x1 y1 z1) (Point3d x2 y2 z2) =+-- let dx = x1 - x2+-- dy = y1 - y2+-- dz = z1 - z2+-- in dx * dx + dy * dy + dz * dz+-- @+--+-- We can build a 'KdTree' using our custom distance function as follows:+--+-- @+-- >>> let kdt = 'buildKdTreeWithDistFn' point3dAsList point3dSquaredDistance points+-- @++-- $axisvaluetypes+--+-- In the above examples, we used a point type with axis values of+-- type 'Double'. We can in fact use axis values of any type that is+-- an instance of the 'Real' typeclass. This means you can use points+-- that are composed of 'Double's, 'Int's, 'Float's, and so on:+--+-- @+-- data Point2i = Point2i Int Int+--+-- point2iAsList :: Point2i -> [Int]+-- point2iAsList (Point2i x y) = [x, y]+--+-- kdt :: [Point2i] -> KdTree Int Point2i+-- kdt dataPoints = buildKdTree point2iAsList dataPoints+-- @++-- | A /k/-d tree structure that stores points of type @p@ with axis+-- values of type @a@.+newtype KdTree a p = KdTree (KDM.KdMap a p ()) deriving Generic+instance (NFData a, NFData p) => NFData (KdTree a p) where rnf = genericRnf++instance Foldable (KdTree a) where+ foldr f z (KdTree kdMap) = KDM.foldrKdMap (f . fst) z kdMap++-- | Builds a 'KdTree' from a list of data points using a default+-- squared distance function 'defaultDistSqrFn'.+--+-- Average complexity: /O(n * log(n))/ for /n/ data points.+--+-- Worst case time complexity: /O(n^2)/ for /n/ data points.+--+-- Worst case space complexity: /O(n)/ for /n/ data points.+--+-- Throws an error if given an empty list of data points.+buildKdTree :: Real a => PointAsListFn a p+ -> [p] -- ^ non-empty list of data points to be stored in the /k/-d tree+ -> KdTree a p+buildKdTree _ [] = error "KdTree must be built with a non-empty list."+buildKdTree pointAsList ps =+ KdTree $ KDM.buildKdMap pointAsList $ zip ps $ repeat ()++-- | Builds a 'KdTree' from a list of data points using a+-- user-specified squared distance function.+--+-- Average time complexity: /O(n * log(n))/ for /n/ data points.+--+-- Worst case time complexity: /O(n^2)/ for /n/ data points.+--+-- Worst case space complexity: /O(n)/ for /n/ data points.+--+-- Throws an error if given an empty list of data points.+buildKdTreeWithDistFn :: Real a => PointAsListFn a p+ -> SquaredDistanceFn a p+ -> [p]+ -> KdTree a p+buildKdTreeWithDistFn _ _ [] = error "KdTree must be built with a non-empty list."+buildKdTreeWithDistFn pointAsList distSqr ps =+ KdTree $ KDM.buildKdMapWithDistFn pointAsList distSqr $ zip ps $ repeat ()++-- | Given a 'KdTree' and a query point, returns the nearest point+-- in the 'KdTree' to the query point.+--+-- Average time complexity: /O(log(n))/ for /n/ data points.+--+-- Worst case time complexity: /O(n)/ for /n/ data points.+nearestNeighbor :: Real a => KdTree a p -> p -> p+nearestNeighbor (KdTree t) query = fst $ KDM.nearestNeighbor t query++-- | Given a 'KdTree', a query point, and a radius, returns all+-- points in the 'KdTree' that are within the given radius of the+-- query point.+--+-- Points are not returned in any particular order.+--+-- Worst case time complexity: /O(n)/ for /n/ data points and+-- a radius that subsumes all points in the structure.+pointsInRadius :: Real a => KdTree a p+ -> a -- ^ radius+ -> p -- ^ query point+ -> [p] -- ^ list of points in tree with+ -- given radius of query point+pointsInRadius (KdTree t) radius query = map fst $ KDM.pointsInRadius t radius query++-- | Given a 'KdTree', a query point, and a number @k@, returns the+-- @k@ nearest points in the 'KdTree' to the query point.+--+-- Neighbors are returned in order of increasing distance from query+-- point.+--+-- Average time complexity: /log(k) * log(n)/ for /k/ nearest+-- neighbors on a structure with /n/ data points.+--+-- Worst case time complexity: /n * log(k)/ for /k/ nearest+-- neighbors on a structure with /n/ data points.+kNearestNeighbors :: Real a => KdTree a p -> Int -> p -> [p]+kNearestNeighbors (KdTree t) k query = map fst $ KDM.kNearestNeighbors t k query++-- | Finds all points in a 'KdTree' with points within a given range,+-- where the range is specified as a set of lower and upper bounds.+--+-- Points are not returned in any particular order.+--+-- Worst case time complexity: /O(n)/ for n data points and a range+-- that spans all the points.+pointsInRange :: Real a => KdTree a p+ -> p -- ^ lower bounds of range+ -> p -- ^ upper bounds of range+ -> [p] -- ^ all points within given range+pointsInRange (KdTree t) lower upper = map fst $ KDM.pointsInRange t lower upper++-- | Returns a list of all the points in the 'KdTree'.+--+-- Time complexity: /O(n)/ for /n/ data points.+points :: KdTree a p -> [p]+points (KdTree t) = KDM.points t++-- | Returns the number of elements in the 'KdTree'.+--+-- Time complexity: /O(1)/+size :: KdTree a p -> Int+size (KdTree t) = KDM.size t
+ lib-src/Data/Point2d.hs view
@@ -0,0 +1,27 @@+{-# OPTIONS_HADDOCK hide #-}++{-# LANGUAGE DeriveGeneric #-}++module Data.Point2d where++import Control.DeepSeq+import Control.DeepSeq.Generics (genericRnf)+import GHC.Generics+import Test.QuickCheck++data Point2d = Point2d Double Double deriving (Show, Eq, Ord, Generic)+instance NFData Point2d where rnf = genericRnf++pointAsList2d :: Point2d -> [Double]+pointAsList2d (Point2d x y) = [x, y]++distSqr2d :: Point2d -> Point2d -> Double+distSqr2d (Point2d x1 y1) (Point2d x2 y2) = let dx = x2 - x1+ dy = y2 - y1+ in dx*dx + dy*dy++instance Arbitrary Point2d where+ arbitrary = do+ x <- arbitrary+ y <- arbitrary+ return (Point2d x y)