rp-tree 0.2.1.0 → 0.3
raw patch · 7 files changed
+98/−137 lines, 7 filesdep ~conduit
Dependency ranges changed: conduit
Files
- README.md +1/−1
- app/Main.hs +2/−25
- bench/time/Main.hs +54/−68
- rp-tree.cabal +2/−1
- src/Data/RPTree.hs +16/−24
- src/Data/RPTree/Conduit.hs +20/−15
- src/Data/RPTree/Internal.hs +3/−3
README.md view
@@ -1,5 +1,5 @@ # rp-tree -+ Random projection trees for approximate nearest neighbor search in high-dimensional vector spaces
app/Main.hs view
@@ -27,7 +27,7 @@ import qualified Data.Vector as V (Vector, toList, fromList, replicate, zip) import Control.Monad (replicateM)-import Data.RPTree (knn, candidates, Inner(..), RPTree, RPForest, leaves, SVector, fromListSv, DVector, fromListDv, dense, writeCsv, tree, forest, dataSource, sparse, normal2, normalSparse2)+import Data.RPTree (knn, candidates, Embed(..), Inner(..), RPTree, RPForest, leaves, SVector, fromListSv, DVector, fromListDv, dense, writeCsv, tree, forest, dataSource, sparse, normal2, normalSparse2) import Data.RPTree.Internal.Testing (datS, datD) main :: IO ()@@ -40,7 +40,7 @@ d = 100 pnz = 0.3 chunk = 20- src = datS n d pnz+ src = datS n d pnz .| C.map (\ x -> Embed x ()) seed = 1234 (q, tts) <- sampleT seed $ do tts <- C.runConduit $@@ -110,31 +110,8 @@ dataset n = V.fromList $ sample 1234 $ replicateM n (dense 2 $ normal 0 1) --- treeC0 :: MonadThrow m =>--- Int -> GenT m (RPTree Double (V.Vector (DVector Double)))--- treeC0 n = treeSink 1234 10 20 100 1.0 2 (srcC n) -{--λ> nn0 10000 (fromListDv [0,0])-[(0.13092191004810114,DV [-8.771274989760332e-2,9.71957819868927e-2]),(0.14722273682679538,DV [-4.767722969780902e-2,0.13928896584839093]),(0.1626065099556818,DV [-4.57842765697381e-2,0.15602780873598454]),(0.22082909577433263,DV [-3.62336905451185e-2,0.21783619811681887]),(0.22085935710897311,DV [0.21196201255823421,-6.2056110535964756e-2]),(0.2636139991233282,DV [-0.24290511334764195,0.10241799862994452]),(0.3869415454995779,DV [-0.3658837577279577,0.12590804368455188]),(0.3951528583078011,DV [-0.3543713488257354,0.1748334308999686]),(0.6174219338196472,DV [-0.4952807707701239,0.3686553979897009]),(0.6968774335522048,DV [-0.6408548616154526,0.2737575638007956])]--}-nn0 :: (Inner SVector v, Inner DVector v) =>- Int -> v Double -> V.Vector (Double, DVector Double)-nn0 n q = case ttsm of- Just tts -> knn metricL2 10 tts q -- FIXME voting search size ?!- -- Nothing -> mempty- where- ttsm = sampleT 1234 $ forestC0 n -cs0 n q = case sampleT 1234 $ forestC0 n of- Just tts -> (`candidates` q) <$> tts--forestC0 :: MonadThrow m =>- Int- -> GenT- m- (IM.IntMap (RPTree Double (V.Vector (DVector Double))))-forestC0 n = forest 1234 10 20 10 100 1.0 2 (srcC n) srcC :: Monad m => Int -> C.ConduitT i (DVector Double) (GenT m) () srcC n = dataSource n normal2
bench/time/Main.hs view
@@ -31,7 +31,7 @@ import qualified Data.Vector as V (Vector, replicateM, fromList) import qualified Data.Vector.Unboxed as VU (Unbox, Vector, map) -import Data.RPTree (tree, forest, recallWith, knn, fromVectorSv, fromListSv, RPForest, RPTree, SVector, Inner(..), normalSparse2, liftC)+import Data.RPTree (tree, forest, recallWith, knn, fromVectorSv, fromListSv, RPForest, RPTree, SVector, Inner(..), normalSparse2, liftC, Embed(..)) import Data.RPTree.Internal.Testing (BenchConfig(..), randSeed, datD, datS) main :: IO ()@@ -49,7 +49,7 @@ chunk <- [100], nzd <- [0.2], n <- [1000],- nq <- [10, 100]+ nq <- [10] ] -- -- Binary mixture@@ -68,7 +68,7 @@ -- | Measure recall @ 10 and mean time binMixFQBench1 :: BenchConfig -> IO (Double, Double)-binMixFQBench1 cfg = forestBenchGen seed (datS n d nzData) act 2 cfg+binMixFQBench1 cfg = forestBenchGen seed src act 2 cfg where n = bcDataSize cfg d = bcVectorDim cfg@@ -77,6 +77,7 @@ nzData = 0.8 -- nz density of data k = 10 -- number of ANN's to return seed = 1234+ src = datS n d nzData .| C.map (\r -> Embed r ()) qs = samples nq seed $ normalSparse2 nzData d act tt = do -- pure $! recallWith metricL2 tt k `map` qs@@ -89,37 +90,38 @@ -- -- MNIST dataset -mnistBench :: IO ()-mnistBench = do- let- cfgs = benchConfigs "MNIST dataset" 784- mnfpath = "assets/mnist/train-images-idx3-ubyte"- forM_ cfgs $ \cfg -> do- stats <- mnistFQBench1 mnfpath cfg- print cfg- print stats+-- mnistBench :: IO ()+-- mnistBench = do+-- let+-- cfgs = benchConfigs "MNIST dataset" 784+-- mnfpath = "assets/mnist/train-images-idx3-ubyte"+-- forM_ cfgs $ \cfg -> do+-- stats <- mnistFQBench1 mnfpath cfg+-- print cfg+-- print stats --- | Measure recall @ 10 and mean time-mnistFQBench1 :: FilePath -> BenchConfig -> IO (Double, Double)-mnistFQBench1 fp cfg = forestBench (mnist fp n) act 1 cfg- where- n = bcDataSize cfg- k = 10 -- number of ANN's to return- d = bcVectorDim cfg- nzData = 0.8 -- nz density of data - act x = do- tt <- runResourceT x- let q = sample 1234 $ normalSparse2 nzData d- pure $! recallWith metricL2 tt k q+-- -- | Measure recall @ 10 and mean time+-- -- mnistFQBench1 :: FilePath -> BenchConfig -> IO (Double, Double)+-- mnistFQBench1 fp cfg = forestBench (mnist fp n) act 1 cfg+-- where+-- n = bcDataSize cfg+-- k = 10 -- number of ANN's to return+-- d = bcVectorDim cfg+-- nzData = 0.8 -- nz density of data +-- act x = do+-- tt <- runResourceT x+-- let q = sample 1234 $ normalSparse2 nzData d+-- pure $! recallWith metricL2 tt k q mnist :: MonadResource m => FilePath -- path to uncompressed MNIST IDX data file -> Int -- number of data items- -> C.ConduitT a (SVector Double) m ()+ -> C.ConduitT a (Embed SVector Double ()) m () mnist fp n = C.takeExactly n src where src = sourceIdxSparse fp .|- C.map (\r -> fromVectorSv (sBufSize r) (VU.map f $ sNzComponents r))+ C.map (\r -> fromVectorSv (sBufSize r) (VU.map f $ sNzComponents r)) .|+ C.map (\r -> Embed r ()) f (i, x) = (i, toUnitRange x) toUnitRange :: Word8 -> Double@@ -130,13 +132,13 @@ -- -- UTILS --- | runs a benchmark on a newly created RPForest initialized with a random seed-forestBench :: (MonadThrow m, Inner SVector v) =>- C.ConduitT () (v Double) m ()- -> (m (RPForest Double (V.Vector (v Double))) -> IO c) -- ^ allows for both deterministic and random data sources- -> Int -- ^ number of replicates- -> BenchConfig- -> IO (c, Double) -- ^ result, mean wall-clock time measurement+-- -- | runs a benchmark on a newly created RPForest initialized with a random seed+-- forestBench :: (MonadThrow m, Inner SVector v) =>+-- C.ConduitT () (v Double) m ()+-- -> (m (RPForest Double (V.Vector (v Double))) -> IO c) -- ^ allows for both deterministic and random data sources+-- -> Int -- ^ number of replicates+-- -> BenchConfig+-- -> IO (c, Double) -- ^ result, mean wall-clock time measurement forestBench src go n cfg = benchmark n setup (const $ pure ()) go where setup = do@@ -145,26 +147,10 @@ pure $ growForest s cfg src --- -- forestBenchRes :: (MonadResource m, Inner SVector v) =>--- -- C.ConduitT--- -- ()--- -- (v Double)--- -- m--- -- ()--- -- -> (RPForest Double (V.Vector (v Double)) -> IO c)--- -- -> Int--- -- -> BenchConfig--- -- -> IO (c, Double)-forestBenchRes src go n cfg = benchmarkM n setup go- where- setup = do- s <- randSeed- runResourceT $ growForest s cfg src--forestBenchGen :: (MonadIO m, Inner SVector v, NFData (v Double)) =>+forestBenchGen :: (MonadIO m, Inner SVector v, NFData x, NFData (v Double)) => Word64- -> C.ConduitT () (v Double) (GenT m) ()- -> (RPForest Double (V.Vector (v Double)) -> m a2)+ -> C.ConduitT () (Embed v Double x) (GenT m) ()+ -> (RPForest Double (V.Vector (Embed v Double x)) -> m a2) -> Int -> BenchConfig -> m (a2, Double)@@ -177,12 +163,12 @@ -treeBench :: (Monad m, Inner SVector v) =>- C.ConduitT () (v Double) m ()- -> (m (RPTree Double (V.Vector (v Double))) -> IO c)- -> Int- -> BenchConfig- -> IO (c, Double)+-- treeBench :: (Monad m, Inner SVector v) =>+-- C.ConduitT () (v Double) m ()+-- -> (m (RPTree Double (V.Vector (v Double))) -> IO c)+-- -> Int+-- -> BenchConfig+-- -> IO (c, Double) treeBench src go n cfg = benchmark n setup (const $ pure ()) go where setup = do@@ -190,19 +176,19 @@ -- let src' = C.transPipe lift src pure $ growTree s cfg src -growTree :: (Monad m, Inner SVector v) =>- Word64- -> BenchConfig- -> C.ConduitT () (v Double) m ()- -> m (RPTree Double (V.Vector (v Double)))+-- growTree :: (Monad m, Inner SVector v) =>+-- Word64+-- -> BenchConfig+-- -> C.ConduitT () (v Double) m ()+-- -> m (RPTree Double (V.Vector (v Double))) growTree seed (BenchConfig _ maxd minl _ chunksize pnz pdim _ _) = tree seed maxd minl chunksize pnz pdim -growForest :: (Monad m, Inner SVector v) =>- Word64- -> BenchConfig- -> C.ConduitT () (v Double) m ()- -> m (RPForest Double (V.Vector (v Double)))+-- growForest :: (Monad m, Inner SVector v) =>+-- Word64+-- -> BenchConfig+-- -> C.ConduitT () (v Double) m ()+-- -> m (RPForest Double (V.Vector (v Double))) growForest seed (BenchConfig _ maxd minl ntrees chunksize pnz pdim _ _) = forest seed maxd minl ntrees chunksize pnz pdim
rp-tree.cabal view
@@ -1,5 +1,5 @@ name: rp-tree-version: 0.2.1.0+version: 0.3 synopsis: Random projection trees description: Random projection trees for approximate nearest neighbor search in high-dimensional vector spaces homepage: https://github.com/ocramz/rp-tree@@ -57,6 +57,7 @@ main-is: Spec.hs build-depends: base , rp-tree+ , conduit , hspec , QuickCheck , splitmix-distributions
src/Data/RPTree.hs view
@@ -26,6 +26,7 @@ -- * Access , levels, points, leaves, candidates -- * Types+ , Embed(..) -- ** RPTree , RPTree, RPForest -- *@@ -90,54 +91,45 @@ import Data.RPTree.Conduit (tree, forest, dataSource, liftC) import Data.RPTree.Gen (sparse, dense, normal2, normalSparse2)-import Data.RPTree.Internal (RPTree(..), RPForest, RPT(..), levels, points, leaves, RT(..), Inner(..), Scale(..), scaleS, scaleD, (/.), innerDD, innerSD, innerSS, metricSSL2, metricSDL2, SVector(..), fromListSv, fromVectorSv, DVector(..), fromListDv, fromVectorDv, partitionAtMedian, Margin, getMargin, sortByVG, serialiseRPForest, deserialiseRPForest)+import Data.RPTree.Internal (RPTree(..), RPForest, RPT(..), Embed(..), levels, points, leaves, RT(..), Inner(..), Scale(..), scaleS, scaleD, (/.), innerDD, innerSD, innerSS, metricSSL2, metricSDL2, SVector(..), fromListSv, fromVectorSv, DVector(..), fromListDv, fromVectorDv, partitionAtMedian, Margin, getMargin, sortByVG, serialiseRPForest, deserialiseRPForest) import Data.RPTree.Internal.Testing (BenchConfig(..), randSeed) import Data.RPTree.Draw (draw, writeCsv) -- | k nearest neighbors knn :: (Ord p, Inner SVector v, VU.Unbox d, Real d) =>- (v2 -> v d -> p) -- ^ distance function+ (u d -> v d -> p) -- ^ distance function -> Int -- ^ k neighbors- -> RPForest d (V.Vector v2) -- ^ random projection forest+ -> RPForest d (V.Vector (Embed u d x)) -- ^ random projection forest -> v d -- ^ query point- -> V.Vector (p, v2) -- ^ ordered in increasing distance order+ -> V.Vector (p, Embed u d x) -- ^ ordered in increasing distance order knn distf k tts q = sortByVG fst cs where- cs = VG.map (\x -> (x `distf` q, x)) $ VG.take k $ fold $ (`candidates` q) <$> tts+ cs = VG.map (\xe -> (eEmbed xe `distf` q, xe)) $ VG.take k $ fold $ (`candidates` q) <$> tts -- | average recall-at-k, computed over a set of trees-recallWith :: (Inner SVector v, VU.Unbox a, Fractional a, Ord a, Ord d, Ord p) =>- (p -> v a -> d)- -> RPForest a (V.Vector p)+recallWith :: (Inner SVector v, VU.Unbox d, Fractional a1, Ord d, Ord a2, Ord x, Ord (u d), Num d) =>+ (u d -> v d -> a2)+ -> RPForest d (V.Vector (Embed u d x)) -> Int -- ^ k : number of nearest neighbors to consider- -> v a -- ^ query point- -> a+ -> v d -- ^ query point+ -> a1 recallWith distf tt k q = sum rs / fromIntegral n where rs = fmap (\t -> recallWith1 distf t k q) tt n = length tt --- -- | Recall computed wrt the Euclidean distance--- recallEuclid :: (Inner SVector v, Inner u v, VU.Unbox a, Ord a, Ord (u a), Floating a) =>--- RPTree a (V.Vector (u a))--- -> Int -- ^ k : number of nearest neighbors to consider--- -> v a -- ^ query point--- -> Double--- recallEuclid = recallWith metricL2--recallWith1 :: (Fractional a1, Inner SVector v, Ord d, VU.Unbox d,- Num d, Ord a3, Ord a2) =>- (a2 -> v d -> a3) -- ^ distance function- -> RPTree d (V.Vector a2)+recallWith1 :: (Inner SVector v, Ord d, VU.Unbox d, Fractional p, Ord a, Ord x, Ord (u d), Num d) =>+ (u d -> v d -> a) -- ^ distance function+ -> RPTree d (V.Vector (Embed u d x)) -> Int -- ^ k : number of nearest neighbors to consider -> v d -- ^ query point- -> a1+ -> p recallWith1 distf tt k q = fromIntegral (length aintk) / fromIntegral k where xs = points tt- dists = sortBy (comparing snd) $ toList $ fmap (\x -> (x, x `distf` q)) xs+ dists = sortBy (comparing snd) $ toList $ fmap (\x -> (x, eEmbed x `distf` q)) xs kk = S.fromList $ map fst $ take k dists -- first k nn's aa = set $ candidates tt q aintk = aa `S.intersection` kk
src/Data/RPTree/Conduit.hs view
@@ -4,7 +4,8 @@ {-# options_ghc -Wno-unused-imports #-} {-# options_ghc -Wno-unused-top-binds #-} module Data.RPTree.Conduit- (tree,+ (+ tree, forest -- ** helpers , dataSource@@ -39,7 +40,7 @@ import qualified Data.Vector.Storable as VS (Vector) import Data.RPTree.Gen (sparse, dense)-import Data.RPTree.Internal (RPTree(..), RPForest, RPT(..), levels, points, Inner(..), innerSD, innerSS, metricSSL2, metricSDL2, SVector(..), fromListSv, DVector(..), fromListDv, partitionAtMedian, RPTError(..))+import Data.RPTree.Internal (RPTree(..), RPForest, RPT(..), levels, points, Inner(..), innerSD, innerSS, metricSSL2, metricSDL2, SVector(..), fromListSv, DVector(..), fromListDv, partitionAtMedian, RPTError(..), Embed(..)) liftC :: (Monad m, MonadTrans t) => C.ConduitT i o m r -> C.ConduitT i o (t m) r@@ -63,8 +64,8 @@ -> Int -- ^ data chunk size -> Double -- ^ nonzero density of projection vectors -> Int -- ^ dimension of projection vectors- -> C.ConduitT () (v Double) m () -- ^ data source- -> m (RPTree Double (V.Vector (v Double)))+ -> C.ConduitT () (Embed v Double x) m () -- ^ data source+ -> m (RPTree Double (V.Vector (Embed v Double x))) tree seed maxDepth minLeaf n pnz dim src = do let rvs = sample seed $ V.replicateM maxDepth (sparse pnz dim stdNormal)@@ -82,7 +83,11 @@ -> Int -- ^ min leaf size -> Int -- ^ data chunk size -> V.Vector (u d) -- ^ random projection vectors- -> C.ConduitT (v d) o m (RPT d (V.Vector (v d))) + -> C.ConduitT+ (Embed v d x)+ o+ m+ (RPT d (V.Vector (Embed v d x))) insertC maxDepth minLeaf n rvs = chunkedAccum n z (insert maxDepth minLeaf rvs) where z = Tip mempty@@ -108,8 +113,8 @@ -> Int -- ^ data chunk size -> Double -- ^ nonzero density of projection vectors -> Int -- ^ dimension of projection vectors- -> C.ConduitT () (v Double) m () -- ^ data source- -> m (RPForest Double (V.Vector (v Double)))+ -> C.ConduitT () (Embed v Double x) m () -- ^ data source+ -> m (RPForest Double (V.Vector (Embed v Double x))) forest seed maxd minl ntrees chunksize pnz dim src = do let rvss = sample seed $ do@@ -130,10 +135,10 @@ -> Int -- ^ chunk size -> IM.IntMap (v1 (u d)) -- one entry per tree -> C.ConduitT- (v d)+ (Embed v d x) o m- (IM.IntMap (RPT d (V.Vector (v d))))+ (IM.IntMap (RPT d (V.Vector (Embed v d x)))) insertMultiC maxd minl n rvss = chunkedAccum n im0 (insertMulti maxd minl rvss) where im0 = IM.map (const z) rvss@@ -145,9 +150,9 @@ Int -> Int -> IM.IntMap (v1 (u d)) -- ^ projection vectors- -> IM.IntMap (RPT d (V.Vector (v d))) -- ^ accumulator of subtrees- -> V.Vector (v d) -- ^ data chunk- -> IM.IntMap (RPT d (V.Vector (v d)))+ -> IM.IntMap (RPT d (V.Vector (Embed v d x))) -- ^ accumulator of subtrees+ -> V.Vector (Embed v d x) -- ^ data chunk+ -> IM.IntMap (RPT d (V.Vector (Embed v d x))) insertMulti maxd minl rvss tacc xs = flip IM.mapWithKey tacc $ \ !i !t -> case IM.lookup i rvss of Just !rvs -> insert maxd minl rvs t xs@@ -158,9 +163,9 @@ Int -- ^ max tree depth -> Int -- ^ min leaf size -> v1 (u d) -- ^ projection vectors- -> RPT d (V.Vector (v d)) -- ^ accumulator- -> V.Vector (v d) -- ^ data chunk- -> RPT d (V.Vector (v d))+ -> RPT d (V.Vector (Embed v d x)) -- ^ accumulator+ -> V.Vector (Embed v d x) -- ^ data chunk+ -> RPT d (V.Vector (Embed v d x)) insert maxDepth minLeaf rvs = loop 0 where z = Tip mempty
src/Data/RPTree/Internal.hs view
@@ -365,8 +365,8 @@ -- | Partition the data wrt the median value of the inner product partitionAtMedian :: (Ord a, Inner u v, VU.Unbox a, Fractional a) => u a -- ^ projection vector- -> V.Vector (v a) -- ^ dataset (3 or more elements)- -> (a, Margin a, V.Vector (v a), V.Vector (v a)) -- ^ median, margin, smaller, larger+ -> V.Vector (Embed v a x) -- ^ dataset (3 or more elements)+ -> (a, Margin a, V.Vector (Embed v a x), V.Vector (Embed v a x)) -- ^ median, margin, smaller, larger partitionAtMedian r xs = (thr, margin, ll, rr) where (ll, rr) = (VG.take nh xs', VG.drop nh xs')@@ -376,7 +376,7 @@ thr = inns VG.! nh -- inner product threshold, mgl < thr < mgr n = VG.length xs -- total data size nh = n `div` 2 -- size of left partition- projs = sortByVG snd $ VG.map (\x -> (x, r `inner` x)) xs+ projs = sortByVG snd $ VG.map (\xe -> (xe, r `inner` (eEmbed xe))) xs (xs', inns) = VG.unzip projs sortByVG :: (VG.Vector v a, Ord b) => (a -> b) -> v a -> v a