KdTree 0.1 → 0.2

raw patch · 3 files changed

+90/−55 lines, 3 filesPVP ok

version bump matches the API change (PVP)

API changes (from Hackage documentation)

Files

• Data/Trees/KdTree.hs +64/−47
• KdTree.cabal +1/−1
• KdTreeTest.hs +25/−7

@@ -1,7 +1,6 @@-module Data.Trees.KdTree where+-- http://en.wikipedia.org/wiki/K-d_tree --- Haskell implementation of http://en.wikipedia.org/wiki/K-d_tree--- by Issac Trotts+module Data.Trees.KdTree where  import Data.Maybe @@ -19,7 +18,7 @@       -- |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+        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@@ -37,9 +36,9 @@   data KdTree point = KdNode { kdLeft :: KdTree point,-			     kdPoint :: point,+                             kdPoint :: point,                              kdRight :: KdTree point,-			     kdAxis :: Int }+                             kdAxis :: Int }                   | KdEmpty      deriving (Eq, Ord, Show) @@ -50,8 +49,8 @@ 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+        where init3 = f x init2+              init2 = F.foldr f init r  fromList :: Point p => [p] -> KdTree p fromList points = fromListWithDepth points 0@@ -60,66 +59,84 @@ fromListWithDepth :: Point p => [p] -> Int -> KdTree p fromListWithDepth [] _ = KdEmpty fromListWithDepth points depth = node-    where   axis = axisFromDepth (head points) depth--	    -- 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-	-	    -- Create node and construct subtrees-	    node = KdNode { kdLeft = fromListWithDepth (take medianIndex sortedPoints) (depth+1),-			    kdPoint = sortedPoints !! medianIndex,-			    kdRight = fromListWithDepth (drop (medianIndex+1) sortedPoints) (depth+1),-			    kdAxis = axis }+    where axis = depth `mod` dimension (head points)  -axisFromDepth :: Point p => p -> Int -> Int-axisFromDepth p depth = depth `mod` k-    where k = dimension p+          -- 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+        +          -- Create node and construct subtrees+          node = KdNode { kdLeft = fromListWithDepth (take medianIndex sortedPoints) (depth+1),+                          kdPoint = sortedPoints !! medianIndex,+                          kdRight = fromListWithDepth (drop (medianIndex+1) sortedPoints) (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 p r axis) probe =-    if xProbe <= xp then doStuff l r else doStuff r l+    if xProbe <= xp then findNearest l r else findNearest r l     where xProbe = coord axis probe-	  xp = coord axis p-          doStuff tree1 tree2 =-		let candidates1 = case nearestNeighbor tree1 probe of-				    Nothing -> [p]-				    Just best1 -> [best1, p]-		    sphereIntersectsPlane = (xProbe - xp)^2 <= dist2 probe p-		    candidates2 = if sphereIntersectsPlane-				    then candidates1 ++ maybeToList (nearestNeighbor tree2 probe)-				    else candidates1 in-		Just . L.minimumBy (compareDistance probe) $ candidates2+          xp = coord axis p+          findNearest tree1 tree2 =+                let candidates1 = case nearestNeighbor tree1 probe of+                                    Nothing -> [p]+                                    Just best1 -> [best1, p]+                    sphereIntersectsPlane = (xProbe - xp)^2 <= dist2 probe p+                    candidates2 = if sphereIntersectsPlane+                                    then candidates1 ++ maybeToList (nearestNeighbor tree2 probe)+                                    else candidates1 in+                Just . L.minimumBy (compareDistance probe) $ candidates2 --- |invariant tells whether the KD tree property holds for a given tree and--- all its subtrees.+-- |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.-invariant :: Point p => KdTree p -> Bool-invariant KdEmpty = True-invariant (KdNode l p r axis) = leftIsGood && rightIsGood+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)+          leftIsGood = all ((<= x) . coord axis) (toList l)+          rightIsGood = all ((>= x) . coord axis) (toList r) -invariant' :: Point p => KdTree p -> Bool-invariant' = all invariant . subtrees+-- |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)+        x <- arbitrary+        y <- arbitrary+        z <- arbitrary+        return (Point3d x y z) 

@@ -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.1+Version:             0.2 Synopsis:            KdTree, for efficient search in K-dimensional point clouds. Description:              This is a simple library for k-d trees in Haskell. It enables efficient

@@ -10,26 +10,44 @@  import qualified Data.Trees.KdTree as Kd -prop_invariant :: [Kd.Point3d] -> Bool-prop_invariant points = Kd.invariant' . Kd.fromList $ points+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_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_pointsAreClosestToThemselves :: [Kd.Point3d] -> Bool prop_pointsAreClosestToThemselves points =     map Just points == map (Kd.nearestNeighbor tree) points     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_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  main = $quickCheckAll