diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,30 @@ +Copyright Author name here (c) 2020 + +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 Author name here 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. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,19 @@ +# vp-tree + +Vantage point trees, as described in + +Data structures and algorithms for nearest neighbor search in general metric spaces - P. N. Yianilos + +http://web.cs.iastate.edu/~honavar/nndatastructures.pdf + + +## Usage + +Import 'Data.VPTree', which also contains usage instructions and comments + + +## Benchmarks + +Cumulative memory usage and garbage collection cycles : + + $ stack bench -- vp-tree:bench-memory 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/bench/memory/Main.hs b/bench/memory/Main.hs new file mode 100644 --- /dev/null +++ b/bench/memory/Main.hs @@ -0,0 +1,83 @@ +{-# options_ghc -Wno-unused-imports #-} +module Main where + +-- vector +import qualified Data.Vector as V (Vector, map) +-- weigh +import Weigh (mainWith, wgroup, func) + +import Data.VPTree.Build (build) +import Data.VPTree.Draw (draw) +import Data.VPTree.Internal (VT, VPTree, withST, withST_, withIO) +import Data.VPTree.Query (range, distances) +import Data.VPTree.TestData (buildP, binDiskSamples, gaussMixSamples, P(..)) + + +main :: IO () +main = mainWith $ do + wgroup "Data.VPTree.Build" $ do + let + go n = buildP (binDiskSamples n) + func "build 10" go 10 + func "build 100" go 100 + func "build 1000" go 1000 + func "build 10.000" go 10000 + -- func "build 100.000" go 100000 + wgroup "Data.VPTree.Query : index size 100" $ do + let + index = buildIndex binDiskSamples 100 + go n = queryIndex index binDiskSamples 1.0 n + func "range 10" go 10 + func "range 100" go 100 + func "range 1000" go 1000 + func "range 10.000" go 10000 + wgroup "Data.VPTree.Query : index size 1000" $ do + let + index = buildIndex binDiskSamples 1000 + go n = queryIndex index binDiskSamples 1.0 n + func "range 10" go 10 + func "range 100" go 100 + func "range 1000" go 1000 + func "range 10.000" go 10000 + wgroup "Data.VPTree.Query : index size 10.000" $ do + let + index = buildIndex binDiskSamples 10000 + go n = queryIndex index binDiskSamples 1.0 n + func "range 10" go 10 + func "range 100" go 100 + func "range 1000" go 1000 + func "range 10.000" go 10000 + + +buildIndex :: (t -> V.Vector P) -> t -> VPTree Double P +buildIndex genTree m = buildP (genTree m) + +queryIndex :: (Num p1, Ord p1) => + VPTree p1 a -> (p2 -> V.Vector a) -> p1 -> p2 -> V.Vector [(p1, a)] +queryIndex index genData thr n = V.map (range index thr) qrys + where + qrys = genData n + +-- rangeWith :: (p1 -> V.Vector P) +-- -> (p2 -> V.Vector P) +-- -> Double +-- -> p1 +-- -> p2 +-- -> V.Vector [(Double, P)] +-- rangeWith genTree genData thr m n = V.map (range tree thr) qrys +-- where +-- tree = buildP (genTree m) +-- qrys = genData n + + + +-- main = +-- mainWith (do func "integers count 0" count 0 +-- func "integers count 1" count 1 +-- func "integers count 2" count 2 +-- func "integers count 3" count 3 +-- func "integers count 10" count 10 +-- func "integers count 100" count 100) +-- where count :: Integer -> () +-- count 0 = () +-- count a = count (a - 1) diff --git a/src/Data/VPTree.hs b/src/Data/VPTree.hs new file mode 100644 --- /dev/null +++ b/src/Data/VPTree.hs @@ -0,0 +1,72 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE RankNTypes #-} +{-# language BangPatterns #-} +{-# language DeriveFunctor, DeriveFoldable, DeriveTraversable, GeneralizedNewtypeDeriving #-} +{-# language LambdaCase #-} +{-# language DeriveDataTypeable #-} +{-# language DeriveGeneric #-} + +{-# options_ghc -Wno-type-defaults #-} +{-# options_ghc -Wno-unused-top-binds #-} +{-# options_ghc -Wno-unused-imports #-} +{-# options_ghc -Wno-name-shadowing #-} +{- | + +This library provides an implementation of Vantage Point Trees [1], a data structure useful for indexing data points that exist in some metric space. + +The current implementation is not particolarly optimized and assumes the data resides entirely in memory but it seems to work decently well for index sizes in the 10's of thousands. + += Usage + +* 'range' : construct an index from a dataset and a distance function + +* 'range' : find points in the index that lie within a given distance from the query + + += References + +1) P. N. Yianilos - Data structures and algorithms for nearest neighbor search in general metric spaces - http://web.cs.iastate.edu/~honavar/nndatastructures.pdf +-} +module Data.VPTree + (VPTree + -- * Construction + , build + -- * Query + , range + -- , nearest + -- * Utilities + -- ** Rendering trees + , draw + ) + where + +import Control.Applicative (Alternative(..)) +import Control.Monad.IO.Class (MonadIO(..)) +-- import Data.Ord (Down(..)) +import Data.Word (Word32) +-- import Control.Exception (Exception(..)) +import Control.Monad.ST (ST, runST) +import Text.Printf (PrintfArg, printf) + +-- -- deepseq +-- import Control.DeepSeq (NFData(..)) +-- mtl +import Control.Monad.Writer (MonadWriter(..)) +-- mwc-probability +import qualified System.Random.MWC.Probability as P (Gen, Prob, withSystemRandom, asGenIO, GenIO, create, initialize, sample, samples, normal, bernoulli, uniformR) +-- primitive +import Control.Monad.Primitive (PrimMonad(..), PrimState) +-- transformers +import Control.Monad.Trans.Maybe (MaybeT(..), runMaybeT) +import Control.Monad.Trans.Writer (WriterT(..), runWriterT, execWriterT) +-- vector +import qualified Data.Vector as V (Vector, map, filter, length, toList, replicate, partition, zipWith, head, tail, fromList, thaw, freeze, (!), foldl) +import qualified Data.Vector.Generic as VG (Vector(..)) + +-- import qualified Data.MaxPQ as MQ (MaxPQ, empty, insert, size, findMax, toList) + +import Data.VPTree.Internal (VT, VPTree, withST, withST_, withIO) +import Data.VPTree.Build (build, buildVT) +import Data.VPTree.Query (range) +import Data.VPTree.Draw (draw) + diff --git a/src/Data/VPTree/Build.hs b/src/Data/VPTree/Build.hs new file mode 100644 --- /dev/null +++ b/src/Data/VPTree/Build.hs @@ -0,0 +1,241 @@ +{-# options_ghc -Wno-unused-imports #-} +{-# options_ghc -Wno-type-defaults #-} +module Data.VPTree.Build (build + -- * Internal + , buildVT + ) where + +import Control.Monad.ST (ST, runST) +import qualified Data.Foldable as F (Foldable(..)) +import Data.Foldable (foldlM) +import Data.Maybe (fromMaybe) + +-- containers +import qualified Data.Set as S (Set, fromList, difference) +-- import qualified Data.Sequence as SQ (Seq) +-- deepseq +-- import Control.DeepSeq (NFData (rnf)) +-- mwc-probability +import qualified System.Random.MWC.Probability as P (Gen, Prob, withSystemRandom, asGenIO, GenIO, create, initialize, sample, samples, normal, bernoulli) +-- primitive +import Control.Monad.Primitive (PrimMonad(..), PrimState) +-- sampling +import Numeric.Sampling (sample) +-- vector +import qualified Data.Vector as V (Vector, map, filter, length, toList, replicate, partition, zipWith, head, tail, fromList, thaw, freeze, (!), foldl) +-- import qualified Data.Vector.Generic as VG (Vector(..)) +-- import Data.Vector.Generic.Mutable (MVector) +-- vector-algorithms +import qualified Data.Vector.Algorithms.Merge as V (sort, Comparison) + +import Data.VPTree.Internal (VT(..), VPTree(..), withST_) + +-- * Construction + +-- | Build a 'VPTree' +-- +-- The supplied distance function @d@ must satisfy the definition of a metric, i.e. +-- +-- * identity of indiscernible elements : \( d(x, y) = 0 \leftrightarrow x \equiv y \) +-- +-- * symmetry : \( d(x, y) = d(y, x) \) +-- +-- * triangle inequality : \( d(x, y) + d(y, z) >= d(x, z) \) +-- +-- The current implementation makes multiple passes over the whole dataset, which is why the indexing data must all be present in memory (currently packed as a 'V.Vector'). +-- +-- Implementation detail : construction of a VP-tree requires a randomized algorithm, but we run that in the ST monad so the result is pure. +build :: (RealFrac p, Floating d, Ord d, Eq a) => + (a -> a -> d) -- ^ distance function + -> p -- ^ proportion of remaining dataset to sample at each level, \(0 < p <= 1 \) + -> V.Vector a -- ^ dataset used for constructing the index + -> VPTree d a +build distf prop xss = withST_ $ \gen -> do + vt <- buildVT distf prop xss gen + pure $ VPT vt distf + + +-- | Build a VP-tree with the given distance function +buildVT :: (PrimMonad m, RealFrac b, Floating d, Eq a, Ord d) => + (a -> a -> d) -- ^ distance function + -> b -- ^ proportion of remaining dataset to sample at each level + -> V.Vector a -- ^ dataset + -> P.Gen (PrimState m) -- ^ PRNG + -> m (VT d a) +buildVT distf prop xss gen = go xss + where + go xs + | length xs < 10 = pure $ Tip xs + | otherwise = do + (vp, xs') <- selectVP distf prop xs gen + let + mu = median $ V.map (`distf` vp) xs' -- median distance to the vantage point + (ll, rr) = V.partition (\x -> distf x vp < mu) xs' + + ltree <- go ll + rtree <- go rr + pure $ Bin mu vp ltree rtree + + +-- | Select a vantage point +selectVP :: (PrimMonad m, RealFrac b, Floating d, Ord d) => + (a -> a -> d) + -> b -> V.Vector a -> P.Gen (PrimState m) -> m (a, V.Vector a) +selectVP distf prop xs gen = do + (pstart, pstail, pscl) <- vpRandSplitInit n xs gen + let pickMu (spread_curr, p_curr, acc) p = do + ds <- sampleId n2 pscl gen -- sample n2 < n points from pscl + let + spread = varianceWrt distf p (V.fromList ds) + if spread > spread_curr + then pure (spread, p, p_curr : acc) + else pure (spread_curr, p_curr, p : acc) + (vp, vrest) <- tail3 <$> foldlM pickMu (0, pstart, mempty) pstail + pure (vp, V.fromList vrest) + where + n = max 1 $ floor (prop * fromIntegral ndata) + n2 = max 1 $ floor (prop * fromIntegral n) + ndata = length xs -- size of dataset at current level + tail3 (_, x, xs) = (x, xs) + + +-- | sample the initialization for picking a vantage point +-- +-- samples a random split of the input dataset, and from the first half further samples a head element, which will be used as candidate vantage point +vpRandSplitInit :: PrimMonad m => + Int + -> V.Vector a -- ^ cannot be less than 3 elements + -> P.Gen (PrimState m) + -> m (a, [a], [a]) -- (head of C, tail of C, complement of C) +vpRandSplitInit n sset gen = do + (ps, psc) <- uniformSplit n sset gen + (pstartv, pstail) <- randomSplit 0.5 1 ps gen -- Pick a random starting point from ps + let + -- this is load-bearing, do not change + pstart = if null pstartv then pstail !! 1 else head pstartv + pure (pstart, pstail, psc) + +-- | Split a dataset in two, returning a ~ uniform sample +-- +-- the Bernoulli parameter depends on the size of the desired sample and that of the dataset +uniformSplit :: (PrimMonad m, Foldable t) => + Int -> t a -> P.Gen (PrimState m) -> m ([a], [a]) +uniformSplit n vv = randomSplit p n vv + where + p = 1 - (fromIntegral n / fromIntegral (length vv)) + +-- | Sample a random split of the dataset in a single pass, by repeatedly tossing a coin +-- +-- Invariant : the concatenation of the two resulting vectors is a permutation of the input vector +-- +-- NB : the second vector in the result tuple will be empty if the requested sample size is larger than the input vector +randomSplit :: (Foldable t, PrimMonad m) => + Double -- ^ Bernoulli parameter + -> Int -- ^ Size of sample + -> t a -- ^ dataset + -> P.Gen (PrimState m) -- ^ PRNG + -> m ([a], [a]) +randomSplit p n vv = P.sample $ foldlM insf ([], []) vv + where + insf (al, ar) x = do + coin <- P.bernoulli p + if length al == n || coin + then pure (al, x : ar) + else pure (x : al, ar) + + + + +-- | Sample _without_ replacement. Returns the input list if the required sample size is too large +sampleId :: (PrimMonad m, Foldable t) => + Int -- ^ Size of sample + -> t a + -> P.Gen (PrimState m) + -> m [a] +sampleId n xs g = fromMaybe (F.toList xs) <$> sample n xs g +{-# INLINE sampleId #-} + +-- | Variance of the distance btw the dataset and a given query point +-- +-- NB input vector must have at least 1 element +varianceWrt :: (Floating a, Ord a) => + (t -> p -> a) -- ^ distance function + -> p -- ^ query point + -> V.Vector t + -> a +varianceWrt distf p ds = variance dists (V.replicate n2 mu) where + dists = V.map (`distf` p) ds + mu = median dists + n2 = V.length ds +{-# INLINE varianceWrt #-} + +-- | NB input vector must have at least 1 element +median :: Ord a => V.Vector a -> a +median xs + | null xs = error "median : input array must have at least 1 element" + | n == 1 = V.head xs + | otherwise = sortV xs V.! floor (fromIntegral n / 2) + where n = length xs +{-# INLINE median #-} + +variance :: (Floating a) => V.Vector a -> V.Vector a -> a +variance xs mus = mean $ V.zipWith sqdiff xs mus + where + sqdiff x y = (x - y) ** 2 +{-# INLINE variance #-} + +mean :: (Fractional a) => V.Vector a -> a +mean xs = sum xs / fromIntegral (length xs) +{-# INLINE mean #-} + +sortV :: Ord a => V.Vector a -> V.Vector a +sortV v = runST $ do + vm <- V.thaw v + V.sort vm + V.freeze vm +{-# INLINE sortV #-} + + + + + +-- -- OLD + + +-- selectVP :: (PrimMonad m, RealFrac b, Ord d, Floating d) => +-- (a -> a -> d) -- ^ distance function +-- -> b -- ^ proportion of dataset to sample +-- -> V.Vector a -- ^ dataset +-- -> P.Gen (PrimState m) +-- -> m a +-- selectVP distf prop sset gen = do +-- (pstart, pstail, pscl) <- vpRandSplitInit n sset gen +-- let pickMu (spread_curr, p_curr) p = do +-- ds <- sampleId n2 pscl gen -- sample n2 < n points from pscl +-- let +-- spread = varianceWrt distf p (V.fromList ds) +-- if spread > spread_curr +-- then pure (spread, p) +-- else pure (spread_curr, p_curr) +-- snd <$> foldlM pickMu (0, pstart) pstail +-- where +-- n = floor (prop * fromIntegral ndata) +-- n2 = floor (prop * fromIntegral n) +-- ndata = length sset -- size of dataset at current level + +-- randomSplit :: (PrimMonad f) => +-- Int -- ^ Size of sample +-- -> V.Vector a -- ^ dataset +-- -> P.Gen (PrimState f) -- ^ PRNG +-- -> f (V.Vector a, V.Vector a) +-- randomSplit n vv gen = split <$> sampleId n ixs gen +-- where +-- split xs = (vxs, vxsc) +-- where +-- ixss = S.fromList xs +-- ixsc = S.fromList ixs `S.difference` ixss +-- vxs = pickItems ixss +-- vxsc = pickItems ixsc +-- m = V.length vv +-- ixs = [0 .. m - 1] +-- pickItems = V.fromList . foldl (\acc i -> vv V.! i : acc) [] diff --git a/src/Data/VPTree/Draw.hs b/src/Data/VPTree/Draw.hs new file mode 100644 --- /dev/null +++ b/src/Data/VPTree/Draw.hs @@ -0,0 +1,50 @@ +{-# LANGUAGE LambdaCase #-} +{-# options_ghc -Wno-unused-imports #-} +module Data.VPTree.Draw ( + draw, drawVT + -- * helpers + , toStringVT + ) where + +import Text.Printf (PrintfArg, printf) +import Data.VPTree.Internal (VPTree(..), VT(..)) + +-- boxes +import qualified Text.PrettyPrint.Boxes as B (Box, render, emptyBox, vcat, hcat, text, top, bottom, center1) + +-- | Render a tree to stdout +-- +-- Useful for debugging +-- +-- This should be called only for small trees, otherwise the printed result quickly overflows the screen and becomes hard to read. +-- +-- NB : prints distance information rounded to two decimal digits +draw :: (Show a, PrintfArg d) => VPTree d a -> IO () +draw = drawVT . vpTree + +drawVT :: (Show a, PrintfArg d) => VT d a -> IO () +drawVT = putStrLn . toStringVT + +toStringVT :: (Show a, PrintfArg d) => VT d a -> String +toStringVT = B.render . toBox + +toBox :: (Show a, PrintfArg d) => VT d a -> B.Box +toBox = \case + (Bin d x tl tr) -> + txt (node x d) `stack` (toBox tl `byside` toBox tr) + Tip x -> txt $ show x + -- Nil -> txt "*" + where + node x d = printf "%s,%5.2f" (show x) d + -- nodeBox x d = + -- txt (printf "%s,%5.2f" (show x) d) + +txt :: String -> B.Box +txt t = spc `byside` B.text t `byside` spc + where spc = B.emptyBox 1 1 + +byside :: B.Box -> B.Box -> B.Box +byside l r = B.hcat B.top [l, r] + +stack :: B.Box -> B.Box -> B.Box +stack t b = B.vcat B.center1 [t, b] diff --git a/src/Data/VPTree/Internal.hs b/src/Data/VPTree/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Data/VPTree/Internal.hs @@ -0,0 +1,118 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE RankNTypes #-} +{-# language DeriveGeneric #-} +{-# language LambdaCase #-} +{-# language DeriveFoldable, DeriveTraversable, DeriveFunctor #-} +module Data.VPTree.Internal where + +import Control.Monad.ST (ST, runST) +import Data.Word (Word32) +import GHC.Generics (Generic(..)) + +-- deepseq +import Control.DeepSeq (NFData(..)) +-- mwc-probability +import qualified System.Random.MWC.Probability as P (Gen, withSystemRandom, asGenIO, GenIO, create, initialize) +-- serialise +import Codec.Serialise (Serialise(..)) +-- vector +import qualified Data.Vector as V (Vector) +import qualified Data.Vector.Generic as VG (Vector(..)) + +-- | Vantage point tree +data VPTree d a = VPT { + vpTree :: VT d a + , vptDistFun :: a -> a -> d -- ^ Distance function used to construct the tree + } deriving (Generic) + +instance (Eq d, Eq a) => Eq (VPTree d a) where + (VPT t1 _) == (VPT t2 _) = t1 == t2 +instance (Show d, Show a) => Show (VPTree d a) where + show (VPT t _) = show t +instance (NFData d, NFData a) => NFData (VPTree d a) where +instance Foldable (VPTree d) where + foldMap f (VPT t _) = foldMap f t + + +-- | Vantage point tree (internal representation) +data VT d a = Bin { + _mu :: !d -- ^ median distance to vantage point + , _vp :: !a -- ^ vantage point + , _near :: !(VT d a) -- ^ points at a distance < mu + , _far :: !(VT d a) -- ^ points farther than mu + } + | Tip (V.Vector a) + deriving (Eq, Generic, Functor, Foldable, Traversable) +instance (Show d, Show a) => Show (VT d a) where + show = \case + -- Nil -> "" + Tip x -> unwords [""] + Bin m v ll rr -> unwords [""] +instance (Serialise d, Serialise a) => Serialise (VT d a) + +instance (NFData d, NFData a) => NFData (VT d a) where + rnf (Bin d x tl tr) = rnf d `seq` rnf x `seq` rnf tl `seq` rnf tr + rnf (Tip x) = rnf x + -- rnf Nil = () + + + +-- | Runs a PRNG action in IO +-- +-- NB : uses 'withSystemRandom' internally +withIO :: (P.GenIO -> IO a) -- ^ Memory bracket for the PRNG + -> IO a +withIO = P.withSystemRandom . P.asGenIO + +-- | Runs a PRNG action in the 'ST' monad, using a fixed seed +-- +-- NB : uses 'P.create' internally +withST_ :: (forall s . P.Gen s -> ST s a) -- ^ Memory bracket for the PRNG + -> a +withST_ st = runST $ do + g <- P.create + st g + +-- | Runs a PRNG action in the 'ST' monad, using a given random seed +-- +-- NB : uses 'P.initialize' internally +withST :: (VG.Vector v Word32) => + v Word32 -- ^ Random seed + -> (forall s . P.Gen s -> ST s a) -- ^ Memory bracket for the PRNG + -> a +withST seed st = runST $ do + g <- P.initialize seed + st g + + +-- + +-- newtype App w m a = App { +-- unApp :: MaybeT (WriterT w m) a +-- } deriving (Functor, Applicative, Monad, Alternative, MonadIO, MonadWriter w) + +-- runApp :: App w m a -> m (Maybe a, w) +-- runApp a = runWriterT $ runMaybeT (unApp a) + +-- runAppST :: (forall s . P.Gen s -> App w (ST s) a) -> (Maybe a, w) +-- runAppST a = withST_ (runApp . a) + +-- -- testApp :: PrimMonad m => P.Gen (PrimState m) -> App m [Double] () +-- testApp g = App $ do +-- z <- P.samples 5 (P.normal 0 1) g +-- tell z +-- pure z + +-- sampleApp :: (Foldable f, PrimMonad m) => +-- Int -> f a -> P.Gen (PrimState m) -> App m [String] [a] +-- sampleApp n ixs g = App $ do +-- zm <- sample n ixs g +-- case zm of +-- Nothing -> do +-- tell ["derp"] +-- empty +-- Just xs -> pure xs + + +-- runAppST :: (forall s . P.Gen s -> WriterT w (ST s) a) -> (a, w) +-- runAppST a = withST_ (runWriterT . a) diff --git a/src/Data/VPTree/Query.hs b/src/Data/VPTree/Query.hs new file mode 100644 --- /dev/null +++ b/src/Data/VPTree/Query.hs @@ -0,0 +1,307 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} +{-# options_ghc -Wno-unused-imports #-} +module Data.VPTree.Query ( + range + -- * Utilities + , distances + ) where + +import Control.Monad.IO.Class (MonadIO(..)) +import Data.Foldable (toList, foldrM, foldlM) + + +-- containers +import Data.Sequence as SQ (Seq) +import Data.Sequence ((|>)) +-- mtl +import Control.Monad.State (MonadState(..)) +-- psqueues +import qualified Data.IntPSQ as PQ (IntPSQ, insert, size, empty, toList, minView) +-- transformers +import Control.Monad.Trans.State (State, evalState, runState) +-- vector +import qualified Data.Vector as V (Vector) + + +import Data.VPTree.Internal (VT(..), VPTree(..)) + +psqList :: (Ord p) => + PQ.IntPSQ p b -> [(p, b)] +psqList q = case PQ.minView q of + Nothing -> mempty + Just (_, p, v, qrest) -> (p, v) : psqList qrest + +-- | All distances to a query point +distances :: VPTree b a + -> a -- ^ query + -> [b] +distances (VPT tt distf) x = map (distf x) $ toList tt + +-- | Range query : find all points in the tree closer to the query point than a given threshold +range :: (Num p, Ord p) => + VPTree p a + -> p -- ^ proximity threshold + -> a -- ^ query point + -> [(p, a)] +range (VPT tt distf) eps x = psqList $ rangeVT eps x distf tt +-- range (VPT tt distf) eps x = rangeVT' eps x distf tt + + +rangeVT :: (Num b, Ord b) => + b -- ^ proximity threshold + -> a -> (a -> a -> b) -> VT b a -> PQ.IntPSQ b a +rangeVT eps x distf = flip evalState 0 . go PQ.empty + where + go acc = \case + Tip ts -> + foldlM insf acc ts + where + insf ac t + | d < eps = do + i <- get + let ac' = PQ.insert i d t ac + put (i + 1) + pure ac' + | otherwise = pure ac + where + d = distf x t + + Bin mu v ll rr + | d < eps -> do + i <- get + let acc' = PQ.insert i d v acc + put (i + 1) + go acc' ll + | d > mu + eps -> go acc rr + | d <= mu + eps && d > mu - eps -> do + accl <- go acc ll + accr <- go acc rr + union accl accr + | otherwise -> go acc ll + where + d = distf x v + + + + +-- rekey starting from the current index +union :: (MonadState Int m, Ord b) => + PQ.IntPSQ b c -> PQ.IntPSQ b c -> m (PQ.IntPSQ b c) +union q1 q2 = do + i0 <- get + pure $ flip evalState i0 $ foldrM f PQ.empty $ l1 <> l2 + where + f (_, p, v) acc = do + i <- get + let acc' = PQ.insert i p v acc + put $ succ i + pure acc' + l1 = PQ.toList q1 + l2 = PQ.toList q2 + + + +-- rangeVT' :: (Ord a, Num a) => +-- a -> p -> (p -> b -> a) -> VT a b -> [(a, b)] +-- rangeVT' eps x distf = go mempty +-- where +-- insert v qry acc = if d < eps +-- then (d, v) : acc +-- else acc +-- where d = distf qry v +-- go acc = \case +-- Nil -> acc +-- Tip t -> insert t x acc +-- Bin mu v ll rr +-- | d < eps -> go ((d, v) : acc) ll +-- | eps < d - mu -> go acc rr +-- | otherwise -> go acc ll <> go acc rr +-- where +-- d = distf x v + + + + +-- nearest :: (Num d, Ord d) => +-- VPTree d a +-- -> Int +-- -> a +-- -> PQ.IntPSQ d a +-- nearest (VPT t df) k x = nearestVT df k t x + + +{- +variable tau keeps track of closest neighbour yet encounteres + +subtrees are then pruned when the metric information stored in the tree suffices to prove that further consideration is futile, i.e. cannot yield a closer neighbor +-} + +-- -- nearestVT :: (Ord p1, Fractional p1) => +-- -- (p2 -> v -> p1) -> Int -> VT p1 v -> p2 -> SQ.Seq (Int, p1, v) +-- nearestVT :: (Ord p1, Fractional p1) => +-- (p2 -> a -> p1) -> p2 -> VT p1 a -> DQ.DEPQ p1 a +-- nearestVT distf x = z +-- where +-- z = go DQ.empty 0 tau0 +-- tau0 = 1/0 -- initial search radius +-- go acc _ _ Tip = acc +-- go acc i tau (Bin mu v ll rr) +-- | xmu < 0 = go acc i tau' rr -- query point is in outer half-population +-- | d < tau = go acc' (succ i) tau' ll +-- | otherwise = go acc i tau' ll +-- where +-- d = distf x v -- x to vp +-- xmu = mu - d -- x to outer shell +-- acc' = DQ.insert i d v acc +-- tau' = min tau d -- updated search radius + + + +-- nearest1 :: (Ord d, Fractional d) => +-- (a -> a -> d) -> a -> VT d a -> Maybe a +-- nearest1 distf x = go 0 tau0 +-- where +-- tau0 = 1/0 -- initial search radius +-- go _ _ Tip = Nothing +-- go i tau (Bin mu v ll rr) +-- | xmu < 0 = go i tau' rr -- query point is in outer half-population +-- | d < tau = Just v +-- | otherwise = go i tau' ll +-- where +-- d = distf x v -- x to vp +-- xmu = mu - d -- x to outer shell +-- tau' = min tau d -- updated search radius + + +-- nearestIO1 distf x = go tau0 +-- where +-- tau0 = 1/0 -- initial search radius +-- go _ (Tip _) = pure Nothing +-- go tau (Bin mu v ll rr) = do +-- logVar "mu" mu +-- logVar "tau" tau +-- logVar "d" d +-- logVar "xmu" xmu +-- if xmu < 0 +-- then do +-- putStrLn "next : R\n" +-- go tau' rr -- query point is in outer half-population +-- else if d < tau +-- then do +-- logVar "v" v +-- pure $ Just v +-- else do +-- putStrLn "next : L\n" +-- go tau' ll +-- where +-- d = distf x v -- x to vp +-- xmu = mu - d -- x to outer shell +-- tau' = min tau d -- updated search radius + + +-- | Query a 'VPTree' for nearest neighbors +-- +-- NB : the distance function used here should be the same as the one used to construct the tree in the first place + + + + + +-- nearest :: (Fractional d, Ord d) => +-- (a -> a -> d) -- ^ Distance function +-- -- -> Int -- ^ Number of nearest neighbors to return +-- -> a -- ^ Query point +-- -> VPTree d a +-- -> PQ.IntPSQ d a +-- nearest distf x = go PQ.empty 0 (1/0) +-- where +-- go acc _ _ Tip = acc +-- go acc i srad (Bin mu v ll rr) +-- | d < srad' = go acc' (succ i) srad' ll +-- | xmu < 0 = go acc i srad rr +-- | otherwise = go acc i srad ll +-- where +-- acc' = PQ.insert i d v acc +-- d = distf x v -- x to vantage point +-- xmu = mu - d -- x to the outer shell +-- srad' = min srad (abs xmu) -- new search radius + +-- nearestVT :: (Num d, Ord d) => +-- (a -> a -> d) +-- -> Int +-- -> VT d a +-- -> a +-- -> PQ.IntPSQ d a +-- nearestVT distf k tr x = go PQ.empty 0 maxd0 tr +-- where +-- maxd0 = 0 -- initial search radius +-- go acc _ _ Tip = acc +-- go acc i maxd (Bin mu v ll rr) +-- | xmu < 0 = go acc i maxd' rr -- query point is in outer half-population +-- | otherwise = +-- let +-- q1 = xmu > maxd' -- x is farther from the outer shell than farthest point +-- q2 = PQ.size acc == k +-- in if q1 || q2 +-- then acc +-- else go acc' (succ i) maxd' ll +-- where +-- d = distf x v -- x to vp +-- xmu = mu - d -- x to outer shell +-- acc' = PQ.insert i d v acc +-- maxd' = max maxd d -- next search radius + +-- logVar :: (MonadIO io, Show a) => String -> a -> io () +-- logVar w x = liftIO $ putStrLn $ unwords [w, "=", show x] + +{- +At any given step we are working with a node of the tree that has a + +vantage point v +threshold distance mu. + +The query point x will be some distance d from v. + +If d is less than mu then use the algorithm recursively to search the subtree of the node that contains the points closer to v than mu; otherwise recurse to the subtree of the node that contains the points that are farther than the vantage point than mu. + +If the recursive use of the algorithm finds a neighboring point n with distance to x that is less than |mu − d| then it cannot help to search the other subtree of this node; the discovered node n is returned. Otherwise, the other subtree also needs to be searched recursively. +-} + +-- nnnn distf k tr x = z +-- where +-- (z, _, _) = go PQ.empty 0 maxd0 tr +-- maxd0 = 0 +-- go acc i maxd Nil = (acc, i, maxd) +-- go acc i maxd (Bin mu v ll rr) +-- | q1 || q2 = go acc' (succ i) maxd' ll -- x closer to v than to shell +-- | d < mu = -- x inside shell but not closer to v +-- let +-- (accl, il, maxdl) = go acc i maxd' ll +-- in go accl il maxdl rr +-- | otherwise = go acc i maxd' rr -- x outside shell +-- where +-- d = distf x v +-- xmu = mu - d +-- acc' = PQ.insert i d v acc +-- maxd' = max maxd (abs xmu) -- next search radius +-- q1 = d < xmu +-- q2 = PQ.size acc == k + + + +-- nearest distf x = go PQ.empty 0 (1/0) +-- where +-- go acc _ _ Tip = acc +-- go acc i srad (Bin mu v ll rr) +-- | xmu < 0 = go acc i srad rr -- query point is outside the radius mu + +-- -- | xv < xmu = go acc i srad ll +-- -- | otherwise = let +-- -- acc' = PQ.insert i xv v acc +-- -- srad' = min mu srad -- new search radius +-- -- in go acc' (i + 1) srad' ll -- FIXME double check this + +-- where +-- xv = distf x v -- x to vantage point +-- xmu = mu - xv -- x to the outer shell diff --git a/src/Data/VPTree/TestData.hs b/src/Data/VPTree/TestData.hs new file mode 100644 --- /dev/null +++ b/src/Data/VPTree/TestData.hs @@ -0,0 +1,90 @@ +{-# language DeriveGeneric #-} +{-# options_ghc -Wno-unused-imports #-} +module Data.VPTree.TestData where + +-- import Data.Foldable (toList) +import GHC.Generics (Generic(..)) +import Text.Printf (printf) + +-- deepseq +import Control.DeepSeq (NFData()) +-- mwc-probability +import qualified System.Random.MWC.Probability as P (Gen, Prob, withSystemRandom, asGenIO, GenIO, create, initialize, sample, samples, normal, bernoulli, uniformR) +-- primitive +import Control.Monad.Primitive (PrimMonad(..), PrimState) +-- vector +import qualified Data.Vector as V (Vector, map, filter, length, toList, replicate, partition, zipWith, head, tail, fromList, thaw, freeze, (!), foldl) + + +import Data.VPTree.Build (build) +-- import Data.VPTree.Draw (draw) +import Data.VPTree.Internal (VT, VPTree, withST, withST_, withIO) +-- import Data.VPTree.Query (range, distances) + + +-- test data + +data P = P !Double !Double deriving (Eq, Generic) +instance NFData P +instance Show P where + show (P x y) = printf "(%2.2f, %2.2f)" x y --show (x,y) + +(.+.) :: P -> P -> P +P x1 y1 .+. P x2 y2 = P (x1 + x2) (y1 + y2) + +distp :: P -> P -> Double +distp (P x1 y1) (P x2 y2) = sqrt $ (x1 - x2)**2 + (y1 - y2)**2 + + + +-- t2, t2', t3 :: VPTree Double P +-- t3 = buildP $ genN3 12 +-- t2 = buildP $ genN2 12 +-- t2' = buildP $ genN2 10000 + +genN1, gaussMixSamples, binDiskSamples :: Int -> V.Vector P +gaussMixSamples n = V.fromList $ withST_ (P.samples n (gaussMix 0 25 1 1)) + +genN1 n = V.fromList $ withST_ (P.samples n (isoNormal2d 0 1)) + +binDiskSamples n = V.fromList $ withST_ $ P.samples n (binDisk 1 1 z fv) + where + z = P 0 0 + fv = P 5 5 + + +-- | binary mixture of isotropic 2d normal distribs +gaussMix :: PrimMonad m => + Double -> Double -> Double -> Double -> P.Prob m P +gaussMix mu1 mu2 sig1 sig2 = do + b <- coin + if b + then isoNormal2d mu1 sig1 + else isoNormal2d mu2 sig2 + +coin :: PrimMonad m => P.Prob m Bool +coin = P.bernoulli 0.5 + +isoNormal2d :: PrimMonad m => Double -> Double -> P.Prob m P +isoNormal2d mu sig = P <$> P.normal mu sig <*> P.normal mu sig + +binDisk :: PrimMonad m => Double -> Double -> P -> P -> P.Prob m P +binDisk r0 r1 p0 p1 = do + b <- coin + if b + then uniformDisk r0 p0 + else uniformDisk r1 p1 + +-- point in a disk of radius r and centered at P +uniformDisk :: PrimMonad m => Double -> P -> P.Prob m P +uniformDisk rmax p = do + r <- P.uniformR (0, rmax) + aa <- P.uniformR (0, 2 * pi) + let + x = r * cos aa + y = r * sin aa + p0 = P x y + pure $ p0 .+. p + +buildP :: V.Vector P -> VPTree Double P +buildP = build distp (1.0 :: Double) diff --git a/test/Spec.hs b/test/Spec.hs new file mode 100644 --- /dev/null +++ b/test/Spec.hs @@ -0,0 +1,1 @@ +{-# OPTIONS_GHC -F -pgmF hspec-discover #-} diff --git a/vp-tree.cabal b/vp-tree.cabal new file mode 100644 --- /dev/null +++ b/vp-tree.cabal @@ -0,0 +1,90 @@ +name: vp-tree +version: 0.1.0.0 +synopsis: Vantage Point Trees +description: Vantage Point Trees enable fast nearest-neighbor queries in metric spaces +homepage: https://github.com/ocramz/vp-tree +license: BSD3 +license-file: LICENSE +author: Marco Zocca +maintainer: ocramz +copyright: 2020-2021 Marco Zocca +category: Data, Data Mining, Data Structures, Machine Learning +build-type: Simple +extra-source-files: README.md +cabal-version: >=1.10 +tested-with: GHC == 8.6.5 + +library + default-language: Haskell2010 + ghc-options: -Wall + hs-source-dirs: src + exposed-modules: Data.VPTree + Data.VPTree.Build + Data.VPTree.Query + Data.VPTree.Internal + Data.VPTree.Draw + Data.VPTree.TestData + -- other-modules: + build-depends: base >= 4.7 && < 5 + , boxes >= 0.1.5 + -- , conduit + , containers >= 0.6.0.1 + , deepseq >= 1.4.4.0 + , depq >= 0.3 + -- , exceptions + , mtl >= 2.2.2 + , mwc-probability >= 2.1.0 + , primitive >= 0.6.4.0 + , psqueues >= 0.2.7.2 + , sampling >= 0.3.3 + , serialise >= 0.2.2.0 + , transformers >= 0.5.6.2 + , vector >= 0.12.1.2 + , vector-algorithms >= 0.8.0.3 + -- DEBUG + -- , hspec + -- , weigh + +-- executable vp-tree +-- default-language: Haskell2010 +-- ghc-options: -threaded -rtsopts -with-rtsopts=-N +-- hs-source-dirs: app +-- main-is: Main.hs +-- build-depends: base +-- , vp-tree + +test-suite spec + default-language: Haskell2010 + ghc-options: -Wall + type: exitcode-stdio-1.0 + hs-source-dirs: test + main-is: Spec.hs + build-depends: base + , vp-tree + , hspec + , mwc-probability + , primitive + , QuickCheck + , vector + +benchmark bench-memory + type: exitcode-stdio-1.0 + hs-source-dirs: bench/memory + main-is: Main.hs + build-depends: base + , vp-tree + , bytestring + , conduit + , containers + , deepseq + , vector + , weigh + ghc-options: -Wall + -rtsopts + -with-rtsopts=-T + default-language: Haskell2010 + + +source-repository head + type: git + location: https://github.com/ocramz/vp-tree