diff --git a/Changelog.md b/Changelog.md --- a/Changelog.md +++ b/Changelog.md @@ -1,3 +1,7 @@ +0.7 + +- add batch index creation 'treeBatch', 'forestBatch', 'dataBatch' and corresponding unit tests (#5) + 0.6 - add 'knnPQ' based on 'heaps', to be confirmed how efficient it is diff --git a/app/Main.hs b/app/Main.hs --- a/app/Main.hs +++ b/app/Main.hs @@ -15,7 +15,7 @@ import qualified Data.Conduit.Combinators as C (map, mapM, scanl, scanlM, last, print, sinkVector, sinkList) import qualified Data.Conduit.List as C (chunksOf, unfold, unfoldM) -- containers -import qualified Data.IntMap as IM (IntMap, fromList, insert, lookup, map, mapWithKey, traverseWithKey, foldlWithKey, foldrWithKey) +import qualified Data.IntMap as IM (IntMap, fromList, insert, lookup, map, mapWithKey, traverseWithKey, foldlWithKey, foldrWithKey, singleton) -- -- exceptions -- import Control.Monad.Catch (MonadThrow(..)) -- -- mnist-idx-conduit @@ -27,42 +27,65 @@ import Control.Monad.Trans.Class (MonadTrans(..)) -- vector import qualified Data.Vector as V (Vector, toList, fromList, replicate, zip, zipWith) - -import Control.Monad (replicateM) -import Data.RPTree (knn, candidates, rpTreeCfg, RPTreeConfig(..), Embed(..), Inner(..), RPTree, RPForest, SVector, fromListSv, DVector, fromListDv, dense, writeCsv, tree, forest, dataSource, sparse, normal2, normalSparse2, datS, datD, circle2d, leaves, levels, treeSize, leafSizes, writeDot) +import qualified Data.Vector.Unboxed as VU (Unbox) +import Data.RPTree (knn, knnH, candidates, rpTreeCfg, RPTreeConfig(..), Embed(..), Inner(..), RPTree, RPForest, SVector, fromListSv, DVector, fromListDv, dense, writeCsv, knnWriteCsv, tree, forest, dataSource, sparse, normal2, normalSparse2, datS, datD, circle2d, leaves, levels, treeSize, leafSizes, writeDot) -- import Data.RPTree.Internal.Testing (datS, datD) main :: IO () main = do let - n = 30000 - maxd = 5 + n = 10000 minl = 10 - chunk = 500 dim = 2 -- cfg = rpTreeCfg n dim - cfg = RPCfg maxd chunk 1.0 - csvTree0 n minl cfg - tree0dot n minl cfg + -- (RPCfg maxd chunk _) = rpTreeCfg minl n dim + maxd = 5 + chunk = 100 + tt = tree0 n maxd minl chunk + csvTree0 tt + tree0dot tt + csvKnnTree0 tt +csvKnnTree0 :: (Show a1, VU.Unbox a1, RealFloat a1) => + RPTree a1 () (V.Vector (Embed DVector a1 a)) -> IO () +csvKnnTree0 tt = do + let + ttlab = prep 0 tt -- label leaves starting from 0 + q = fromListDv [1, 1] -- query + tts = IM.singleton 0 tt + k = 10 + labf v = (v, -1) -- labelling function for KNN points + hits = labf <$> knnL2 k tts q + hitsH = labf <$> knnHL2 k tts q + knnWriteCsv "r/scatter_knn.csv" ttlab hits + knnWriteCsv "r/scatter_knnH.csv" ttlab hitsH +knnL2, knnHL2 :: (VU.Unbox p, RealFloat p, Inner SVector v, Inner u v) => + Int + -> RPForest p (V.Vector (Embed u p x)) + -> v p + -> V.Vector (u p) +knnL2 n ff q = eEmbed . snd <$> knn metricL2 n ff q +knnHL2 n ff q = eEmbed . snd <$> knnH metricL2 n ff q -tree0dot :: Int -> Int -> RPTreeConfig -> IO () -tree0dot n minl (RPCfg maxd chunk _) = - writeDot f fpath "tree0" $ tree0 n maxd minl chunk +-- render the tree with graphviz +tree0dot :: (Ord (t a), Foldable t) => RPTree d x (t a) -> IO () +tree0dot = writeDot f fpath "tree0" where f = show . length fpath = "tree0.dot" -csvTree0 :: Int -> Int -> RPTreeConfig -> IO () -csvTree0 n minl (RPCfg maxd chunk _) = do +-- scatter the whole dataset with distinct colors for the contents of each leaf, render as a CSV +csvTree0 :: (VU.Unbox a1, Show a1, Traversable t) => + t (V.Vector (Embed DVector a1 a2)) -> IO () +csvTree0 tt = do let - tt = tree0 n maxd minl chunk - ttlab = prep tt + ttlab = prep A tt writeCsv "r/scatter_data_2.csv" ttlab -prep :: (Traversable t) => t (V.Vector (Embed v e a)) -> t (V.Vector (v e, Pal)) -prep = flip evalState A . traverse labeled +prep :: (Traversable t, Enum s) => + s -> t (V.Vector (Embed v e a)) -> t (V.Vector (v e, s)) +prep x0 = flip evalState x0 . traverse labeled labeled :: (Enum b) => V.Vector (Embed v e a) @@ -159,7 +182,7 @@ -- liftC = C.transPipe lift embedC :: Monad m => C.ConduitT (v e) (Embed v e ()) m () -embedC = C.map (\ x -> Embed x ()) +embedC = C.map (`Embed` ()) diff --git a/rp-tree.cabal b/rp-tree.cabal --- a/rp-tree.cabal +++ b/rp-tree.cabal @@ -1,5 +1,5 @@ name: rp-tree -version: 0.6 +version: 0.7 synopsis: Random projection trees description: Random projection trees for approximate nearest neighbor search in high-dimensional vector spaces . @@ -25,10 +25,12 @@ exposed-modules: Data.RPTree other-modules: Data.RPTree.Internal + Data.RPTree.Batch Data.RPTree.Gen Data.RPTree.Draw Data.RPTree.Conduit Data.RPTree.Internal.Testing + Data.RPTree.Internal.MedianHeap build-depends: base >= 4.7 && < 5 , boxes , bytestring @@ -62,6 +64,7 @@ , hspec , QuickCheck , splitmix-distributions + , vector benchmark bench-time default-language: Haskell2010 diff --git a/src/Data/RPTree.hs b/src/Data/RPTree.hs --- a/src/Data/RPTree.hs +++ b/src/Data/RPTree.hs @@ -49,14 +49,18 @@ -} module Data.RPTree ( -- * Construction - tree + -- ** Batch + treeBatch + , forestBatch + -- ** Incremental (Conduit-based) + , tree , forest -- ** Parameters , rpTreeCfg, RPTreeConfig(..) -- , ForestParams -- * Query , knn - , knnPQ + , knnH -- * I/O , serialiseRPForest , deserialiseRPForest @@ -89,7 +93,7 @@ -- * Rendering -- , draw -- ** CSV - , writeCsv + , writeCsv, knnWriteCsv -- ** GraphViz dot , writeDot -- * Testing @@ -97,6 +101,8 @@ , liftC -- ** Random generation , randSeed + -- *** Batch + , dataBatch -- *** Conduit , dataSource , datS, datD @@ -125,7 +131,7 @@ -- deepseq import Control.DeepSeq (NFData(..)) -- heaps -import qualified Data.Heap as H (Heap, fromList, insert, Entry(..), empty, group, viewMin, map) +import qualified Data.Heap as H (Heap, fromList, insert, Entry(..), empty, group, viewMin, map, union) -- -- psqueues -- import qualified Data.IntPSQ as PQ (IntPSQ, findMin, minView, empty, insert, fromList, toList) -- transformers @@ -139,13 +145,16 @@ -- vector-algorithms import qualified Data.Vector.Algorithms.Merge as V (sortBy) +import Data.RPTree.Batch (treeBatch, forestBatch, dataBatch) import Data.RPTree.Conduit (tree, forest, dataSource, liftC, rpTreeCfg, RPTreeConfig(..)) import Data.RPTree.Gen (sparse, dense, normal2, normalSparse2, circle2d) import Data.RPTree.Internal (RPTree(..), RPForest, RPT(..), Embed(..), leaves, levels, points, 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, datS, datD) -import Data.RPTree.Draw (writeDot, writeCsv) +import Data.RPTree.Draw (writeDot, writeCsv, knnWriteCsv) + + -- | Look up the \(k\) nearest neighbors to a query point -- -- The supplied distance function @d@ must satisfy the definition of a metric, i.e. @@ -178,10 +187,38 @@ where 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 + insf acc x = H.insert (H.Entry p x) acc + where + p = eEmbed x `distf` q +-- | Same as 'knn' but based on min-'H.Heap' +-- +-- Leaves are prioritized according to their margin +knnH :: (Ord p, Inner SVector v, VU.Unbox d, Fractional d, Ord 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 to the query point +knnH distf k tts q = VG.map (\xe -> (eEmbed xe `distf` q, xe)) $ go mempty 0 htot + where + htot = unions $ (`candidatesH` q) <$> tts + go acc n hh = case H.viewMin hh of + Nothing -> acc + Just ((H.Entry _ xsh), hrest) -> + let + nels = length xsh + ntot = nels + n + in + if ntot > k && not (null acc) + then acc + else go (xsh <> acc) ntot hrest + + +unions :: Foldable t => t (H.Heap a) -> H.Heap a +unions = foldl H.union H.empty + -- | deduplicate nub :: (Ord a) => H.Heap a -> [a] nub = foldMap maybeToList . H.map view . H.group @@ -276,9 +313,82 @@ | otherwise -> go i' rtree +-- | RPTree margins used as search priority (using a min-heaps : lower margin leaves will be ranked highest and therefore searched first) +candidatesH :: (Inner SVector v, VU.Unbox d, Ord d, Fractional d) => + RPTree d l a + -> v d -- ^ query point + -> H.Heap (H.Entry d a) +candidatesH (RPTree rvs tt) x = go 0 tt H.empty infty + where + infty = 1 / 0 + go _ (Tip _ xs) acc p = insertp p xs acc + go ixLev (Bin _ thr margin ltree rtree) acc p = + 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 + pl = p `min` dl + pr = p `min` dr + in if + | proj < thr && + dl > dr -> H.union (go i' ltree acc pl) (go i' rtree acc pr) + | proj < thr -> go i' ltree acc pl + | proj > thr && + dl < dr -> H.union (go i' ltree acc pl) (go i' rtree acc pr) + | otherwise -> go i' rtree acc pr +insertp :: Ord p => + p -> a -> H.Heap (H.Entry p a) -> H.Heap (H.Entry p a) +insertp p x = H.insert (H.Entry p x) + + +data RPTreeStats = RPTreeStats { + rptsLength :: Int + } deriving (Eq, Show) + +treeStats :: RPTree d l a -> RPTreeStats +treeStats (RPTree _ tt) = RPTreeStats l + where + l = length tt + + +-- | How many data items are stored in the 'RPTree' +treeSize :: (Foldable t) => RPTree d l (t a) -> Int +treeSize = sum . leafSizes + +-- | How many data items are stored in each leaf of the 'RPTree' +leafSizes :: Foldable t => RPTree d l (t a) -> RPT d l Int +leafSizes (RPTree _ tt) = length <$> tt + + + +-- candidatesPQ :: (Ord a, Ord d, Inner SVector v, VU.Unbox d, Num d) => +-- RPTree d l a -> v d -> H.Heap a +-- candidatesPQ (RPTree rvs tt) x = go 0 tt H.empty +-- where +-- go _ (Tip _ xs) acc = H.insert xs acc +-- go ixLev (Bin _ thr margin ltree rtree) acc = +-- 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 +-- in if +-- | proj < thr && +-- dl > dr -> H.union (go i' ltree acc) (go i' rtree acc) +-- | proj < thr -> go i' ltree acc +-- | proj > thr && +-- dl < dr -> H.union (go i' ltree acc) (go i' rtree acc) +-- | otherwise -> go i' rtree acc + + -- -- | 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 @@ -340,23 +450,7 @@ -data RPTreeStats = RPTreeStats { - rptsLength :: Int - } deriving (Eq, Show) -treeStats :: RPTree d l a -> RPTreeStats -treeStats (RPTree _ tt) = RPTreeStats l - where - l = length tt - - --- | How many data items are stored in the 'RPTree' -treeSize :: (Foldable t) => RPTree d l (t a) -> Int -treeSize = sum . leafSizes - --- | How many data items are stored in each leaf of the 'RPTree' -leafSizes :: Foldable t => RPTree d l (t a) -> RPT d l Int -leafSizes (RPTree _ tt) = length <$> tt -- pqSeq :: Ord a => PQ.IntPSQ a b -> Seq (a, b) -- pqSeq pqq = go pqq mempty diff --git a/src/Data/RPTree/Batch.hs b/src/Data/RPTree/Batch.hs new file mode 100644 --- /dev/null +++ b/src/Data/RPTree/Batch.hs @@ -0,0 +1,75 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# options_ghc -Wno-unused-imports #-} +module Data.RPTree.Batch ( + treeBatch, forestBatch, + -- * utils + dataBatch + ) where + +import Control.Monad (replicateM) +import GHC.Word (Word64) + +-- containers +import qualified Data.IntMap.Strict as IM (IntMap, fromList, insert, lookup, map, mapWithKey, traverseWithKey, foldlWithKey, foldrWithKey, intersectionWith) +-- splitmix-distributions +import System.Random.SplitMix.Distributions (Gen, sample, GenT, sampleT, stdNormal) +-- 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.Unboxed as VU (Vector, Unbox, fromList) + +import Data.RPTree.Gen (sparse, dense) +import Data.RPTree.Internal (RPTree(..), RPForest, RPT(..), create, createMulti, SVector, Inner(..), Embed(..)) + +-- | Populate a tree from a dataset +-- +-- Assumptions on the data source: +-- +-- * non-empty : contains at least one value +treeBatch :: Inner SVector v => + Word64 -- ^ random seed + -> Int -- ^ max tree depth + -> Int -- ^ min leaf size + -> Double -- ^ nonzero density of projection vectors + -> Int -- ^ dimension of projection vectors + -> V.Vector (Embed v Double x) -- ^ dataset + -> RPTree Double () (V.Vector (Embed v Double x)) +treeBatch seed maxDepth minLeaf pnz dim src = + let + rvs = sample seed $ V.replicateM maxDepth (sparse pnz dim stdNormal) + t = create maxDepth minLeaf rvs src + in RPTree rvs t + +-- | Populate a forest from a data stream +-- +-- Assumptions on the data source: +-- +-- * non-empty : contains at least one value +forestBatch :: (Inner SVector v) => + Word64 -- ^ random seed + -> Int -- ^ max tree depth, \(l > 1\) + -> Int -- ^ min leaf size, \(m_{leaf} > 1\) + -> Int -- ^ number of trees, \(n_t > 1\) + -> Double -- ^ nonzero density of projection vectors, \(p_{nz} \in (0, 1)\) + -> Int -- ^ dimension of projection vectors, \(d > 1\) + -> V.Vector (Embed v Double x) -- ^ dataset + -> RPForest Double (V.Vector (Embed v Double x)) +forestBatch seed maxd minl ntrees pnz dim src = + let + rvss = sample seed $ do + rvs <- replicateM ntrees $ V.replicateM maxd (sparse pnz dim stdNormal) + pure $ IM.fromList $ zip [0 .. ] rvs + ts = createMulti maxd minl rvss src + in IM.intersectionWith RPTree rvss ts + +-- | Batch random data points +dataBatch :: (Monad m, VG.Vector v a) => + Int -- ^ number of points to generate + -> GenT m a -- ^ random point generator + -> GenT m (v a) +dataBatch n gg = flip VG.unfoldrM 0 $ \i -> do + if i == n + then pure Nothing + else do + x <- gg + pure $ Just (x, i + 1) diff --git a/src/Data/RPTree/Conduit.hs b/src/Data/RPTree/Conduit.hs --- a/src/Data/RPTree/Conduit.hs +++ b/src/Data/RPTree/Conduit.hs @@ -40,8 +40,8 @@ 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(..), Embed(..)) - +import Data.RPTree.Internal (RPTree(..), RPForest, RPT(..), insert, insertMulti, levels, points, Inner(..), innerSD, innerSS, metricSSL2, metricSDL2, SVector(..), fromListSv, DVector(..), fromListDv, partitionAtMedian, RPTError(..), Embed(..)) +import qualified Data.RPTree.Internal.MedianHeap as MH (MedianHeap, insert, median) liftC :: (Monad m, MonadTrans t) => C.ConduitT i o m r -> C.ConduitT i o (t m) r liftC = C.transPipe lift @@ -160,60 +160,7 @@ z = Tip () mempty -{-# SCC insertMulti #-} -insertMulti :: (Ord d, Inner u v, VU.Unbox d, Fractional d, VG.Vector v1 (u d)) => - Int - -> Int - -> IM.IntMap (v1 (u d)) -- ^ projection vectors - -> 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 - _ -> t -{-# SCC insert #-} -insert :: (VG.Vector v1 (u d), Ord d, Inner u v, VU.Unbox d, Fractional d) => - Int -- ^ max tree depth - -> Int -- ^ min leaf size - -> v1 (u d) -- ^ projection vectors - -> 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 - loop ixLev !tt xs = - let - r = rvs VG.! ixLev -- proj vector for current level - in - case tt of - - b@(Bin _ thr0 margin0 tl0 tr0) -> - if ixLev >= maxDepth - then b -- return current subtree - else - case partitionAtMedian r xs of - Nothing -> Tip () mempty - Just (thr, margin, ll, rr) -> Bin () thr' margin' tl tr - where - margin' = margin0 <> margin - thr' = (thr0 + thr) / 2 - tl = loop (ixLev + 1) tl0 ll - tr = loop (ixLev + 1) tr0 rr - - Tip _ xs0 -> do - let xs' = xs <> xs0 - if ixLev >= maxDepth || length xs' <= minLeaf - then Tip () xs' -- concat data in leaf - else - case partitionAtMedian r xs' of - Nothing -> Tip () mempty - Just (thr, margin, ll, rr) -> Bin () thr margin tl tr - where - tl = loop (ixLev + 1) z ll - tr = loop (ixLev + 1) z rr diff --git a/src/Data/RPTree/Draw.hs b/src/Data/RPTree/Draw.hs --- a/src/Data/RPTree/Draw.hs +++ b/src/Data/RPTree/Draw.hs @@ -5,6 +5,7 @@ module Data.RPTree.Draw ( -- * CSV writeCsv + , knnWriteCsv -- * GraphViz dot , writeDot -- , draw @@ -39,7 +40,16 @@ import Data.RPTree.Internal (RPTree(..), RPT(..), DVector, toListDv) - +knnWriteCsv :: + (Foldable t1, Foldable t2, Show a1, Show b1, Show a2, Show b2, VU.Unbox a1, VU.Unbox a2) => + FilePath + -> t1 (V.Vector (DVector a1, b1)) + -> t2 (DVector a2, b2) -- ^ points obtained from knn + -> IO () +knnWriteCsv fp ds knnds = TL.writeFile fp $ TLB.toLazyText bld + where + bld = toCsvV ds <> + toCsv knnds -- | Encode dataset as CSV and save into file writeCsv :: (Foldable t, VU.Unbox a, Show a, Show b) => @@ -50,8 +60,11 @@ toCsvV :: (Foldable t, VU.Unbox a, Show a, Show b) => t (V.Vector (DVector a, b)) -> TLB.Builder -toCsvV = foldMap (\v -> foldMap (\(r, i) -> toCsvRow r i <> newline ) v) +toCsvV = foldMap toCsv +toCsv :: (Foldable t, Show a, Show b, VU.Unbox a) => t (DVector a, b) -> TLB.Builder +toCsv = foldMap (\(r, i) -> toCsvRow r i <> newline) + toCsvRow :: (Show a, Show b, VU.Unbox a) => DVector a -> b @@ -75,7 +88,7 @@ -> String -- ^ graph name -> RPTree d x t -> IO () -writeDot f fp name tt = TL.writeFile fp (toDot f name tt) +writeDot f fp name tt = TL.writeFile fp (toDot f name tt) toDot :: Ord a => (a -> String) -> String -> RPTree d x a -> TL.Text toDot f name (RPTree _ tt) = TLB.toLazyText $ open <> x <> close diff --git a/src/Data/RPTree/Internal.hs b/src/Data/RPTree/Internal.hs --- a/src/Data/RPTree/Internal.hs +++ b/src/Data/RPTree/Internal.hs @@ -33,7 +33,7 @@ -- bytestring import qualified Data.ByteString.Lazy as LBS (ByteString, toStrict, fromStrict) -- containers -import qualified Data.IntMap.Strict as IM (IntMap, fromList) +import qualified Data.IntMap.Strict as IM (IntMap, fromList, map, lookup, mapWithKey) -- deepseq import Control.DeepSeq (NFData(..)) -- -- microlens @@ -206,6 +206,98 @@ -- | Set of data points used to construct the index points :: Monoid m => RPTree d l m -> m points (RPTree _ t) = fold t + + + + + + + +-- | Batch (= non-incremental) creation of a 'RPT' +create :: (Ord d, Inner u v, VU.Unbox d, Fractional d, VG.Vector v1 (u d)) => + Int -- ^ max tree depth + -> Int -- ^ min leaf size + -> v1 (u d) -- ^ projection vectors + -> V.Vector (Embed v d x) -- ^ dataset + -> RPT d () (V.Vector (Embed v d x)) +create maxDepth minLeaf rvs = insert maxDepth minLeaf rvs z + where + z = Tip () mempty + +-- | Batch (= non-incremental) creation of multiple 'RPT's +createMulti :: (Ord d, Inner u v, VU.Unbox d, Fractional d, VG.Vector v1 (u d)) => + Int + -> Int + -> IM.IntMap (v1 (u d)) -- ^ projection vectors + -> V.Vector (Embed v d x) -- ^ dataset + -> IM.IntMap (RPT d () (V.Vector (Embed v d x))) +createMulti maxd minl rvss xs = + flip IM.mapWithKey im0 $ \i t -> case IM.lookup i rvss of + Just rvs -> create maxd minl rvs xs + _ -> t + where + im0 = IM.map (const z) rvss + z = Tip () mempty + + + +{-# SCC insertMulti #-} +insertMulti :: (Ord d, Inner u v, VU.Unbox d, Fractional d, VG.Vector v1 (u d)) => + Int + -> Int + -> IM.IntMap (v1 (u d)) -- ^ projection vectors + -> 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 + _ -> t + +{-# SCC insert #-} +insert :: (VG.Vector v1 (u d), Ord d, Inner u v, VU.Unbox d, Fractional d) => + Int -- ^ max tree depth + -> Int -- ^ min leaf size + -> v1 (u d) -- ^ projection vectors + -> 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 + loop ixLev !tt xs = + let + r = rvs VG.! ixLev -- proj vector for current level + in + case tt of + + b@(Bin _ thr0 margin0 tl0 tr0) -> + if ixLev >= maxDepth + then b -- return current subtree + else + case partitionAtMedian r xs of + Nothing -> Tip () mempty + Just (thr, margin, ll, rr) -> Bin () thr' margin' tl tr + where + margin' = margin0 <> margin + thr' = (thr0 + thr) / 2 + tl = loop (ixLev + 1) tl0 ll + tr = loop (ixLev + 1) tr0 rr + + Tip _ xs0 -> do + let xs' = xs <> xs0 + if ixLev >= maxDepth || length xs' <= minLeaf + then Tip () xs' -- concat data in leaf + else + case partitionAtMedian r xs' of + Nothing -> Tip () mempty + Just (thr, margin, ll, rr) -> Bin () thr margin tl tr + where + tl = loop (ixLev + 1) z ll + tr = loop (ixLev + 1) z rr + + + -- | Scale a vector diff --git a/src/Data/RPTree/Internal/MedianHeap.hs b/src/Data/RPTree/Internal/MedianHeap.hs new file mode 100644 --- /dev/null +++ b/src/Data/RPTree/Internal/MedianHeap.hs @@ -0,0 +1,77 @@ +{-# language CPP #-} +{-# OPTIONS_GHC -Wno-unused-imports #-} +module Data.RPTree.Internal.MedianHeap ( + MedianHeap, + insert, fromList, + median) where + +import Data.Bifunctor (Bifunctor(..)) +import Data.Ord (Down(..)) + +import qualified Data.Heap as H (Heap, Entry(..), viewMin, insert, singleton) + +type MinHeap p a = H.Heap (H.Entry p a) +type MaxHeap p a = H.Heap (H.Entry (Down p) a) +-- | The items @a@ in a 'MedianHeap' have type @p@ +data MedianHeap p a = MedianHeap (MaxHeap p a) (MinHeap p a) deriving (Eq, Show) +instance Semigroup (MedianHeap p a) where + MedianHeap l1 r1 <> MedianHeap l2 r2 = MedianHeap (l1 <> l2) (r1 <> r2) +instance Monoid (MedianHeap p a) where + mempty = MedianHeap mempty mempty +insertMax :: Ord p => p -> a -> MaxHeap p a -> MaxHeap p a +insertMax p x = H.insert (H.Entry (Down p) x) +insertMin :: Ord p => p -> a -> MinHeap p a -> MinHeap p a +insertMin p x = H.insert (H.Entry p x) + +fromList :: (Foldable t, Fractional p, Ord p) => t (p, a) -> MedianHeap p a +fromList = foldr (uncurry insert) mempty + +-- | Insert a weighted entry in the heap +insert :: (Fractional p, Ord p) => + p -- ^ weight + -> a + -> MedianHeap p a + -> MedianHeap p a +insert p x heap@(MedianHeap ll rr) = + case median heap of + Nothing -> MedianHeap (insertMax p x mempty) rr + Just q -> + if q > p + then balance $ MedianHeap ll (insertMin p x rr) + else balance $ MedianHeap (insertMax p x ll) rr + +balance :: Ord p => MedianHeap p a -> MedianHeap p a +balance (MedianHeap ll rr) + | length ll == length rr = MedianHeap ll rr + | length ll > length rr = + let (H.Entry (Down p) x, xs) = deconstruct ll + in MedianHeap xs (insertMin p x rr) + | otherwise = + let (H.Entry p x, xs) = deconstruct rr + in MedianHeap (insertMax p x ll) xs + where + deconstruct heap = case H.viewMin heap of + Nothing -> error "cannot view empty heap" + Just (x, xs) -> (x, xs) + +-- | Compute the median weight +median :: Fractional p => MedianHeap p a -> Maybe p +median (MedianHeap lesser greater) + | length lesser > length greater = H.priority . getHD <$> viewHead lesser + | length lesser < length greater = H.priority <$> viewHead greater + | otherwise = do + leftHead <- getHD <$> viewHead lesser + rightHead <- viewHead greater + return $ (H.priority leftHead + H.priority rightHead) / 2 + +getHD :: H.Entry (Down b) c -> H.Entry b c +getHD = first getDown + +#if MIN_VERSION_base(4,14,0) +#else +getDown :: Down a -> a +getDown (Down x) = x +#endif + +viewHead :: H.Heap b -> Maybe b +viewHead h = fst <$> H.viewMin h