KdTree (empty) → 0.1

raw patch · 6 files changed

+248/−0 lines, 6 filesdep +QuickCheckdep +basesetup-changed

Dependencies added: QuickCheck, base

Files

Data/Trees/KdTree.hs +125/−0
KdTree.cabal +35/−0
KdTreeTest.hs +35/−0
LICENSE +30/−0
README +21/−0
Setup.hs +2/−0

+ Data/Trees/KdTree.hs view

@@ -0,0 +1,125 @@+module Data.Trees.KdTree where++-- Haskell implementation of http://en.wikipedia.org/wiki/K-d_tree+-- by Issac Trotts++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++-- Select axis based on depth so that axis cycles through all valid values+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 }++axisFromDepth :: Point p => p -> Int -> Int+axisFromDepth p depth = depth `mod` k+    where k = dimension p++toList :: KdTree p -> [p]+toList t = F.foldr (:) [] t++subtrees :: KdTree p -> [KdTree p]+subtrees KdEmpty = [KdEmpty]+subtrees t@(KdNode l x r axis) = subtrees l ++ [t] ++ subtrees r++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+    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++-- |invariant tells whether the KD tree property holds for a given tree and+-- all its subtrees.+-- 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+    where x = coord axis p+	  leftIsGood = all ((<= x) . coord axis) (toList l)+	  rightIsGood = all ((>= x) . coord axis) (toList r)++invariant' :: Point p => KdTree p -> Bool+invariant' = all invariant . subtrees++instance Arbitrary Point3d where+    arbitrary = do+	x <- arbitrary+	y <- arbitrary+	z <- arbitrary+	return (Point3d x y z)+

+ KdTree.cabal view

@@ -0,0 +1,35 @@+Name:                KdTree++-- 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+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+    searching through collections of points in O(log N) time on average,+    using the nearestNeighbor function.++Homepage:            https://github.com/ijt/kdtree+License:             BSD3+License-file:        LICENSE+Author:              Issac Trotts+Maintainer:          issac.trotts@gmail.com+Copyright:           Copyright 2011, Issac Trotts+Category:            Graphics+Build-type:          Simple+Cabal-version:       >=1.6+Extra-source-files:  README++source-repository head+  type:     git+  location: git@github.com:ijt/kdtree.git++Library+  Exposed-modules: Data.Trees.KdTree+  Build-depends: base < 5+  +Executable KdTreeTest+  Main-is: KdTreeTest.hs+  Build-depends: QuickCheck+

+ KdTreeTest.hs view

@@ -0,0 +1,35 @@+{-# 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_invariant :: [Kd.Point3d] -> Bool+prop_invariant points = Kd.invariant' . 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++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++main = $quickCheckAll+

+ LICENSE view

@@ -0,0 +1,30 @@+Copyright (c)2011, Issac Trotts++All rights reserved.++Redistribution and use in source and binary forms, with or without+modification, are permitted provided that the following conditions are met:++    * Redistributions of source code must retain the above copyright+      notice, this list of conditions and the following disclaimer.++    * Redistributions in binary form must reproduce the above+      copyright notice, this list of conditions and the following+      disclaimer in the documentation and/or other materials provided+      with the distribution.++    * Neither the name of Issac Trotts nor the names of other+      contributors may be used to endorse or promote products derived+      from this software without specific prior written permission.++THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS+"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT+LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR+A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT+OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT+LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY+THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

+ README view

@@ -0,0 +1,21 @@+This is a simple library for k-d trees in Haskell, based on the algorithms+at http://en.wikipedia.org/wiki/K-d_tree.++It enables efficient searching through collections of points in O(log N) time+for randomly distributed points, using the nearestNeighbor function.++Here is an example of an interactive session using this module:++[ ~/haskell/KdTree ] ghci+GHCi, version 7.0.3: http://www.haskell.org/ghc/  :? for help+...+Prelude> :m Data.Trees.KdTree +Prelude Data.Trees.KdTree> import Test.QuickCheck+Prelude Data.Trees.KdTree Test.QuickCheck> points <- sample' arbitrary :: IO [Point3d]+...+Prelude Data.Trees.KdTree Test.QuickCheck> let tree = fromList points+Prelude Data.Trees.KdTree Test.QuickCheck> nearestNeighbor tree (head points)+Just (Point3d {p3x = 0.0, p3y = 0.0, p3z = 0.0})+Prelude Data.Trees.KdTree Test.QuickCheck> head points+Point3d {p3x = 0.0, p3y = 0.0, p3z = 0.0}+

+ Setup.hs view

@@ -0,0 +1,2 @@+import Distribution.Simple+main = defaultMain