rp-tree 0.3 → 0.3.1
raw patch · 7 files changed
+125/−52 lines, 7 filesbinary-added
Files
- README.md +2/−2
- app/Main.hs +2/−2
- r/scatter.png binary
- rp-tree.cabal +7/−3
- src/Data/RPTree.hs +48/−11
- src/Data/RPTree/Conduit.hs +0/−4
- src/Data/RPTree/Internal.hs +66/−30
README.md view
@@ -1,5 +1,5 @@ # rp-tree -+ -Random projection trees for approximate nearest neighbor search in high-dimensional vector spaces+Random projection trees for approximate nearest neighbor search in high-dimensional vector spaces.
app/Main.hs view
@@ -27,8 +27,8 @@ import qualified Data.Vector as V (Vector, toList, fromList, replicate, zip) import Control.Monad (replicateM)-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)+import Data.RPTree (knn, candidates, Embed(..), Inner(..), RPTree, RPForest, leaves, SVector, fromListSv, DVector, fromListDv, dense, writeCsv, forest, dataSource, sparse, normal2, normalSparse2, datS, datD)+-- import Data.RPTree.Internal.Testing (datS, datD) main :: IO () main = do -- putStrLn "hello!"
+ r/scatter.png view
binary file changed (absent → 264365 bytes)
rp-tree.cabal view
@@ -1,7 +1,9 @@ name: rp-tree-version: 0.3+version: 0.3.1 synopsis: Random projection trees description: Random projection trees for approximate nearest neighbor search in high-dimensional vector spaces+ .+ The entry point to the library and all documentation is in the "Data.RPTree" module. homepage: https://github.com/ocramz/rp-tree license: BSD3 license-file: LICENSE@@ -12,7 +14,8 @@ build-type: Simple extra-source-files: README.md Changelog.md-cabal-version: >=1.10+extra-doc-files: r/scatter.png+cabal-version: >=1.18 tested-with: GHC == 8.10.4 library@@ -20,11 +23,12 @@ ghc-options: -Wall hs-source-dirs: src exposed-modules: Data.RPTree- Data.RPTree.Internal.Testing+ other-modules: Data.RPTree.Internal Data.RPTree.Gen Data.RPTree.Draw Data.RPTree.Conduit+ Data.RPTree.Internal.Testing build-depends: base >= 4.7 && < 5 , boxes , bytestring
src/Data/RPTree.hs view
@@ -11,11 +11,46 @@ {-# options_ghc -Wno-unused-top-binds #-} {-|-Random projection trees for approximate nearest neighbor search in high-dimensional vector spaces+Random projection trees for approximate nearest neighbor search in high-dimensional vector spaces.++== Introduction++Similarity search is a common problem in many fields (imaging, natural language processing, ..), and is often one building block of a larger data processing system.++There are many ways to /embed/ data in a vector space such that similarity search can be recast as a geometrical nearest neighbor lookup.++In turn, the efficiency and effectiveness of querying such a vector database strongly depends on how internally the data index is represented, graphs and trees being two common approaches.++The naive, all-pairs exact search becomes impractical even at moderate data sizes, which motivated research into approximate indexing methods.+++== Overview++This library provides a /tree/-based approach to approximate nearest neighbor search. The database is recursively partitioned according to a series of random projections, and this partitioning is logically arranged as a tree which allows for rapid lookup.++Internally, a single random projection vector is sampled per tree level, as proposed in [1]. The projection vectors in turn can be sparse with a tunable sparsity parameter, which can help compressing the database at a small accuracy cost.++Retrieval accuracy can be improved by populating multiple trees (i.e. a /random forest/), and intersecting the results of the same query against each of them.++== Quick Start++1) Build an index with 'forest'++2) Lookup the \(k\) nearest neighbors to a query point with 'knn'++3) The database can be serialised and restored with 'serialiseRPForest' and 'deserialiseRPForest', respectively.++++== References++1) Hyvonen, V., et al., Fast Nearest Neighbor Search through Sparse Random Projections and Voting, https://www.cs.helsinki.fi/u/ttonteri/pub/bigdata2016.pdf+ -} module Data.RPTree ( -- * Construction- tree, forest+ -- tree,+ forest -- * Query , knn -- * I/O@@ -41,13 +76,8 @@ , metricSSL2, metricSDL2 -- *** Scale , scaleS, scaleD- -- * Conduit- , dataSource- -- * Random generation- -- ** vector- , sparse, dense- , normal2 + -- * Rendering , draw -- * CSV@@ -55,6 +85,13 @@ -- * Testing , randSeed, BenchConfig(..), normalSparse2 , liftC+ -- ** Random generation+ -- *** Conduit+ , dataSource+ , datS, datD+ -- *** Vector data+ , sparse, dense+ , normal2 ) where import Control.Monad (replicateM)@@ -92,17 +129,17 @@ import Data.RPTree.Conduit (tree, forest, dataSource, liftC) import Data.RPTree.Gen (sparse, dense, normal2, normalSparse2) 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.Internal.Testing (BenchConfig(..), randSeed, datS, datD) import Data.RPTree.Draw (draw, writeCsv) --- | k nearest neighbors+-- | Look up the \(k\) nearest neighbors to a query point knn :: (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) -- ^ ordered in increasing distance order+ -> V.Vector (p, Embed u d x) -- ^ ordered in increasing distance order to the query point knn distf k tts q = sortByVG fst cs where cs = VG.map (\xe -> (eEmbed xe `distf` q, xe)) $ VG.take k $ fold $ (`candidates` q) <$> tts
src/Data/RPTree/Conduit.hs view
@@ -55,8 +55,6 @@ -- * stationary : each chunk is representative of the whole dataset -- -- * bounded : we wait until the end of the stream to produce a result------ Throws 'EmptyResult' if the conduit is empty tree :: (Monad m, Inner SVector v) => Word64 -- ^ random seed -> Int -- ^ max tree depth@@ -103,8 +101,6 @@ -- * stationary : each chunk is representative of the whole dataset -- -- * bounded : we wait until the end of the stream to produce a result------ Throws 'EmptyResult' if the conduit is empty forest :: (Monad m, Inner SVector v) => Word64 -- ^ random seed -> Int -- ^ max tree depth
src/Data/RPTree/Internal.hs view
@@ -44,16 +44,18 @@ import Control.Monad.Trans.State (StateT(..), runStateT, evalStateT, State, runState, evalState) -- vector import qualified Data.Vector as V (Vector, replicateM, fromList)-import qualified Data.Vector.Generic as VG (Vector(..), map, sum, unfoldr, unfoldrM, length, replicateM, (!), take, drop, unzip, freeze, thaw, foldl, foldr, toList, zipWith, last, head)+import qualified Data.Vector.Generic as VG (Vector(..), map, sum, unfoldr, unfoldrM, length, replicateM, (!), (!?), take, drop, unzip, freeze, thaw, foldl, foldr, toList, zipWith, last, head, imap) import qualified Data.Vector.Unboxed as VU (Vector, Unbox, fromList, toList) import qualified Data.Vector.Storable as VS (Vector) -- vector-algorithms import qualified Data.Vector.Algorithms.Merge as V (sortBy) --- | Pair a datum with a vector embedding+-- | Pairing of a data item with its vector embedding+--+-- The vector is used internally for indexing data Embed v e a = Embed {- eEmbed :: !(v e) -- ^ the embedding is a vector- , eData :: !a+ eEmbed :: !(v e) -- ^ vector embedding+ , eData :: !a -- ^ data item } deriving (Eq, Ord, Show, Generic, Functor) instance (NFData (v e), NFData a) => NFData (Embed v e a) instance (Serialise (v e), Serialise a) => Serialise (Embed v e a)@@ -180,6 +182,7 @@ rpTreeData :: Traversal' (RPTree d a) a rpTreeData = rpTree . rpData +-- | All data buckets stored at the leaves of the tree leaves :: RPTree d a -> [a] leaves = (^.. rpTreeData) @@ -210,8 +213,8 @@ class (Scale u, Scale v) => Inner u v where inner :: (VU.Unbox a, Num a) => u a -> v a -> a metricL2 :: (VU.Unbox a, Floating a) => u a -> v a -> a- (^+^) :: (VU.Unbox a, Num a) => u a -> v a -> u a- (^-^) :: (VU.Unbox a, Num a) => u a -> v a -> u a+ (^+^) :: (VU.Unbox a, Num a) => u a -> v a -> v a+ (^-^) :: (VU.Unbox a, Num a) => u a -> v a -> v a instance Inner SVector SVector where inner (SV _ v1) (SV _ v2) = innerSS v1 v2@@ -221,13 +224,13 @@ instance Inner SVector VU.Vector where inner (SV _ v1) v2 = innerSD v1 v2 metricL2 (SV _ v1) v2 = metricSDL2 v1 v2- (SV n v1) ^+^ v2 = SV n $ sumSD v1 v2- (SV n v1) ^-^ v2 = SV n $ diffSD v1 v2+ (SV _ v1) ^+^ v2 = sumSD v1 v2+ (SV _ v1) ^-^ v2 = diffSD v1 v2 instance Inner SVector DVector where inner (SV _ v1) (DV v2) = innerSD v1 v2 metricL2 (SV _ v1) (DV v2) = metricSDL2 v1 v2- (SV n v1) ^+^ (DV v2) = SV n $ sumSD v1 v2- (SV n v1) ^-^ (DV v2) = SV n $ diffSD v1 v2+ (SV _ v1) ^+^ (DV v2) = DV $ sumSD v1 v2+ (SV _ v1) ^-^ (DV v2) = DV $ diffSD v1 v2 instance Inner DVector DVector where inner (DV v1) (DV v2) = innerDD v1 v2 metricL2 (DV v1) (DV v2) = metricDDL2 v1 v2@@ -287,10 +290,9 @@ duv = u `diffSS` v -- | Vector distance induced by the L2 norm (sparse-dense)-metricSDL2 :: (Floating a, VG.Vector v1 a, VU.Unbox a,- VG.Vector v1 (Int, a), VG.Vector v2 a) =>- v1 (Int, a) -> v2 a -> a-metricSDL2 u v = sqrt $ VG.sum $ VG.map (\(_, x) -> x ** 2) duv+metricSDL2 :: (Floating a, VU.Unbox a, VG.Vector u (Int, a), VG.Vector v a) =>+ u (Int, a) -> v a -> a+metricSDL2 u v = sqrt $ VG.sum $ VG.map (** 2) duv where duv = u `diffSD` v @@ -308,8 +310,8 @@ -- | Vector sum sumSD :: (VG.Vector u (Int, a), VG.Vector v a, VU.Unbox a, Num a) =>- u (Int, a) -> v a -> u (Int, a)-sumSD = binSD (-)+ u (Int, a) -> v a -> v a+sumSD = binSDD (+) 0 -- | Vector sum sumSS :: (VG.Vector u (Int, a), VG.Vector v (Int, a), VU.Unbox a, Num a) =>@@ -318,8 +320,8 @@ -- | Vector difference diffSD :: (VG.Vector u (Int, a), VG.Vector v a, VU.Unbox a, Num a) =>- u (Int, a) -> v a -> u (Int, a)-diffSD = binSD (-)+ u (Int, a) -> v a -> v a+diffSD = binSDD (-) 0 -- | Vector difference diffSS :: (VG.Vector u (Int, a), VG.Vector v (Int, a), VU.Unbox a, Num a) =>@@ -344,23 +346,37 @@ LT -> Just ((il, f xl z ), (succ i1, i2 )) GT -> Just ((ir, f z xr), (i1 , succ i2)) ---- FIXME the return type of a sparse-dense binary operation depends on the operator itself (S * D = S , S + D = D ), so 'binSD' must be changed-binSD :: (VG.Vector u (Int, a), VG.Vector v a, VU.Unbox a) =>- (a -> a -> a) -> u (Int, a) -> v a -> u (Int, a)-binSD f vv1 vv2 = VG.unfoldr go 0+-- | sparse * dense -> dense+--+-- e.g. vector sum, difference+binSDD :: (VG.Vector v1 a, VG.Vector v2 p, VG.Vector v3 (Int, p)) =>+ (p -> p -> a) -> p -> v3 (Int, p) -> v2 p -> v1 a+binSDD f z vv1 vv2 = VG.unfoldr go (0, 0) where nz1 = VG.length vv1 nz2 = VG.length vv2- go i- | i >= nz1 || i >= nz2 = Nothing- | otherwise = Just ((il, y), succ i)- where- (il, xl) = vv1 VG.! i- xr = vv2 VG.! il- y = f xl xr+ go (i1, i2)+ | i1 >= nz1 || i2 >= nz2 = Nothing+ | otherwise =+ let+ (il, xl) = vv1 VG.! i1+ xr = vv2 VG.! i2+ in case il `compare` i2 of+ EQ -> Just (f xl xr, (succ i1, succ i2))+ LT -> Just (f xl z , (succ i1, i2 ))+ GT -> Just (f z xr, (i1 , succ i2)) +{-+- b0+- b1+a2 b2+- b3+a4 b4+-} +++ {-# SCC partitionAtMedian #-} -- | Partition the data wrt the median value of the inner product partitionAtMedian :: (Ord a, Inner u v, VU.Unbox a, Fractional a) =>@@ -420,3 +436,23 @@ -- makeLensesFor [("_idD", "idD")] ''Id -- instance (Eq a) => Ord (Id a) where -- Id _ u1 <= Id _ u2 = u1 <= u2+++++++-- -- FIXME the return type of a sparse-dense binary operation depends on the operator itself (S * D = S , S + D = D ), so 'binSD' must be changed+-- binSD :: (VG.Vector u (Int, a), VG.Vector v a) =>+-- (a -> a -> a) -> u (Int, a) -> v a -> u (Int, a)+-- binSD f vv1 vv2 = VG.unfoldr go 0+-- where+-- nz1 = VG.length vv1+-- nz2 = VG.length vv2+-- go i+-- | i >= nz1 || i >= nz2 = Nothing+-- | otherwise = Just ((il, y), succ i)+-- where+-- (il, xl) = vv1 VG.! i+-- xr = vv2 VG.! il+-- y = f xl xr