KdTree 0.2.2.0 → 0.2.2.1
raw patch · 6 files changed
+242/−242 lines, 6 filesdep +KdTreesetup-changedPVP ok
version bump matches the API change (PVP)
Dependencies added: KdTree
API changes (from Hackage documentation)
Files
- Data/Trees/KdTree.hs +0/−166
- KdTree.cabal +11/−10
- KdTreeTest.hs +0/−65
- Setup.hs +1/−1
- src/Data/Trees/KdTree.hs +166/−0
- test/KdTreeTest.hs +64/−0
− Data/Trees/KdTree.hs
@@ -1,166 +0,0 @@--- http://en.wikipedia.org/wiki/K-d_tree--module Data.Trees.KdTree where--import Data.Maybe--import qualified Data.Foldable as F-import qualified Data.List as L-import Test.QuickCheck--class Point p where- -- |dimension returns the number of coordinates of a point.- dimension :: p -> Int-- -- |coord gets the k'th coordinate, starting from 0.- coord :: Int -> p -> Double-- -- |dist2 returns the squared distance between two points.- dist2 :: p -> p -> Double- dist2 a b = sum . map diff2 $ [0..dimension a - 1]- where diff2 i = (coord i a - coord i b)^2---- |compareDistance p a b compares the distances of a and b to p.-compareDistance :: (Point p) => p -> p -> p -> Ordering-compareDistance p a b = dist2 p a `compare` dist2 p b--data Point3d = Point3d { p3x :: Double, p3y :: Double, p3z :: Double }- deriving (Eq, Ord, Show)--instance Point Point3d where- dimension _ = 3-- coord 0 p = p3x p- coord 1 p = p3y p- coord 2 p = p3z p---data KdTree point = KdNode { kdLeft :: KdTree point,- kdPoint :: point,- kdRight :: KdTree point,- kdAxis :: Int }- | KdEmpty- deriving (Eq, Ord, Show)--instance Functor KdTree where- fmap _ KdEmpty = KdEmpty- fmap f (KdNode l x r axis) = KdNode (fmap f l) (f x) (fmap f r) axis--instance F.Foldable KdTree where- foldr f init KdEmpty = init- foldr f init (KdNode l x r _) = F.foldr f init3 l- where init3 = f x init2- init2 = F.foldr f init r--fromList :: Point p => [p] -> KdTree p-fromList points = fromListWithDepth points 0---- |fromListWithDepth selects an axis based on depth so that the axis cycles--- through all valid values.-fromListWithDepth :: Point p => [p] -> Int -> KdTree p-fromListWithDepth [] _ = KdEmpty-fromListWithDepth points depth = node- where axis = depth `mod` dimension (head points) -- -- Sort point list and choose median as pivot element- sortedPoints =- L.sortBy (\a b -> coord axis a `compare` coord axis b) points- medianIndex = length sortedPoints `div` 2- medianCoordinate = coord axis (sortedPoints !! medianIndex)- - leftPoints = filter (\p -> coord axis p < medianCoordinate) sortedPoints- trueMedianIndex = length leftPoints- rightPoints = drop (trueMedianIndex+1) sortedPoints- - -- Create node and construct subtrees- node = KdNode { kdLeft = fromListWithDepth leftPoints (depth+1),- kdPoint = sortedPoints !! trueMedianIndex,- kdRight = fromListWithDepth rightPoints (depth+1),- kdAxis = axis }--toList :: KdTree p -> [p]-toList t = F.foldr (:) [] t---- |subtrees t returns a list containing t and all its subtrees, including the--- empty leaf nodes.-subtrees :: KdTree p -> [KdTree p]-subtrees KdEmpty = [KdEmpty]-subtrees t@(KdNode l x r axis) = subtrees l ++ [t] ++ subtrees r---- |nearestNeighbor tree p returns the nearest neighbor of p in tree.-nearestNeighbor :: Point p => KdTree p -> p -> Maybe p-nearestNeighbor KdEmpty probe = Nothing-nearestNeighbor (KdNode KdEmpty p KdEmpty _) probe = Just p-nearestNeighbor (KdNode l pivot r axis) probe =- if xProbe < xPivot then findNearest l r else findNearest r l- where xProbe = coord axis probe- xPivot = coord axis pivot- findNearest tree1 tree2 =- let candidate1 = case nearestNeighbor tree1 probe of- Nothing -> pivot- Just best -> L.minimumBy (compareDistance probe) [best, pivot]- sphereIntersectsPlane = (xProbe - xPivot)^2 <= dist2 probe candidate1- candidates2 = if sphereIntersectsPlane- then [candidate1] ++ maybeToList (nearestNeighbor tree2 probe)- else [candidate1] in- Just . L.minimumBy (compareDistance probe) $ candidates2---- |nearNeighbors tree p returns all neighbors within distance r from p in tree.-nearNeighbors :: Point p => KdTree p -> Double -> p -> [p]-nearNeighbors KdEmpty radius probe = []-nearNeighbors (KdNode KdEmpty p KdEmpty _) radius probe = if dist2 p probe <= radius^2 then [p] else []-nearNeighbors (KdNode l p r axis) radius probe =- if xProbe <= xp- then let nearest = maybePivot ++ nearNeighbors l radius probe- in if xProbe + abs radius > xp- then nearNeighbors r radius probe ++ nearest- else nearest- else let nearest = maybePivot ++ nearNeighbors r radius probe- in if xProbe - abs radius < xp- then nearNeighbors l radius probe ++ nearest- else nearest- where xProbe = coord axis probe- xp = coord axis p- maybePivot = if dist2 probe p <= radius^2 then [p] else []---- |isValid tells whether the K-D tree property holds for a given tree.--- Specifically, it tests that all points in the left subtree lie to the left--- of the plane, p is on the plane, and all points in the right subtree lie to--- the right.-isValid :: Point p => KdTree p -> Bool-isValid KdEmpty = True-isValid (KdNode l p r axis) = leftIsGood && rightIsGood- where x = coord axis p- leftIsGood = all ((<= x) . coord axis) (toList l)- rightIsGood = all ((>= x) . coord axis) (toList r)---- |allSubtreesAreValid tells whether the K-D tree property holds for the given--- tree and all subtrees.-allSubtreesAreValid :: Point p => KdTree p -> Bool-allSubtreesAreValid = all isValid . subtrees---- |kNearestNeighbors tree k p returns the k closest points to p within tree.-kNearestNeighbors :: (Eq p, Point p) => KdTree p -> Int -> p -> [p]-kNearestNeighbors KdEmpty _ _ = []-kNearestNeighbors _ k _ | k <= 0 = []-kNearestNeighbors tree k probe = nearest : kNearestNeighbors tree' (k-1) probe- where nearest = fromJust $ nearestNeighbor tree probe- tree' = tree `remove` nearest---- |remove t p removes the point p from t.-remove :: (Eq p, Point p) => KdTree p -> p -> KdTree p-remove KdEmpty _ = KdEmpty-remove (KdNode l p r axis) pKill =- if p == pKill- then fromListWithDepth (toList l ++ toList r) axis- else if coord axis pKill <= coord axis p- then KdNode (remove l pKill) p r axis- else KdNode l p (remove r pKill) axis--instance Arbitrary Point3d where- arbitrary = do- x <- arbitrary- y <- arbitrary- z <- arbitrary- return (Point3d x y z)-
KdTree.cabal view
@@ -3,7 +3,7 @@ -- The package version. See the Haskell package versioning policy -- (http://www.haskell.org/haskellwiki/Package_versioning_policy) for -- standards guiding when and how versions should be incremented.-Version: 0.2.2.0+Version: 0.2.2.1 Synopsis: KdTree, for efficient search in K-dimensional point clouds. Description: This is a simple library for k-d trees in Haskell. It enables@@ -26,16 +26,17 @@ location: git@github.com:binarysunrise-io/kdtree.git Library- Default-language: Haskell2010- Exposed-modules: Data.Trees.KdTree- Build-depends: base < 5+ default-language: Haskell2010+ hs-source-dirs: src+ exposed-modules: Data.Trees.KdTree+ build-depends: base < 5 , QuickCheck Test-suite KdTreeTest- Type: exitcode-stdio-1.0- Main-is: KdTreeTest.hs- Default-language: Haskell2010- other-modules: Data.Trees.KdTree-- Build-depends: base >= 4.8 && < 5+ type: exitcode-stdio-1.0+ main-is: KdTreeTest.hs+ default-language: Haskell2010+ hs-source-dirs: test+ build-depends: base >= 4.8 && < 5+ , KdTree , QuickCheck
− KdTreeTest.hs
@@ -1,65 +0,0 @@-{-# LANGUAGE TemplateHaskell #-}--module Main where--import Data.Maybe-import qualified Data.List as L--import Test.QuickCheck-import Test.QuickCheck.All--import qualified Data.Trees.KdTree as Kd--prop_constructionProducesValidTrees :: [Kd.Point3d] -> Bool-prop_constructionProducesValidTrees points =- Kd.allSubtreesAreValid . Kd.fromList $ points--prop_samePoints :: [Kd.Point3d] -> Bool-prop_samePoints points =- L.sort points == (L.sort . Kd.toList . Kd.fromList $ points)--prop_nearestNeighbor :: [Kd.Point3d] -> Kd.Point3d -> Bool-prop_nearestNeighbor points probe =- Kd.nearestNeighbor tree probe == bruteNearestNeighbor points probe - where tree = Kd.fromList points- bruteNearestNeighbor :: [Kd.Point3d] -> Kd.Point3d -> Maybe Kd.Point3d- bruteNearestNeighbor [] _ = Nothing- bruteNearestNeighbor points probe =- Just . head . L.sortBy (Kd.compareDistance probe) $ points--prop_nearNeighbors :: [Kd.Point3d] -> Kd.Point3d -> Double -> Bool-prop_nearNeighbors points probe radius =- (L.sort (Kd.nearNeighbors tree radius probe) ==- L.sort (bruteNearNeighbors points radius probe))- where tree = Kd.fromList points- bruteNearNeighbors :: [Kd.Point3d] -> Double -> Kd.Point3d -> [Kd.Point3d]- bruteNearNeighbors [] radius _ = []- bruteNearNeighbors points radius probe =- filter (withinDistance probe radius) points- withinDistance probe radius point = Kd.dist2 probe point <= radius^2--prop_pointsAreClosestToThemselves :: [Kd.Point3d] -> Bool-prop_pointsAreClosestToThemselves points =- map Just points == map (Kd.nearestNeighbor tree) points- where tree = Kd.fromList points--prop_kNearestNeighborsMatchesBrute :: [Kd.Point3d] -> Int -> Kd.Point3d -> Bool-prop_kNearestNeighborsMatchesBrute points k p =- L.sort (Kd.kNearestNeighbors tree k p) == L.sort (bruteKnearestNeighbors points k p)- where tree = Kd.fromList points- bruteKnearestNeighbors points k p =- take k . L.sortBy (Kd.compareDistance p) $ points--prop_removeReallyRemovesPoints :: [Kd.Point3d] -> Property-prop_removeReallyRemovesPoints points = points /= [] ==>- L.sort (Kd.toList (tree `Kd.remove` (head points))) == L.sort (tail points)- where tree = Kd.fromList points--prop_removePreservesInvariant :: [Kd.Point3d] -> Kd.Point3d -> Bool-prop_removePreservesInvariant points pKill =- Kd.allSubtreesAreValid $ tree `Kd.remove` pKill- where tree = Kd.fromList points--return []-main = $quickCheckAll-
Setup.hs view
@@ -1,2 +1,2 @@-import Distribution.Simple+import Distribution.Simple main = defaultMain
+ src/Data/Trees/KdTree.hs view
@@ -0,0 +1,166 @@+-- http://en.wikipedia.org/wiki/K-d_tree++module Data.Trees.KdTree where++import Data.Maybe++import qualified Data.Foldable as F+import qualified Data.List as L+import Test.QuickCheck++class Point p where+ -- |dimension returns the number of coordinates of a point.+ dimension :: p -> Int++ -- |coord gets the k'th coordinate, starting from 0.+ coord :: Int -> p -> Double++ -- |dist2 returns the squared distance between two points.+ dist2 :: p -> p -> Double+ dist2 a b = sum . map diff2 $ [0..dimension a - 1]+ where diff2 i = (coord i a - coord i b)^2++-- |compareDistance p a b compares the distances of a and b to p.+compareDistance :: (Point p) => p -> p -> p -> Ordering+compareDistance p a b = dist2 p a `compare` dist2 p b++data Point3d = Point3d { p3x :: Double, p3y :: Double, p3z :: Double }+ deriving (Eq, Ord, Show)++instance Point Point3d where+ dimension _ = 3++ coord 0 p = p3x p+ coord 1 p = p3y p+ coord 2 p = p3z p+++data KdTree point =+ KdNode {+ kdLeft :: KdTree point,+ kdPoint :: point,+ kdRight :: KdTree point,+ kdAxis :: Int+ }+ | KdEmpty deriving (Eq, Ord, Show)++instance Functor KdTree where+ fmap _ KdEmpty = KdEmpty+ fmap f (KdNode l x r axis) = KdNode (fmap f l) (f x) (fmap f r) axis++instance F.Foldable KdTree where+ foldr f init KdEmpty = init+ foldr f init (KdNode l x r _) = F.foldr f init3 l+ where init3 = f x init2+ init2 = F.foldr f init r++fromList :: Point p => [p] -> KdTree p+fromList points = fromListWithDepth points 0++-- |fromListWithDepth selects an axis based on depth so that the axis cycles+-- through all valid values.+fromListWithDepth :: Point p => [p] -> Int -> KdTree p+fromListWithDepth [] _ = KdEmpty+fromListWithDepth points depth = node+ where axis = depth `mod` dimension (head points)++ -- Sort point list and choose median as pivot element+ sortedPoints =+ L.sortBy (\a b -> coord axis a `compare` coord axis b) points+ medianIndex = length sortedPoints `div` 2+ medianCoordinate = coord axis (sortedPoints !! medianIndex)++ leftPoints = filter (\p -> coord axis p < medianCoordinate) sortedPoints+ trueMedianIndex = length leftPoints+ rightPoints = drop (trueMedianIndex+1) sortedPoints++ -- Create node and construct subtrees+ node = KdNode { kdLeft = fromListWithDepth leftPoints (depth+1),+ kdPoint = sortedPoints !! trueMedianIndex,+ kdRight = fromListWithDepth rightPoints (depth+1),+ kdAxis = axis }++toList :: KdTree p -> [p]+toList = F.foldr (:) []++-- |subtrees t returns a list containing t and all its subtrees, including the+-- empty leaf nodes.+subtrees :: KdTree p -> [KdTree p]+subtrees KdEmpty = [KdEmpty]+subtrees t@(KdNode l x r axis) = subtrees l ++ [t] ++ subtrees r++-- |nearestNeighbor tree p returns the nearest neighbor of p in tree.+nearestNeighbor :: Point p => KdTree p -> p -> Maybe p+nearestNeighbor KdEmpty probe = Nothing+nearestNeighbor (KdNode KdEmpty p KdEmpty _) probe = Just p+nearestNeighbor (KdNode l pivot r axis) probe =+ if xProbe < xPivot then findNearest l r else findNearest r l+ where xProbe = coord axis probe+ xPivot = coord axis pivot+ findNearest tree1 tree2 =+ let candidate1 = case nearestNeighbor tree1 probe of+ Nothing -> pivot+ Just best -> L.minimumBy (compareDistance probe) [best, pivot]+ sphereIntersectsPlane = (xProbe - xPivot)^2 <= dist2 probe candidate1+ candidates2 = if sphereIntersectsPlane+ then candidate1 : maybeToList (nearestNeighbor tree2 probe)+ else [candidate1] in+ Just . L.minimumBy (compareDistance probe) $ candidates2++-- |nearNeighbors tree p returns all neighbors within distance r from p in tree.+nearNeighbors :: Point p => KdTree p -> Double -> p -> [p]+nearNeighbors KdEmpty radius probe = []+nearNeighbors (KdNode KdEmpty p KdEmpty _) radius probe = [p | dist2 p probe <= radius^2]+nearNeighbors (KdNode l p r axis) radius probe =+ if xProbe <= xp+ then let nearest = maybePivot ++ nearNeighbors l radius probe+ in if xProbe + abs radius > xp+ then nearNeighbors r radius probe ++ nearest+ else nearest+ else let nearest = maybePivot ++ nearNeighbors r radius probe+ in if xProbe - abs radius < xp+ then nearNeighbors l radius probe ++ nearest+ else nearest+ where xProbe = coord axis probe+ xp = coord axis p+ maybePivot = [p | dist2 probe p <= radius^2]++-- |isValid tells whether the K-D tree property holds for a given tree.+-- Specifically, it tests that all points in the left subtree lie to the left+-- of the plane, p is on the plane, and all points in the right subtree lie to+-- the right.+isValid :: Point p => KdTree p -> Bool+isValid KdEmpty = True+isValid (KdNode l p r axis) = leftIsGood && rightIsGood+ where x = coord axis p+ leftIsGood = all ((<= x) . coord axis) (toList l)+ rightIsGood = all ((>= x) . coord axis) (toList r)++-- |allSubtreesAreValid tells whether the K-D tree property holds for the given+-- tree and all subtrees.+allSubtreesAreValid :: Point p => KdTree p -> Bool+allSubtreesAreValid = all isValid . subtrees++-- |kNearestNeighbors tree k p returns the k closest points to p within tree.+kNearestNeighbors :: (Eq p, Point p) => KdTree p -> Int -> p -> [p]+kNearestNeighbors KdEmpty _ _ = []+kNearestNeighbors _ k _ | k <= 0 = []+kNearestNeighbors tree k probe = nearest : kNearestNeighbors tree' (k-1) probe+ where nearest = fromJust $ nearestNeighbor tree probe+ tree' = tree `remove` nearest++-- |remove t p removes the point p from t.+remove :: (Eq p, Point p) => KdTree p -> p -> KdTree p+remove KdEmpty _ = KdEmpty+remove (KdNode l p r axis) pKill+ | p == pKill = fromListWithDepth (toList l ++ toList r) axis+ | coord axis pKill <= coord axis p = KdNode (remove l pKill) p r axis+ | otherwise = KdNode l p (remove r pKill) axis++instance Arbitrary Point3d where+ arbitrary = do+ x <- arbitrary+ y <- arbitrary+ z <- arbitrary+ return (Point3d x y z)+
+ test/KdTreeTest.hs view
@@ -0,0 +1,64 @@+{-# LANGUAGE TemplateHaskell #-}++module Main where++import qualified Data.List as L+import Data.Maybe++import Test.QuickCheck+import Test.QuickCheck.All++import qualified Data.Trees.KdTree as Kd++prop_constructionProducesValidTrees :: [Kd.Point3d] -> Bool+prop_constructionProducesValidTrees = Kd.allSubtreesAreValid . Kd.fromList++prop_samePoints :: [Kd.Point3d] -> Bool+prop_samePoints points =+ L.sort points == (L.sort . Kd.toList . Kd.fromList $ points)++prop_nearestNeighbor :: [Kd.Point3d] -> Kd.Point3d -> Bool+prop_nearestNeighbor points probe =+ Kd.nearestNeighbor tree probe == bruteNearestNeighbor points probe+ where tree = Kd.fromList points+ bruteNearestNeighbor :: [Kd.Point3d] -> Kd.Point3d -> Maybe Kd.Point3d+ bruteNearestNeighbor [] _ = Nothing+ bruteNearestNeighbor points probe =+ Just . L.minimumBy (Kd.compareDistance probe) $ points++prop_nearNeighbors :: [Kd.Point3d] -> Kd.Point3d -> Double -> Bool+prop_nearNeighbors points probe radius =+ L.sort (Kd.nearNeighbors tree radius probe) ==+ L.sort (bruteNearNeighbors points radius probe)+ where tree = Kd.fromList points+ bruteNearNeighbors :: [Kd.Point3d] -> Double -> Kd.Point3d -> [Kd.Point3d]+ bruteNearNeighbors [] radius _ = []+ bruteNearNeighbors points radius probe =+ filter (withinDistance probe radius) points+ withinDistance probe radius point = Kd.dist2 probe point <= radius^2++prop_pointsAreClosestToThemselves :: [Kd.Point3d] -> Bool+prop_pointsAreClosestToThemselves points =+ map Just points == map (Kd.nearestNeighbor tree) points+ where tree = Kd.fromList points++prop_kNearestNeighborsMatchesBrute :: [Kd.Point3d] -> Int -> Kd.Point3d -> Bool+prop_kNearestNeighborsMatchesBrute points k p =+ L.sort (Kd.kNearestNeighbors tree k p) == L.sort (bruteKnearestNeighbors points k p)+ where tree = Kd.fromList points+ bruteKnearestNeighbors points k p =+ take k . L.sortBy (Kd.compareDistance p) $ points++prop_removeReallyRemovesPoints :: [Kd.Point3d] -> Property+prop_removeReallyRemovesPoints points = points /= [] ==>+ L.sort (Kd.toList (tree `Kd.remove` head points)) == L.sort (tail points)+ where tree = Kd.fromList points++prop_removePreservesInvariant :: [Kd.Point3d] -> Kd.Point3d -> Bool+prop_removePreservesInvariant points pKill =+ Kd.allSubtreesAreValid $ tree `Kd.remove` pKill+ where tree = Kd.fromList points++return []+main = $quickCheckAll+