diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -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. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/app-src/Benchmarks/KDTBenchmark.hs b/app-src/Benchmarks/KDTBenchmark.hs new file mode 100644 --- /dev/null +++ b/app-src/Benchmarks/KDTBenchmark.hs @@ -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 + ] diff --git a/app-src/Tests/KdTreeTest.hs b/app-src/Tests/KdTreeTest.hs new file mode 100644 --- /dev/null +++ b/app-src/Tests/KdTreeTest.hs @@ -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 diff --git a/kdt.cabal b/kdt.cabal new file mode 100644 --- /dev/null +++ b/kdt.cabal @@ -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 diff --git a/lib-src/Data/KdMap/Dynamic.hs b/lib-src/Data/KdMap/Dynamic.hs new file mode 100644 --- /dev/null +++ b/lib-src/Data/KdMap/Dynamic.hs @@ -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 diff --git a/lib-src/Data/KdMap/Static.hs b/lib-src/Data/KdMap/Static.hs new file mode 100644 --- /dev/null +++ b/lib-src/Data/KdMap/Static.hs @@ -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 diff --git a/lib-src/Data/KdTree/Dynamic.hs b/lib-src/Data/KdTree/Dynamic.hs new file mode 100644 --- /dev/null +++ b/lib-src/Data/KdTree/Dynamic.hs @@ -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 +-- . +-- +-- 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 diff --git a/lib-src/Data/KdTree/Static.hs b/lib-src/Data/KdTree/Static.hs new file mode 100644 --- /dev/null +++ b/lib-src/Data/KdTree/Static.hs @@ -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 diff --git a/lib-src/Data/Point2d.hs b/lib-src/Data/Point2d.hs new file mode 100644 --- /dev/null +++ b/lib-src/Data/Point2d.hs @@ -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)