diff --git a/Data/Octree/Internal.hs b/Data/Octree/Internal.hs new file mode 100644 --- /dev/null +++ b/Data/Octree/Internal.hs @@ -0,0 +1,299 @@ +{-# LANGUAGE ScopedTypeVariables, RecordWildCards #-} +module Data.Octree.Internal(Vector3(..), dist, + Octree(..), lookup, nearest, withinRange, fromList, toList, insert, + -- internal + ODir, + octreeStep, octantDistance, splitBy', joinStep, splitStep, allOctants, octantDistance', + cmp, + pickClosest + ) where + +import Data.Vector.V3 +import Data.Vector.Class + +--import Text.Show +import Prelude hiding(lookup) +import Data.List(sort, sortBy) +import Data.Maybe(maybeToList, listToMaybe) +import Data.Bits((.&.)) +import Test.QuickCheck.All(quickCheckAll) +import Test.QuickCheck.Arbitrary + +-- | norm of a vector +-- | distance between two vectors +norm :: Vector3 -> Double +norm a = a `vdot` a + +-- | distance between two vectors +dist :: Vector3 -> Vector3 -> Double +dist u v = norm (u - v) + +data Octree a = Node { split :: Vector3, + nwu, nwd, neu, ned, swu, swd, seu, sed :: Octree a } | + Leaf { unLeaf :: [(Vector3, a)] } deriving (Show) + +instance Functor Octree where + fmap f (Leaf l) = Leaf . fmap (\(c, a) -> (c, f a)) $ l + fmap f (Node { split = sp, + nwu = anwu, + nwd = anwd, + neu = aneu, + ned = aned, + swu = aswu, + swd = aswd, + seu = aseu, + sed = ased }) = Node { split = sp, + nwu = fmap f anwu, + nwd = fmap f anwd, + neu = fmap f aneu, + ned = fmap f aned, + swu = fmap f aswu, + swd = fmap f aswd, + seu = fmap f aseu, + sed = fmap f ased } + +-- | Enumerated type to indicate octants in 3D-space relative to given center. +data ODir = SWD | SED | NWD | NED | SWU | SEU | NWU | NEU deriving (Eq, Ord, Enum, Show, Bounded) + +-- | Internal method that gives octant of a first vector relative to the second vector as a center. +cmp :: Vector3 -> Vector3 -> ODir +cmp ca cb = joinStep (cx, cy, cz) + where cx = v3x ca >= v3x cb + cy = v3y ca >= v3y cb + cz = v3z ca >= v3z cb + +-- | Internal method that joins result of three coordinate comparisons and makes an octant name `ODir` +joinStep :: (Enum a1, Enum a3, Enum a2, Enum a) => (a1, a2, a3) -> a +joinStep (cx, cy, cz) = toEnum (fromEnum cx + 2 * fromEnum cy + 4 * fromEnum cz) + +-- | This function converts octant name to a function that steps down in an Octree towards this octant +octreeStep :: Octree a -> ODir -> Octree a +octreeStep ot NWU = nwu ot +octreeStep ot NWD = nwd ot +octreeStep ot NEU = neu ot +octreeStep ot NED = ned ot +octreeStep ot SWU = swu ot +octreeStep ot SWD = swd ot +octreeStep ot SEU = seu ot +octreeStep ot SED = sed ot + +-- | Function that splits octant name into three boolean values, depending of sign of a relative coordinate in that octant. +-- | (Coordinate is relative to a split point within Octree.) +splitStep :: ODir -> (Bool, Bool, Bool) +splitStep step = ((val .&. 1) == 1, (val .&. 2) == 2, (val .&. 4) == 4) + where val = fromEnum step + +-- | Internal function that finds a lower bound for a distance between a point of relative coordinates, +-- | and octant of given name. It works only when coordinates of a given point are all positive. +-- | It is used only by `octantDistance`, which respectively changes octant name depending of signs of +-- | relative coordinates. +-- here we assume that a, b, c > 0 (otherwise we will take abs, and correspondingly invert results) +-- same octant +-- dp = difference between given point and the center of Octree node +octantDistance' :: Vector3 -> ODir -> Scalar +octantDistance' dp NEU = 0.0 +-- adjacent by plane +octantDistance' dp NWU = v3x dp +octantDistance' dp SEU = v3y dp +octantDistance' dp NED = v3z dp +-- adjacent by edge +octantDistance' dp SWU = sqrt ( v3x dp * v3x dp + v3y dp * v3y dp) +octantDistance' dp SED = sqrt ( v3y dp * v3y dp + v3z dp * v3z dp) +octantDistance' dp NWD = sqrt ( v3x dp * v3x dp + v3z dp * v3z dp) +-- adjacent by point +octantDistance' dp SWD = norm dp + +-- | List of all octant names. +allOctants :: [ODir] +allOctants = [minBound..maxBound] + +-- | Internal function that makes code clearer. +xor :: Bool -> Bool -> Bool +xor = (/=) + +-- | Finds a minimum bounds for a distance between a given point +-- | in relative coordinates and a given octant. +octantDistance :: Vector3 -> ODir -> Double +octantDistance dp odir = octantDistance' (abs dp) (toggle dp odir) + +-- | Toggles octant names depending on a signs of vector coordinates +-- | for use in octantDistance. +toggle :: Vector3 -> ODir -> ODir +toggle dp odir = + joinStep ((v3x dp >= 0) `xor` not u, + (v3y dp >= 0) `xor` not v, + (v3z dp >= 0) `xor` not w) + where (u, v, w) = splitStep odir + +-- | Given a point in relative coordinates, gives list of all octants and minimum distances from this point. +octantDistances :: Vector3 -> [(ODir, Double)] +octantDistances dp = [(o, octantDistance dp o) | o <- allOctants] + +-- | splits a list of vectors and "payload" tuples +-- | into a tuple with elements destined for different octants. +-- FIXME: VERY IMPORTANT - add prop_splitBy vs cmp +splitBy :: Vector3 -> [(Vector3, a)] -> ([(Vector3, a)], + [(Vector3, a)], + [(Vector3, a)], + [(Vector3, a)], + [(Vector3, a)], + [(Vector3, a)], + [(Vector3, a)], + [(Vector3, a)]) +splitBy _splitPoint [] = ([], [], [], [], [], [], [], []) +splitBy splitPoint ((pt@(coord, a)):aList) = + case i of + SWD -> (pt:swd, sed, nwd, ned, swu, seu, nwu, neu) + SED -> ( swd, pt:sed, nwd, ned, swu, seu, nwu, neu) + NWD -> ( swd, sed, pt:nwd, ned, swu, seu, nwu, neu) + NED -> ( swd, sed, nwd, pt:ned, swu, seu, nwu, neu) + SWU -> ( swd, sed, nwd, ned, pt:swu, seu, nwu, neu) + SEU -> ( swd, sed, nwd, ned, swu, pt:seu, nwu, neu) + NWU -> ( swd, sed, nwd, ned, swu, seu, pt:nwu, neu) + NEU -> ( swd, sed, nwd, ned, swu, seu, nwu, pt:neu) + where i = cmp coord splitPoint + (swd, sed, nwd, ned, swu, seu, nwu, neu) = splitBy splitPoint aList + +{- +-- | Internal Sums a list of vectors +sumVector3s [(coord, _)] = coord +sumVector3s ((coord, _):aList) = coord + sumVector3s aList +-} +-- | Computes a center of mass for a given list of vectors - used to find a splitPoint. +massCenter :: Fractional a => [(a, b)] -> a +massCenter aList = sum (map fst aList) / count + where + count = fromInteger . toInteger . length $ aList + +-- | Helper function to map over an 8-element tuple +tmap :: (t -> t1)-> (t, t, t, t, t, t, t, t)-> (t1, t1, t1, t1, t1, t1, t1, t1) +tmap t (a, b, c, d, e, f, g, h) = (t a, t b, t c, t d, t e, t f, t g, t h) + +-- | Maximum number of elements before Octree leaf is split. +leafLimit :: Int +leafLimit = 16 + +-- | Creates an Octree from a list of (index, payload) tuples. +fromList :: [(Vector3, a)] -> Octree a +fromList aList = if length aList <= leafLimit + then Leaf aList + else let splitPoint :: Vector3 = massCenter aList + in splitBy' fromList splitPoint aList +-- | Internal method, that splits a list into octants depending on coordinates, +-- | and then applies a specified function to each of these sublists, +-- | in order to create subnodes of the Octree +splitBy' :: ([(Vector3, a)] -> Octree a1)-> Vector3-> [(Vector3, a)]-> Octree a1 +splitBy' f splitPoint aList = Node { split = splitPoint, + nwu = tnwu, + nwd = tnwd, + neu = tneu, + ned = tned, + swu = tswu, + swd = tswd, + seu = tseu, + sed = tsed } + where + (tswd, tsed, tnwd, tned, tswu, tseu, tnwu, tneu) = tmap f $ splitBy splitPoint aList +-- TODO: use arrays for memory savings + +-- | Internal method that prepends contents of the given subtree to a list +-- | given as argument. +toList' :: Octree t -> [(Vector3, t)] -> [(Vector3, t)] +toList' (Leaf l ) tmp = l ++ tmp +toList' (Node { nwu = a, + nwd = b, + neu = c, + ned = d, + swu = e, + swd = f, + seu = g, + sed = h }) tmp = foldr toList' tmp [a, b, c, d, e, f, g, h] +-- | Creates an Octree from list, trying to keep split points near centers +-- | of mass for each subtree. +toList :: Octree t -> [(Vector3, t)] +toList t = toList' t [] + +-- | Finds a path to a Leaf where a given point should be, +-- | and returns a list of octant names. +pathTo :: Vector3 -> Octree a -> [ODir] +pathTo pt (Leaf _) = [] +pathTo pt node = aStep : pathTo pt (octreeStep node aStep) + where aStep = cmp pt (split node) + +-- | Applies a given function to a node specified by a path (list of octant names), +-- | and then returns a modified Octree. +applyByPath :: (Octree a -> Octree a) -> [ODir] -> Octree a -> Octree a +applyByPath f [] ot = f ot +applyByPath f (step:path) node = case step of + NWU -> node{ nwu = applyByPath f path (nwu node) } + NWD -> node{ nwd = applyByPath f path (nwd node) } + NEU -> node{ neu = applyByPath f path (neu node) } + NED -> node{ ned = applyByPath f path (ned node) } + SWU -> node{ swu = applyByPath f path (swu node) } + SWD -> node{ swd = applyByPath f path (swd node) } + SEU -> node{ seu = applyByPath f path (seu node) } + SED -> node{ sed = applyByPath f path (sed node) } + +-- | Inserts a point into an Octree. +-- | NOTE: insert accepts duplicate points, but lookup would not find them - use withinRange in such case. +insert :: (Vector3, a) -> Octree a -> Octree a +insert (pt, dat) ot = applyByPath insert' path ot + where path = pathTo pt ot + insert' (Leaf l) = fromList ((pt, dat) : l) + insert' _ = error "Impossible in insert'" + + +-- | Internal: finds candidates for nearest neighbour lazily for each octant; +-- | they are returned in a list of (octant, min. bound for distance, Maybe candidate) tuples. +candidates' :: Vector3 -> Octree a -> [(ODir, Double, [(Vector3, a)])] +candidates' pt (Leaf l) = [] +candidates' pt node = map findCandidates . sortBy compareDistance . octantDistances $ pt - split node + where + findCandidates (octant, d) = (octant, d, maybeToList . pickClosest pt . toList . octreeStep node $ octant) + compareDistance a b = compare (snd a) (snd b) + +-- | Finds a given point, if it is in the tree. +lookup :: Octree a -> Vector3 -> Maybe (Vector3, a) +lookup (Leaf l) pt = listToMaybe . filter ((==pt) . fst) $ l +lookup node pt = flip lookup pt . octreeStep node . cmp pt . split $ node + +-- | Finds nearest neighbour for a given point. +nearest :: Octree a -> Vector3 -> Maybe (Vector3, a) +nearest (Leaf l) pt = pickClosest pt l +nearest node pt = selectFrom candidates + where candidates = map findCandidate . sortBy compareDistance . octantDistances $ pt - split node + compareDistance a b = compare (snd a) (snd b) + findCandidate (octant, d) = (maybeToList . nearest' $ octreeStep node $ octant, d) + selectFrom (([], _d) : cs) = selectFrom cs + selectFrom (([best], _d) : cs) = selectFrom' best cs + selectFrom [] = Nothing + nearest' n = nearest n pt + + + selectFrom' best (([], d) : cs) = selectFrom' best cs + -- TODO: FAILS: shortcut guard to avoid recursion over whole structure (since d is bound for distance within octant): + selectFrom' best ((c, d) : cs) | d > dist pt (fst best) = Just best + selectFrom' best (([next], d) : cs) = selectFrom' nextBest cs + where nextBest = if dist pt (fst best) <= dist pt (fst next) + then best + else next + selectFrom' best [] = Just best + +-- | Internal method that picks from a given list a point closest to argument, +pickClosest :: Vector3 -> [(Vector3, t)] -> Maybe (Vector3, t) +pickClosest pt [] = Nothing +pickClosest pt (a:as) = Just $ foldr (pickCloser pt) a as +pickCloser pt va@(a, _a) vb@(b, _b) = if dist pt a <= dist pt b + then va + else vb + +-- | Returns all points within Octree that are within a given distance from argument. +withinRange :: Scalar -> Vector3 -> Octree a -> [(Vector3, a)] +withinRange r pt (Leaf l) = filter (\(lpt, _) -> dist pt lpt <= r) l +withinRange r pt node = (concat . -- merge results + map recurseOctant . -- recurse over remaining octants + filter ((<=r) . snd) . -- discard octants that are out of range + octantDistances $ pt - split node) -- find octant distances + where + recurseOctant (octant, _d) = withinRange r pt . octreeStep node $ octant + diff --git a/Octree.cabal b/Octree.cabal --- a/Octree.cabal +++ b/Octree.cabal @@ -1,10 +1,10 @@ name: Octree -version: 0.2 +version: 0.2.0.1 stability: beta homepage: https://github.com/mgajda/octree package-url: http://hackage.haskell.org/package/octree synopsis: Simple unbalanced Octree for storing data about 3D points -description: Octree data structure is relatively shallow data structure for space partitioning. +description: Octree data structure is relatively shallow data structure for space partitioning. [This version is not different that 0.2.1.] category: Data license: BSD3 license-file: LICENSE @@ -27,11 +27,12 @@ Library build-depends: base>=4.0, base <4.6, AC-Vector >= 2.3.0, QuickCheck >= 2.4.0 exposed-modules: Data.Octree + other-modules: Data.Octree.Internal exposed: True extensions: ScopedTypeVariables -- I do not know how to make it work. To test use: runghc tests/test_Octree.hs ---Test-suite test_Octree --- Type: exitcode-stdio-1.0 --- Build-Depends: base>=4.0, base<4.6 --- Main-is: tests/test_Octree.hs +Test-suite test_Octree + Type: exitcode-stdio-1.0 + Build-depends: base>=4.0, base <4.6, AC-Vector >= 2.3.0, QuickCheck >= 2.4.0 + Main-is: tests/test_Octree.hs diff --git a/tests/test_Octree.hs b/tests/test_Octree.hs new file mode 100644 --- /dev/null +++ b/tests/test_Octree.hs @@ -0,0 +1,97 @@ +{-# LANGUAGE TemplateHaskell #-} +module Main(main) where + +import Data.Octree.Internal +--import Data.Octree() -- test that interface module is not broken +import Prelude hiding(lookup) +import Data.List(sort, sortBy) + +import Test.QuickCheck.All(quickCheckAll) +import Test.QuickCheck.Arbitrary + +import Data.Vector.Class + +-- | For testing purposes +instance Ord Vector3 where + a `compare` b = pointwiseOrd $ zipWith compare (vunpack a) (vunpack b) + +pointwiseOrd [] = EQ +pointwiseOrd (LT:cs) = LT +pointwiseOrd (GT:cs) = GT +pointwiseOrd (EQ:cs) = pointwiseOrd cs + +instance Arbitrary Vector3 where + arbitrary = do x <- arbitrary + y <- arbitrary + z <- arbitrary + return $ Vector3 x y z + +-- | These are tests for internal helper functions: + +-- for easier testing +origin :: Vector3 +origin = fromInteger 0 + +prop_cmp1 a b = cmp a b == joinStep (dx >= 0, dy >= 0, dz >= 0) + where Vector3 dx dy dz = a - b + +prop_cmp2 a = cmp a origin == joinStep (dx >= 0, dy >= 0, dz >= 0) + where Vector3 dx dy dz = a + +prop_stepDescription a b = splitStep (cmp a b) == (v3x a >= v3x b, v3y a >= v3y b, v3z a >= v3z b) + +prop_octantDistanceNoGreaterThanInterpointDistance0 ptA ptB = triangleInequality + where triangleInequality = (octantDistance' aptA (cmp ptB origin)) <= (dist aptA ptB) + aptA = abs ptA + +prop_octantDistanceNoGreaterThanInterpointDistance ptA ptB vp = triangleInequality || sameOctant + where triangleInequality = (octantDistance (ptA - vp) (cmp ptB vp)) <= (dist ptA ptB) + sameOctant = (cmp ptA vp) == (cmp ptB vp) + +prop_octantDistanceNoGreaterThanInterpointDistanceZero ptA ptB = triangleInequality || sameOctant + where triangleInequality = (octantDistance ptA (cmp ptB origin)) <= (dist ptA ptB) + sameOctant = (cmp ptA origin) == (cmp ptB origin) + +prop_octantDistanceNoGreaterThanCentroidDistance pt vp = all testFun allOctants + where testFun odir = (octantDistance (pt - vp) odir) <= dist pt vp + +prop_splitByPrime splitPt pt = (unLeaf . octreeStep ot . cmp pt $ splitPt) == [arg] + where ot = splitBy' Leaf splitPt [arg] + arg = (pt, dist pt splitPt) + + +prop_pickClosest :: (Eq a) => [(Vector3, a)] -> Vector3 -> Bool +prop_pickClosest l pt = pickClosest pt l == naiveNearest pt l + +-- | These are tests for exposed functions: + +prop_lookup l = all isIn l + where ot = fromList l + isIn x = lookup ot (fst x) == Just x + +prop_fromToList l = sort l == (sort . toList . fromList $ l) +prop_insertionPreserved l = sort l == (sort . toList . foldr insert (Leaf []) $ l) +prop_nearest l pt = nearest (fromList l) pt == naiveNearest pt l +prop_naiveWithinRange r l pt = naiveWithinRange r pt l == (sort . map fst . withinRange r pt . fromList . tuplify pt $ l) + +tuplify pt = map (\a -> (a, dist pt a)) + +compareDistance pt a b = compare (dist pt (fst a)) (dist pt (fst b)) + +naiveNearest pt l = if byDist == [] then Nothing else Just . head $ byDist + where byDist = sortBy (compareDistance pt) l + +naiveWithinRange r pt l = sort . filter (\p -> dist pt p <= r) $ l + +-- unfortunately there is no Arbitrary for (a -> b) +-- since generic properties are quite common, I wonder how to force Quickcheck to default something reasonable? +prop_fmap1 l = genericProperty_fmap (+1) l +prop_fmap2 l = genericProperty_fmap (*2) l +prop_fmap3 l = genericProperty_fmap show l + +genericProperty_fmap f l = (sort . mapSnd f $ l) == (sort . toList . fmap f . fromList $ l) + where + mapSnd :: (a -> b) -> [(c, a)] -> [(c, b)] + mapSnd f l = map (\(c, a) -> (c, f a)) l + +main = do $quickCheckAll