rp-tree 0.5 → 0.6
raw patch · 4 files changed
+118/−88 lines, 4 filesdep +heapsdep −microlensdep −microlens-thdep −psqueuesdep ~hspec
Dependencies added: heaps
Dependencies removed: microlens, microlens-th, psqueues
Dependency ranges changed: hspec
Files
- Changelog.md +8/−0
- rp-tree.cabal +4/−7
- src/Data/RPTree.hs +100/−67
- src/Data/RPTree/Internal.hs +6/−14
Changelog.md view
@@ -1,3 +1,11 @@+0.6++- add 'knnPQ' based on 'heaps', to be confirmed how efficient it is++0.5.1++- no need for 'microlens' and 'hspec' in the library+ 0.5 - fixed intermittent bug in 'knn', originally due to messing up the order of 'take' and 'sort'.
rp-tree.cabal view
@@ -1,5 +1,5 @@ name: rp-tree-version: 0.5+version: 0.6 synopsis: Random projection trees description: Random projection trees for approximate nearest neighbor search in high-dimensional vector spaces .@@ -35,11 +35,9 @@ , conduit >= 1.3.1 , containers >= 0.6 , deepseq >= 1.4- - , microlens- , microlens-th+ , heaps , mtl >= 2.2.2- , psqueues+ -- , psqueues , serialise >= 0.2 , splitmix , splitmix-distributions >= 0.9@@ -49,8 +47,7 @@ , vector-algorithms >= 0.8 -- -- -- DEBUG -- , benchpress- -- , hspec- , hspec >= 2.7.1+ -- , hspec >= 2.7.1 -- , mnist-idx-conduit test-suite spec
src/Data/RPTree.hs view
@@ -56,7 +56,7 @@ -- , ForestParams -- * Query , knn- -- , knnPQ+ , knnPQ -- * I/O , serialiseRPForest , deserialiseRPForest@@ -111,6 +111,7 @@ import Data.Foldable (Foldable(..), maximumBy, minimumBy) import Data.Functor.Identity (Identity(..)) import Data.List (partition, sortBy)+import Data.Maybe (maybeToList) import Data.Monoid (Sum(..)) import Data.Ord (comparing) import Data.Semigroup (Min(..))@@ -123,14 +124,16 @@ import qualified Data.Set as S (Set, fromList, intersection, insert) -- deepseq import Control.DeepSeq (NFData(..))--- psqueues-import qualified Data.IntPSQ as PQ (IntPSQ, findMin, minView, empty, insert, fromList, toList)+-- heaps+import qualified Data.Heap as H (Heap, fromList, insert, Entry(..), empty, group, viewMin, map)+-- -- psqueues+-- import qualified Data.IntPSQ as PQ (IntPSQ, findMin, minView, empty, insert, fromList, toList) -- transformers import Control.Monad.Trans.State (StateT(..), runStateT, evalStateT, State, runState, evalState, get, put) import Control.Monad.Trans.Class (MonadTrans(..)) -- vector import qualified Data.Vector as V (Vector, replicateM, fromList)-import qualified Data.Vector.Generic as VG (Vector(..), unfoldrM, length, replicateM, (!), map, freeze, thaw, take, drop, unzip)+import qualified Data.Vector.Generic as VG (Vector(..), unfoldrM, length, replicateM, (!), map, freeze, thaw, take, drop, unzip, foldl, fromList) import qualified Data.Vector.Unboxed as VU (Vector, Unbox, fromList) import qualified Data.Vector.Storable as VS (Vector) -- vector-algorithms@@ -162,23 +165,50 @@ where cs = VG.map (\xe -> (eEmbed xe `distf` q, xe)) $ fold $ (`candidates` q) <$> tts --- | Same as 'knn' but accumulating the result in low margin order (following the intuition in 'annoy').+-- | Same as 'knn' but based on 'H.Heap' ----- FIXME to be verified-knnPQ :: (Ord p, Inner SVector v, VU.Unbox d, RealFrac d) =>+-- FIXME uses 'nub' internally so might be slow+knnPQ :: (Ord p, Inner SVector v, VU.Unbox d, Real d) => (u d -> v d -> p) -- ^ distance function -> Int -- ^ k neighbors -> RPForest d (V.Vector (Embed u d x)) -- ^ random projection forest -> v d -- ^ query point- -> V.Vector (p, Embed u d x)-knnPQ distf k tts q = sortByVG fst cs+ -> V.Vector (p, Embed u d x) -- ^ ordered in increasing distance order to the query point+knnPQ distf k tts q = VG.fromList $ take k $ map (\(H.Entry p x) -> (p , x)) $ nub h where- cs = VG.map (\xe -> (eEmbed xe `distf` q, xe)) $ fold cstt- cstt = (takeFromPQ nsing) . (`candidatesPQ` q) <$> tts- nsing = (k `div` n) `max` 1- n = length tts+ xs = fold $ (`candidates` q) <$> tts+ h = VG.foldl insf H.empty xs+ insf acc x = H.insert (H.Entry p x) acc+ where+ p = eEmbed x `distf` q +-- | deduplicate+nub :: (Ord a) => H.Heap a -> [a]+nub = foldMap maybeToList . H.map view . H.group+ where+ view h = fst <$> H.viewMin h +type PQueue p a = H.Heap (H.Entry p a)++-- -- | Same as 'knn' but accumulating the result in low margin order (following the intuition in 'annoy').+-- --+-- -- FIXME to be verified+-- knnPQ :: (Ord p, Inner SVector v, VU.Unbox d, RealFrac d) =>+-- (u d -> v d -> p) -- ^ distance function+-- -> Int -- ^ k neighbors+-- -> RPForest d (V.Vector (Embed u d x)) -- ^ random projection forest+-- -> v d -- ^ query point+-- -> V.Vector (p, Embed u d x)+-- knnPQ distf k tts q = sortByVG fst cs+-- where+-- cs = VG.map (\xe -> (eEmbed xe `distf` q, xe)) $ fold cstt+-- cstt = (takeFromPQ nsing) . (`candidatesPQ` q) <$> tts+-- nsing = (k `div` n) `max` 1+-- n = length tts++++ -- | Average recall-at-k, computed over a set of trees -- -- The supplied distance function @d@ must satisfy the definition of a metric, i.e.@@ -246,63 +276,66 @@ | otherwise -> go i' rtree --- | like 'candidates' but outputs an ordered 'IntPQ' where the margin to the median projection is interpreted as queue priority-candidatesPQ :: (Fractional d, Ord d, Inner SVector v, VU.Unbox d) =>- RPTree d l xs- -> v d -- ^ query point- -> PQ.IntPSQ d xs-candidatesPQ (RPTree rvs tt) x = evalS $ go 0 tt PQ.empty (1/0)- where- go _ (Tip _ xs) acc dprev =- insPQ dprev xs acc- go ixLev (Bin _ thr margin ltree rtree) acc dprev = do- let- (mglo, mghi) = getMargin margin- r = rvs VG.! ixLev- proj = r `inner` x- i' = succ ixLev- dl = abs (mglo - proj) -- left margin- dr = abs (mghi - proj) -- right margin- if | proj < thr &&- dl > dr -> do- ll <- go i' ltree acc (min dprev dl)- lr <- go i' rtree acc (min dprev dr)- pure $ PQ.fromList (PQ.toList ll <> PQ.toList lr)- | proj < thr -> go i' ltree acc (min dprev dl)- | proj > thr &&- dl < dr -> do- ll <- go i' ltree acc (min dprev dl)- lr <- go i' rtree acc (min dprev dr)- pure $ PQ.fromList (PQ.toList ll <> PQ.toList lr)- | otherwise -> go i' rtree acc (min dprev dr) -takeFromPQ :: (Ord p, Foldable t, Monoid (t a)) =>- Int -- ^ number of elements to keep- -> PQ.IntPSQ p (t a)- -> t a-takeFromPQ n pq = foldMap snd $ reverse $ go [] 0 pq- where- go acc nacc q = case PQ.minView q of- Nothing -> acc- Just (_, p, xs, pqRest) ->- let- nxs = length xs- nacc' = nacc + nxs- in if nacc' < n- then go ((p, xs) : acc) nacc' pqRest- else acc -type S = State Int-evalS :: S a -> a-evalS = flip evalState 0 -insPQ :: (Ord p) => p -> v -> PQ.IntPSQ p v -> S (PQ.IntPSQ p v)-insPQ p x pq = do- i <- get- let- pq' = PQ.insert i p x pq- put (succ i)- pure pq'+-- -- | like 'candidates' but outputs an ordered 'IntPQ' where the margin to the median projection is interpreted as queue priority+-- candidatesPQ :: (Fractional d, Ord d, Inner SVector v, VU.Unbox d) =>+-- RPTree d l xs+-- -> v d -- ^ query point+-- -> PQ.IntPSQ d xs+-- candidatesPQ (RPTree rvs tt) x = evalS $ go 0 tt PQ.empty (1/0)+-- where+-- go _ (Tip _ xs) acc dprev =+-- insPQ dprev xs acc+-- go ixLev (Bin _ thr margin ltree rtree) acc dprev = do+-- let+-- (mglo, mghi) = getMargin margin+-- r = rvs VG.! ixLev+-- proj = r `inner` x+-- i' = succ ixLev+-- dl = abs (mglo - proj) -- left margin+-- dr = abs (mghi - proj) -- right margin+-- if | proj < thr &&+-- dl > dr -> do+-- ll <- go i' ltree acc (min dprev dl)+-- lr <- go i' rtree acc (min dprev dr)+-- pure $ PQ.fromList (PQ.toList ll <> PQ.toList lr)+-- | proj < thr -> go i' ltree acc (min dprev dl)+-- | proj > thr &&+-- dl < dr -> do+-- ll <- go i' ltree acc (min dprev dl)+-- lr <- go i' rtree acc (min dprev dr)+-- pure $ PQ.fromList (PQ.toList ll <> PQ.toList lr)+-- | otherwise -> go i' rtree acc (min dprev dr)++-- takeFromPQ :: (Ord p, Foldable t, Monoid (t a)) =>+-- Int -- ^ number of elements to keep+-- -> PQ.IntPSQ p (t a)+-- -> t a+-- takeFromPQ n pq = foldMap snd $ reverse $ go [] 0 pq+-- where+-- go acc nacc q = case PQ.minView q of+-- Nothing -> acc+-- Just (_, p, xs, pqRest) ->+-- let+-- nxs = length xs+-- nacc' = nacc + nxs+-- in if nacc' < n+-- then go ((p, xs) : acc) nacc' pqRest+-- else acc++-- type S = State Int+-- evalS :: S a -> a+-- evalS = flip evalState 0++-- insPQ :: (Ord p) => p -> v -> PQ.IntPSQ p v -> S (PQ.IntPSQ p v)+-- insPQ p x pq = do+-- i <- get+-- let+-- pq' = PQ.insert i p x pq+-- put (succ i)+-- pure pq'
src/Data/RPTree/Internal.hs view
@@ -2,7 +2,6 @@ {-# language DeriveAnyClass #-} {-# LANGUAGE DeriveTraversable #-} {-# LANGUAGE DeriveFoldable #-}--- {-# LANGUAGE UndecidableInstances #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE DeriveFunctor #-}@@ -10,7 +9,6 @@ {-# language LambdaCase #-} {-# language MultiParamTypeClasses #-} {-# language GeneralizedNewtypeDeriving #-}-{-# language TemplateHaskell #-} {-# LANGUAGE BangPatterns #-} {-# options_ghc -Wno-unused-imports #-} module Data.RPTree.Internal where@@ -38,10 +36,10 @@ import qualified Data.IntMap.Strict as IM (IntMap, fromList) -- deepseq import Control.DeepSeq (NFData(..))--- microlens-import Lens.Micro ((^..), Traversal', folded, Getting)--- microlens-th-import Lens.Micro.TH (makeLensesFor)+-- -- microlens+-- import Lens.Micro ((^..), Traversal', folded, Getting)+-- -- microlens-th+-- import Lens.Micro.TH (makeLensesFor) -- serialise import Codec.Serialise (Serialise(..), serialise, deserialiseOrFail) -- vector@@ -162,7 +160,6 @@ Bin x thr mg tl tr -> Bin <$> f x <*> pure thr <*> pure mg <*> bitraverse f g tl <*> bitraverse f g tr Tip x y -> Tip <$> f x <*> g y instance (Serialise a, Serialise l, Serialise d) => Serialise (RPT d l a)-makeLensesFor [("_rpData", "rpData"), ("_rpLabel", "rpLabel")] ''RPT instance (NFData v, NFData l, NFData a) => NFData (RPT v l a) -- | Random projection trees@@ -177,7 +174,6 @@ , _rpTree :: !(RPT d l a) } deriving (Eq, Show, Functor, Foldable, Traversable, Generic) instance (Serialise d, Serialise l, Serialise a, VU.Unbox d) => Serialise (RPTree d l a)-makeLensesFor [("_rpTree", "rpTree")] ''RPTree instance (NFData a, NFData l, NFData d) => NFData (RPTree d l a) -- | A random projection forest is an ordered set of 'RPTree's@@ -199,12 +195,9 @@ Left e -> Left (show e) Right xs -> Right $ IM.fromList $ zip [0 ..] xs -rpTreeData :: Traversal' (RPTree d l a) a-rpTreeData = rpTree . rpData- -- | All data buckets stored at the leaves of the tree leaves :: RPTree d l a -> [a]-leaves = (^.. rpTreeData)+leaves = toList -- | Number of tree levels levels :: RPTree d l a -> Int@@ -214,8 +207,6 @@ points :: Monoid m => RPTree d l m -> m points (RPTree _ t) = fold t --- -- points in 2d--- data P a = P !a !a deriving (Eq, Show) -- | Scale a vector class Scale v where@@ -481,3 +472,4 @@ -- (il, xl) = vv1 VG.! i -- xr = vv2 VG.! il -- y = f xl xr+