tdigest 0 → 0.1
raw patch · 10 files changed
+883/−862 lines, 10 filesdep +semigroupoidsdep −Chartdep −Chart-diagramsdep −bytesdep ~basedep ~base-compatdep ~binary
Dependencies added: semigroupoids
Dependencies removed: Chart, Chart-diagrams, bytes, directory, filepath, machines, mwc-random, optparse-applicative, parallel, statistics, time
Dependency ranges changed: base, base-compat, binary, deepseq, doctest, semigroups, tasty-quickcheck, vector, vector-algorithms
Files
- CHANGELOG.md +6/−0
- README.md +3/−3
- bench/Simple.hs +0/−296
- src/Data/TDigest.hs +48/−10
- src/Data/TDigest/Internal.hs +511/−0
- src/Data/TDigest/Internal/Tree.hs +0/−496
- src/Data/TDigest/NonEmpty.hs +169/−0
- src/Data/TDigest/Postprocess.hs +125/−23
- tdigest.cabal +6/−31
- tests/Tests.hs +15/−3
+ CHANGELOG.md view
@@ -0,0 +1,6 @@+## 0.1++- Add `validateHistogram` and `debugPrint`+- Fix a pointy centroid bug.+- Add `Data.TDigest.NonEmpty` module+- Add `mean`, `variance`, `stddev`
README.md view
@@ -13,7 +13,6 @@ ## Benchmarks - Using 50M exponentially distributed numbers: - average: **16s**; incorrect approximation of median, mostly to measure prng speed@@ -22,11 +21,12 @@ - buffered t-digest: **68s** - sequential t-digest: **65s** -### Example histogram+## Example histogram ``` tdigest-simple -m tdigest -d standard -s 100000 -c 10 -o output.svg -i 34+cp output.svg example.svg inkscape --export-png=example.png --export-dpi=80 --export-background-opacity=0 --without-gui example.svg ``` -+
− bench/Simple.hs
@@ -1,296 +0,0 @@-{-# LANGUAGE DataKinds #-}-{-# LANGUAGE RankNTypes #-}-{-# LANGUAGE ScopedTypeVariables #-}-module Main (main) where--import Prelude ()-import Prelude.Compat-import Control.Monad (join, replicateM)-import Control.Monad.ST (runST)-import Control.Parallel.Strategies (parList, rseq, using)-import Data.Foldable (for_)-import Data.List (sort)-import Data.Machine-import Data.Machine.Runner (runT1)-import Data.Monoid ((<>))-import Data.Proxy (Proxy (..))-import Data.Time (diffUTCTime, getCurrentTime)-import Data.Word (Word32)-import GHC.TypeLits (KnownNat, SomeNat (..), someNatVal)-import Statistics.Distribution (ContGen (..), density)--import Statistics.Distribution.Exponential (exponential)-import Statistics.Distribution.Gamma (gammaDistr)-import Statistics.Distribution.Normal (standard)-import Statistics.Distribution.Uniform (uniformDistr)--import qualified Data.Vector.Algorithms.Intro as Intro-import qualified Data.Vector.Unboxed as V-import qualified Data.Vector.Unboxed.Mutable as VU-import qualified Options.Applicative as O-import qualified System.Random.MWC as MWC--import qualified Graphics.Rendering.Chart.Backend.Diagrams as Chart-import Graphics.Rendering.Chart.Easy ((&), (.~), (^.))-import qualified Graphics.Rendering.Chart.Easy as Chart--import Data.TDigest------------------------------------------------------------------------------------ Data----------------------------------------------------------------------------------data Method- = MethodAverage- | MethodNaive- | MethodVector- | MethodTDigest- | MethodTDigestBuffered- | MethodTDigestSparking- deriving (Show)--data Distrib- = DistribIncr- | DistribUniform- | DistribExponent- | DistribGamma- | DistribStandard- deriving (Show)--timed :: Show a => IO a -> IO ()-timed mx = do- s <- getCurrentTime- x <- mx- print x- e <- getCurrentTime- print (diffUTCTime e s)--action :: Method -> Distrib -> Int -> Int -> Word32 -> Maybe FilePath -> IO ()-action m d s c iseed fp = do- print (m, d, s, c)- let seed = initSeed (V.singleton iseed)- let dens = case d of- DistribIncr -> density $ uniformDistr 0 (fromIntegral s)- DistribUniform -> density $ uniformDistr 0 1 -- median around 0.5- DistribExponent -> density $ exponential $ log 2 -- median around 1.0- DistribGamma -> density $ gammaDistr 0.1 0.1 -- median around .0000593391- DistribStandard -> density standard- let input = take s $ case d of- DistribIncr -> [1 .. fromIntegral s] -- not sure, but end point prevents floating- DistribUniform -> randomStream (uniformDistr 0 1) seed -- median around 0.5- DistribExponent -> randomStream (exponential $ log 2) seed -- median around 1.0- DistribGamma -> randomStream (gammaDistr 0.1 0.1) seed -- median around .0000593391- DistribStandard -> randomStream standard seed- let method = case m of- MethodAverage -> pure . average- MethodNaive -> pure . naiveMedian- MethodVector -> pure . vectorMedian- MethodTDigest -> reifyNat c $ tdigestMachine fp dens- MethodTDigestBuffered -> reifyNat c $ tdigestBufferedMachine fp dens- MethodTDigestSparking -> reifyNat c $ tdigestSparkingMachine fp dens- timed $ method input--reifyNat :: forall x. Int -> (forall n. KnownNat n => Proxy n -> x) -> x-reifyNat n f = case someNatVal (fromIntegral n) of- Nothing -> error "Negative m"- Just (SomeNat cp) -> f cp--actionParser :: O.Parser (IO ())-actionParser = action- <$> O.option (maybeReader readMethod) (- O.short 'm' <> O.long "method" <> O.metavar ":method" <> O.value MethodTDigestBuffered)- <*> O.option (maybeReader readDistrib) (- O.short 'd' <> O.long "distrib" <> O.metavar ":distrib" <> O.value DistribUniform)- <*> O.option O.auto (- O.short 's' <> O.long "size" <> O.metavar ":size" <> O.value 1000000)- <*> O.option O.auto (- O.short 'c' <> O.long "compression" <> O.metavar ":comp" <> O.value 20)- <*> O.option O.auto (- O.short 'i' <> O.long "seed" <> O.metavar ":seed" <> O.value 42)- <*> O.optional (O.strOption (- O.short 'o' <> O.long "output" <> O.metavar ":output.svg"))- where- readMethod "average" = Just MethodAverage- readMethod "naive" = Just MethodNaive- readMethod "vector" = Just MethodVector- readMethod "digest" = Just MethodTDigest- readMethod "tdigest" = Just MethodTDigest- readMethod "buffered" = Just MethodTDigestBuffered- readMethod "sparking" = Just MethodTDigestSparking- readMethod _ = Nothing-- readDistrib "incr" = Just DistribIncr- readDistrib "uniform" = Just DistribUniform- readDistrib "exponent" = Just DistribExponent- readDistrib "standard" = Just DistribStandard- readDistrib "gamma" = Just DistribGamma- readDistrib _ = Nothing---- Only on optparse-applicative-0.13-maybeReader :: (String -> Maybe a) -> O.ReadM a-maybeReader f = O.eitherReader $ \x -> maybe (Left x) Right (f x)--main :: IO ()-main = join (O.execParser opts)- where- opts = O.info (O.helper <*> actionParser)- (O.fullDesc <> O.header "tdigest-simple - a small utility to explore tdigest")------------------------------------------------------------------------------------ Methods----------------------------------------------------------------------------------average :: [Double] -> Maybe Double-average [] = Nothing-average (x:xs) = Just $ go x 1 xs- where- go z _ [] = z- go z n (y : ys) = go ((z * n + y) / (n + 1)) (n + 1) ys--naiveMedian :: [Double] -> Maybe Double-naiveMedian [] = Nothing-naiveMedian xs = Just $ sort xs !! (length xs `div` 2)--vectorMedian :: [Double] -> Maybe Double-vectorMedian l- | null l = Nothing- | otherwise = runST $ do- let v = V.fromList l- mv <- V.thaw v- Intro.sort mv- Just <$> VU.unsafeRead mv (VU.length mv `div` 2)--tdigestMachine- :: forall comp. KnownNat comp- => Maybe FilePath -> (Double -> Double) -> Proxy comp -> [Double] -> IO (Maybe Double)-tdigestMachine fp dens _ input = do- mdigest <- fmap validate <$> runT1 machine- case mdigest of- Nothing -> return Nothing- Just (Left err) -> fail $ "Validation error: " ++ err- Just (Right digest) -> do- printStats fp dens digest- return $ median digest- where- machine :: MachineT IO k (TDigest comp)- machine- = fold (flip insert) mempty- <~ source input--tdigestBufferedMachine- :: forall comp. KnownNat comp- => Maybe FilePath -> (Double -> Double) -> Proxy comp -> [Double] -> IO (Maybe Double)-tdigestBufferedMachine fp dens _ input = do- mdigest <- fmap validate <$> runT1 machine- case mdigest of- Nothing -> return Nothing- Just (Left err) -> fail $ "Validation error: " ++ err- Just (Right digest) -> do- printStats fp dens digest- return $ median digest- where- machine :: MachineT IO k (TDigest comp)- machine- = fold mappend mempty- <~ mapping tdigest- <~ buffered 10000- <~ source input---- Sparking machine doesn't count-tdigestSparkingMachine- :: forall comp. KnownNat comp- => Maybe FilePath -> (Double -> Double) -> Proxy comp -> [Double] -> IO (Maybe Double)-tdigestSparkingMachine fp dens _ input = do- mdigest <- fmap validate <$> runT1 machine- case mdigest of- Nothing -> return Nothing- Just (Left err) -> fail $ "Validation error: " ++ err- Just (Right digest) -> do- printStats fp dens digest- return $ median digest- where- machine :: MachineT IO k (TDigest comp)- machine- = fold mappend mempty- <~ sparking- <~ mapping tdigest- <~ buffered 10000- <~ source input--printStats :: Maybe FilePath -> (Double -> Double) -> TDigest comp -> IO ()-printStats mfp dens digest = do- -- Extra: print quantiles- putStrLn "quantiles"- for_ ([0.1,0.2..0.9] ++ [0.95,0.99,0.999,0.9999,0.99999]) $ \q ->- putStrLn $ show q ++ ":" ++ show (quantile q digest)- putStrLn "cdf"- for_ ([0, 0.25, 0.5, 1, 2]) $ \x ->- putStrLn $ show x ++ ": " ++ show (cdf x digest)- let mi = minimumValue digest- let ma = maximumValue digest- let points = flip map [0,0.01..1] $ \x -> mi + (ma - mi) * x- for_ mfp $ \fp -> do- putStrLn $ "Writing to " ++ fp- Chart.toFile Chart.def fp $ do- Chart.layout_title Chart..= "Histogram"- color <- Chart.takeColor- let lineStyle = Chart.def- & Chart.line_color .~ color- Chart.plot $ pure $ tdigestToPlot lineStyle digest- Chart.plot $ Chart.line "theoretical" [map (\x -> (x, dens x)) points]--tdigestToPlot :: Chart.LineStyle -> TDigest comp -> Chart.Plot Double Double-tdigestToPlot lineStyle digest = Chart.Plot- { Chart._plot_render = renderHistogram- , Chart._plot_legend = []- , Chart._plot_all_points = unzip allPoints- }- where- hist = histogram digest- allPoints = flip map hist $ \(HistBin mi ma w _) ->- let x = (ma + mi) / 2- d = ma - mi- y = w / d / tw- in (x, y)- tw = totalWeight digest-- renderHistogram pmap = do- let fillColor = Chart.blend 0.5 (Chart.opaque Chart.white) (lineStyle ^. Chart.line_color)- let fillStyle = Chart.def & Chart.fill_color .~ fillColor- Chart.withLineStyle lineStyle $ Chart.withFillStyle fillStyle $- for_ hist $ \(HistBin mi ma w _) -> do- let d = ma - mi- y = w / d / tw- path = Chart.rectPath $ Chart.Rect- (Chart.mapXY pmap (mi,0))- (Chart.mapXY pmap (ma,y))- Chart.alignFillPath path >>= Chart.fillPath- Chart.alignStrokePath path >>= Chart.strokePath------------------------------------------------------------------------------------ Machine additions----------------------------------------------------------------------------------sparking :: Process a a-sparking- = asParts- <~ mapping (\x -> x `using` parList rseq)- <~ buffered 10------------------------------------------------------------------------------------ Statistics additions----------------------------------------------------------------------------------randomStream :: ContGen d => d -> MWC.Seed -> [Double]-randomStream d = go- where- continue (xs, seed) = xs ++ go seed- go seed = continue $ runST $ do- g <- MWC.restore seed- -- Generate first 10000 elements- xs <- replicateM 10000 (genContVar d g)- seed' <- MWC.save g- pure (xs, seed')--initSeed :: V.Vector Word32 -> MWC.Seed-initSeed v = runST $ MWC.initialize v >>= MWC.save
src/Data/TDigest.hs view
@@ -10,16 +10,39 @@ -- === Examples -- -- >>> quantile 0.99 (tdigest [1..1000] :: TDigest 25)--- Just 990.499...+-- Just 990.5 -- -- >>> quantile 0.99 (tdigest [1..1000] :: TDigest 3)--- Just 992.7...+-- Just 989.0... -- -- t-Digest is more precise in tails, especially median is imprecise: ----- >>> median (tdigest [1..1000] :: TDigest 25)--- Just 502.5...+-- >>> median (forceCompress $ tdigest [1..1000] :: TDigest 25)+-- Just 497.6... --+-- === Semigroup+--+-- This operation isn't strictly associative, but statistical+-- variables shouldn't be affected.+--+-- >>> let td xs = tdigest xs :: TDigest 10+--+-- >>> median (td [1..500] <> (td [501..1000] <> td [1001..1500]))+-- Just 802...+--+-- >>> median ((td [1..500] <> td [501..1000]) <> td [1001..1500])+-- Just 726...+--+-- The linear is worst-case scenario:+--+-- >>> let td' xs = tdigest (fairshuffle xs) :: TDigest 10+--+-- >>> median (td' [1..500] <> (td' [501..1000] <> td' [1001..1500]))+-- Just 750.3789...+--+-- >>> median ((td' [1..500] <> td' [501..1000]) <> td' [1001..1500])+-- Just 750.3789...+-- module Data.TDigest ( -- * Construction TDigest,@@ -36,20 +59,20 @@ -- -- >>> let digest = foldl' (flip insert') mempty [0..1000] :: TDigest 10 -- >>> (size digest, size $ compress digest)- -- (1001,54)+ -- (1001,52) -- -- >>> (quantile 0.1 digest, quantile 0.1 $ compress digest)- -- (Just 99.6...,Just 90.1...)+ -- (Just 99.6...,Just 89.7...) -- -- /Note:/ when values are inserted in more random order, -- t-Digest self-compresses on the fly: -- -- >>> let digest = foldl' (flip insert') mempty (fairshuffle [0..1000]) :: TDigest 10 -- >>> (size digest, size $ compress digest, size $ forceCompress digest)- -- (77,77,44)+ -- (78,78,48) -- -- >>> quantile 0.1 digest- -- Just 96.9...+ -- Just 98.9... -- compress, forceCompress,@@ -64,6 +87,10 @@ -- ** Percentile median, quantile,+ -- ** Mean & Variance+ mean,+ variance,+ stddev, -- ** CDF cdf, icdf,@@ -71,18 +98,29 @@ -- * Debug valid, validate,+ debugPrint,+ validateHistogram, ) where import Prelude ()-import Prelude.Compat ()+import Prelude.Compat -import Data.TDigest.Internal.Tree+import Data.TDigest.Internal import Data.TDigest.Postprocess +-- | Standard deviation, square root of variance.+--+-- >>> stddev (tdigest $ fairshuffle [0..100] :: TDigest 10)+-- Just 29.1...+--+stddev :: TDigest comp -> Maybe Double+stddev = fmap sqrt . variance+ -- $setup -- >>> :set -XDataKinds -- >>> import Prelude.Compat -- >>> import Data.List.Compat (foldl')+-- >>> import Data.Semigroup ((<>)) -- -- >>> let merge [] ys = []; merge xs [] = xs; merge (x:xs) (y:ys) = x : y : merge xs ys -- >>> let fairshuffle' xs = uncurry merge (splitAt (length xs `div` 2) xs)
+ src/Data/TDigest/Internal.hs view
@@ -0,0 +1,511 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE KindSignatures #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE ScopedTypeVariables #-}+-- | Internals of 'TDigest'.+--+-- Tree implementation is based on /Adams’ Trees Revisited/ by Milan Straka+-- <http://fox.ucw.cz/papers/bbtree/bbtree.pdf>+module Data.TDigest.Internal where++import Prelude ()+import Prelude.Compat+import Control.DeepSeq (NFData (..))+import Control.Monad.ST (ST, runST)+import Data.Binary (Binary (..))+import Data.Either (isRight)+import Data.Foldable (toList)+import Data.List.Compat (foldl')+import Data.Ord (comparing)+import Data.Proxy (Proxy (..))+import Data.Semigroup (Semigroup (..))+import Data.Semigroup.Reducer (Reducer (..))+import GHC.TypeLits (KnownNat, Nat, natVal)++import qualified Data.Vector.Algorithms.Heap as VHeap+import qualified Data.Vector.Unboxed as VU+import qualified Data.Vector.Unboxed.Mutable as MVU++{-# INLINE assert #-}+assert :: Bool -> String -> a -> a+assert _ _ = \x -> x+{-+assert False msg _ = error msg+assert True _ x = x+-}++-------------------------------------------------------------------------------+-- TDigest+-------------------------------------------------------------------------------++-- TODO: make newtypes+type Mean = Double+type Weight = Double+type Centroid = (Mean, Weight)+type Size = Int++-- | 'TDigest' is a tree of centroids.+--+-- @compression@ is a @1/δ@. The greater the value of @compression@ the less+-- likely value merging will happen.+data TDigest (compression :: Nat)+ -- | Tree node+ = Node+ {-# UNPACK #-} !Size -- size of this tree/centroid+ {-# UNPACK #-} !Mean -- mean of the centroid+ {-# UNPACK #-} !Weight -- weight of the centrod+ {-# UNPACK #-} !Weight -- total weight of the tree+ !(TDigest compression) -- left subtree+ !(TDigest compression) -- right subtree+ -- | Empty tree+ | Nil+ deriving (Show)++-- [Note: keep min & max in the tree]+--+-- We tried it, but it seems the alloc/update cost is bigger than+-- re-calculating them on need (it's O(log n) - calculation!)++-- [Note: singleton node]+-- We tried to add one, but haven't seen change in performance++-- [Note: inlining balanceR and balanceL]+-- We probably can squueze some performance by making+-- 'balanceL' and 'balanceR' check arguments only once (like @containers@ do)+-- and not use 'node' function.+-- *But*, the benefit vs. code explosion is not yet worth.++instance KnownNat comp => Semigroup (TDigest comp) where+ (<>) = combineDigest++-- | Both 'cons' and 'snoc' are 'insert'+instance KnownNat comp => Reducer Double (TDigest comp) where+ cons = insert+ snoc = flip insert+ unit = singleton++instance KnownNat comp => Monoid (TDigest comp) where+ mempty = emptyTDigest+ mappend = combineDigest++-- | 'TDigest' has only strict fields.+instance NFData (TDigest comp) where+ rnf x = x `seq` ()++-- | 'TDigest' isn't compressed after de-serialisation,+-- but it can be still smaller.+instance KnownNat comp => Binary (TDigest comp) where+ put = put . getCentroids+ get = foldl' (flip insertCentroid) emptyTDigest . lc <$> get+ where+ lc :: [Centroid] -> [Centroid]+ lc = id++getCentroids :: TDigest comp -> [Centroid]+getCentroids = ($ []) . go+ where+ go Nil = id+ go (Node _ x w _ l r) = go l . ((x,w) : ) . go r++-- | Total count of samples.+--+-- >>> totalWeight (tdigest [1..100] :: TDigest 5)+-- 100.0+--+totalWeight :: TDigest comp -> Weight+totalWeight Nil = 0+totalWeight (Node _ _ _ tw _ _) = tw++size :: TDigest comp -> Int+size Nil = 0+size (Node s _ _ _ _ _) = s++-- | Center of left-most centroid. Note: may be different than min element inserted.+--+-- >>> minimumValue (tdigest [1..100] :: TDigest 3)+-- 1.0+--+minimumValue :: TDigest comp -> Mean+minimumValue = go posInf+ where+ go acc Nil = acc+ go _acc (Node _ x _ _ l _) = go x l++-- | Center of right-most centroid. Note: may be different than max element inserted.+--+-- >>> maximumValue (tdigest [1..100] :: TDigest 3)+-- 99.0+--+maximumValue :: TDigest comp -> Mean+maximumValue = go negInf+ where+ go acc Nil = acc+ go _acc (Node _ x _ _ _ r) = go x r++-------------------------------------------------------------------------------+-- Implementation+-------------------------------------------------------------------------------++emptyTDigest :: TDigest comp+emptyTDigest = Nil++combineDigest+ :: KnownNat comp+ => TDigest comp+ -> TDigest comp+ -> TDigest comp+combineDigest a Nil = a+combineDigest Nil b = b+combineDigest a@(Node n _ _ _ _ _) b@(Node m _ _ _ _ _)+ -- TODO: merge first, then shuffle and insert (part of compress)+ | n < m = compress $ foldl' (flip insertCentroid) b (getCentroids a)+ | otherwise = compress $ foldl' (flip insertCentroid) a (getCentroids b)++insertCentroid+ :: forall comp. KnownNat comp+ => Centroid+ -> TDigest comp+ -> TDigest comp+insertCentroid (x, w) Nil = singNode x w+insertCentroid (mean, weight) td = go 0 mean weight False td+ where+ -- New weight of the tree+ n :: Weight+ n = totalWeight td + weight++ -- 1/delta+ compression :: Double+ compression = fromInteger $ natVal (Proxy :: Proxy comp)++ go+ :: Weight -- weight to the left of this tree+ -> Mean -- mean to insert+ -> Weight -- weight to insert+ -> Bool -- should insert everything.+ -- if we merged somewhere on top, rest is inserted as is+ -> TDigest comp -- subtree to insert/merge centroid into+ -> TDigest comp+ go _ newX newW _ Nil = singNode newX newW+ go cum newX newW e (Node s x w tw l r) = case compare newX x of+ -- Exact match, insert here+ EQ -> Node s x (w + newW) (tw + newW) l r -- node x (w + newW) l r++ -- there is *no* room to insert into this node+ LT | thr <= w -> balanceL x w (go cum newX newW e l) r+ GT | thr <= w -> balanceR x w l (go (cum + totalWeight l + w) newX newW e r)++ -- otherwise go left ... or later right+ LT | e -> balanceL x w (go cum newX newW e l) r+ LT -> case l of+ -- always create a new node+ Nil -> case mrw of+ Nothing -> node' s nx nw (tw + newW) Nil r+ Just rw -> balanceL nx nw (go cum newX rw True Nil) r+ Node _ _ _ _ _ _+ | lmax < newX && abs (newX - x) < abs (newX - lmax) {- && newX < x -} -> case mrw of+ Nothing -> node' s nx nw (tw + nw - w) l r+ -- in this two last LT cases, we have to recalculate size+ Just rw -> balanceL nx nw (go cum newX rw True l) r+ | otherwise -> balanceL x w (go cum newX newW e l) r+ where+ lmax = maximumValue l++ -- ... or right+ GT | e -> balanceR x w l (go (cum + totalWeight l + w) newX newW True r)+ GT -> case r of+ Nil -> case mrw of+ Nothing -> node' s nx nw (tw + newW) l Nil+ Just rw -> balanceR nx nw l (go (cum + totalWeight l + nw) newX rw True Nil)+ Node _ _ _ _ _ _+ | rmin > newX && abs (newX - x) < abs (newX - rmin) {- && newX > x -} -> case mrw of+ Nothing -> node' s nx nw (tw + newW) l r+ -- in this two last GT cases, we have to recalculate size+ Just rw -> balanceR nx nw l (go (cum + totalWeight l + nw) newX rw True r)+ | otherwise -> balanceR x w l (go (cum + totalWeight l + w) newX newW e r)+ where+ rmin = minimumValue r+ where+ -- quantile approximation of current node+ cum' = cum + totalWeight l+ q = (w / 2 + cum') / n++ -- threshold, max size of current node/centroid+ thr = {- traceShowId $ traceShow (n, q) $ -} threshold n q compression++ -- We later use nx, nw and mrw:++ -- max size of current node+ dw :: Weight+ mrw :: Maybe Weight+ (dw, mrw) =+ let diff = assert (thr > w) "threshold should be larger than current node weight"+ $ w + newW - thr+ in if diff < 0 -- i.e. there is room+ then (newW, Nothing)+ else (thr - w, Just $ diff)++ -- the change of current node+ (nx, nw) = {- traceShowId $ traceShow (newX, newW, x, dw, mrw) $ -} combinedCentroid x w x dw++-- | Constructor which calculates size and total weight.+node :: Mean -> Weight -> TDigest comp -> TDigest comp -> TDigest comp+node x w l r = Node+ (1 + size l + size r)+ x w+ (w + totalWeight l + totalWeight r)+ l r++-- | Balance after right insertion.+balanceR :: Mean -> Weight -> TDigest comp -> TDigest comp -> TDigest comp+balanceR x w l r+ | size l + size r <= 1 = node x w l r+ | size r > balOmega * size l = case r of+ Nil -> error "balanceR: impossible happened"+ (Node _ rx rw _ Nil rr) ->+ -- assert (0 < balAlpha * size rr) "balanceR" $+ -- single left rotation+ node rx rw (node x w l Nil) rr+ (Node _ rx rw _ rl rr)+ | size rl < balAlpha * size rr ->+ -- single left rotation+ node rx rw (node x w l rl) rr+ (Node _ rx rw _ (Node _ rlx rlw _ rll rlr) rr) ->+ -- double left rotation+ node rlx rlw (node x w l rll) (node rx rw rlr rr)+ | otherwise = node x w l r++-- | Balance after left insertion.+balanceL :: Mean -> Weight -> TDigest comp -> TDigest comp -> TDigest comp+balanceL x w l r+ | size l + size r <= 1 = node x w l r+ | size l > balOmega * size r = case l of+ Nil -> error "balanceL: impossible happened"+ (Node _ lx lw _ ll Nil) ->+ -- assert (0 < balAlpha * size ll) "balanceL" $+ -- single right rotation+ node lx lw ll (node x w Nil r)+ (Node _ lx lw _ ll lr)+ | size lr < balAlpha * size ll ->+ -- single right rotation+ node lx lw ll (node x w lr r)+ (Node _ lx lw _ ll (Node _ lrx lrw _ lrl lrr)) ->+ -- double left rotation+ node lrx lrw (node lx lw ll lrl) (node x w lrr r)+ | otherwise = node x w l r++-- | Alias to 'Node'+node' :: Int -> Mean -> Weight -> Weight -> TDigest comp -> TDigest comp -> TDigest comp+node' = Node++-- | Create singular node.+singNode :: Mean -> Weight -> TDigest comp+singNode x w = Node 1 x w w Nil Nil++-- | Add two weighted means together.+combinedCentroid+ :: Mean -> Weight+ -> Mean -> Weight+ -> Centroid+combinedCentroid x w x' w' =+ ( (x * w + x' * w') / w'' -- this is probably not num. stable+ , w''+ )+ where+ w'' = w + w'++-- | Calculate the threshold, i.e. maximum weight of centroid.+threshold+ :: Double -- ^ total weight+ -> Double -- ^ quantile+ -> Double -- ^ compression (1/δ)+ -> Double+threshold n q compression = 4 * n * q * (1 - q) / compression++-------------------------------------------------------------------------------+-- Compression+-------------------------------------------------------------------------------++-- | Compress 'TDigest'.+--+-- Reinsert the centroids in "better" order (in original paper: in random)+-- so they have opportunity to merge.+--+-- Compression will happen only if size is both:+-- bigger than @'relMaxSize' * comp@ and bigger than 'absMaxSize'.+--+compress :: forall comp. KnownNat comp => TDigest comp -> TDigest comp+compress Nil = Nil+compress td+ | size td > relMaxSize * compression && size td > absMaxSize+ = forceCompress td+ | otherwise+ = td+ where+ compression = fromInteger $ natVal (Proxy :: Proxy comp)++-- | Perform compression, even if current size says it's not necessary.+forceCompress :: forall comp. KnownNat comp => TDigest comp -> TDigest comp+forceCompress Nil = Nil+forceCompress td =+ foldl' (flip insertCentroid) emptyTDigest $ fmap fst $ VU.toList centroids+ where+ -- Centroids are shuffled based on space+ centroids :: VU.Vector (Centroid, Double)+ centroids = runST $ do+ v <- toMVector td+ -- sort by cumulative weight+ VHeap.sortBy (comparing snd) v+ f <- VU.unsafeFreeze v+ pure f++toMVector+ :: forall comp s. KnownNat comp+ => TDigest comp -- ^ t-Digest+ -> ST s (VU.MVector s (Centroid, Double)) -- ^ return also a "space left in the centroid" value for "shuffling"+toMVector td = do+ v <- MVU.new (size td)+ (i, cum) <- go v (0 :: Int) (0 :: Double) td+ pure $ assert (i == size td && abs (cum - totalWeight td) < 1e-6) "traversal in toMVector:" v+ where+ go _ i cum Nil = pure (i, cum)+ go v i cum (Node _ x w _ l r) = do+ (i', cum') <- go v i cum l+ MVU.unsafeWrite v i' ((x, w), space w cum')+ go v (i' + 1) (cum' + w) r++ n = totalWeight td+ compression = fromInteger $ natVal (Proxy :: Proxy comp)++ space w cum = thr - w+ where+ q = (w / 2 + cum) / n+ thr = threshold n q compression++-------------------------------------------------------------------------------+-- Params+-------------------------------------------------------------------------------++-- | Relative size parameter. Hard-coded value: 25.+relMaxSize :: Int+relMaxSize = 25++-- | Absolute size parameter. Hard-coded value: 1000.+absMaxSize :: Int+absMaxSize = 1000++-------------------------------------------------------------------------------+-- Tree balance parameters+-------------------------------------------------------------------------------++balOmega :: Int+balOmega = 3++balAlpha :: Int+balAlpha = 2++-- balDelta = 0++-------------------------------------------------------------------------------+-- Debug+-------------------------------------------------------------------------------++-- | Output the 'TDigest' tree.+debugPrint :: TDigest comp -> IO ()+debugPrint td = go 0 td+ where+ go i Nil = putStrLn $ replicate (i * 3) ' ' ++ "Nil"+ go i (Node s m w tw l r) = do+ go (i + 1) l+ putStrLn $ replicate (i * 3) ' ' ++ "Node " ++ show (s,m,w,tw)+ go (i + 1) r++-- | @'isRight' . 'validate'@+valid :: TDigest comp -> Bool+valid = isRight . validate++-- | Check various invariants in the 'TDigest' tree.+validate :: TDigest comp -> Either String (TDigest comp)+validate td+ | not (all sizeValid centroids) = Left "invalid sizes"+ | not (all weightValid centroids) = Left "invalid weights"+ | not (all orderValid centroids) = Left "invalid ordering"+ | not (all balanced centroids) = Left "tree is ill-balanced"+ | otherwise = Right td+ where+ centroids = goc td++ goc Nil = []+ goc n@(Node _ _ _ _ l r) = n : goc l ++ goc r++ sizeValid Nil = True+ sizeValid (Node s _ _ _ l r) = s == size l + size r + 1++ weightValid Nil = True+ weightValid (Node _ _ w tw l r) = eq tw $ w + totalWeight l + totalWeight r++ orderValid Nil = True+ orderValid (Node _ _ _ _ Nil Nil) = True+ orderValid (Node _ x _ _ (Node _ lx _ _ _ _) Nil) = lx < x+ orderValid (Node _ x _ _ Nil (Node _ rx _ _ _ _)) = x < rx+ orderValid (Node _ x _ _ (Node _ lx _ _ _ _) (Node _ rx _ _ _ _)) = lx < x && x < rx++ balanced Nil = True+ balanced (Node _ _ _ _ l r) =+ size l <= max 1 (balOmega * size r) &&+ size r <= max 1 (balOmega * size l)++-------------------------------------------------------------------------------+-- Double helpers+-------------------------------------------------------------------------------++eq :: Double -> Double -> Bool+eq a b = abs (a-b) < 1e-6++negInf :: Double+negInf = negate posInf++posInf :: Double+posInf = 1/0++-------------------------------------------------------------------------------+-- Higher level helpers+-------------------------------------------------------------------------------++-- | Insert single value into 'TDigest'.+insert+ :: KnownNat comp+ => Double -- ^ element+ -> TDigest comp+ -> TDigest comp+insert x = compress . insert' x++-- | Insert single value, don't compress 'TDigest' even if needed.+--+-- For sensibly bounded input, it makes sense to let 'TDigest' grow (it might+-- grow linearly in size), and after that compress it once.+insert'+ :: KnownNat comp+ => Double -- ^ element+ -> TDigest comp+ -> TDigest comp+insert' x = insertCentroid (x, 1)++-- | Make a 'TDigest' of a single data point.+singleton :: KnownNat comp => Double -> TDigest comp+singleton x = insert x emptyTDigest++-- | Strict 'foldl'' over 'Foldable' structure.+tdigest :: (Foldable f, KnownNat comp) => f Double -> TDigest comp+tdigest = foldl' insertChunk emptyTDigest . chunks . toList+ where+ -- compress after each chunk, forceCompress at the very end.+ insertChunk td xs =+ compress (foldl' (flip insert') td xs)++ chunks [] = []+ chunks xs =+ let (a, b) = splitAt 1000 xs -- 1000 is totally arbitrary.+ in a : chunks b++-- $setup+-- >>> :set -XDataKinds
− src/Data/TDigest/Internal/Tree.hs
@@ -1,496 +0,0 @@-{-# LANGUAGE DataKinds #-}-{-# LANGUAGE KindSignatures #-}-{-# LANGUAGE MultiParamTypeClasses #-}-{-# LANGUAGE ScopedTypeVariables #-}--- | Internals of 'TDigest'.------ Tree implementation is based on /Adams’ Trees Revisited/ by Milan Straka--- <http://fox.ucw.cz/papers/bbtree/bbtree.pdf>-module Data.TDigest.Internal.Tree where--import Prelude ()-import Prelude.Compat-import Control.DeepSeq (NFData (..))-import Control.Monad.ST (ST, runST)-import Data.Binary (Binary (..))-import Data.Either (isRight)-import Data.Foldable (toList)-import Data.List.Compat (foldl')-import Data.Ord (comparing)-import Data.Proxy (Proxy (..))-import Data.Semigroup (Semigroup (..))-import Data.Semigroup.Reducer (Reducer (..))-import GHC.TypeLits (KnownNat, Nat, natVal)--import qualified Data.Vector.Algorithms.Heap as VHeap-import qualified Data.Vector.Unboxed as VU-import qualified Data.Vector.Unboxed.Mutable as MVU--assert :: Bool -> String -> a -> a-assert False msg _ = error msg-assert True _ x = x------------------------------------------------------------------------------------ TDigest------------------------------------------------------------------------------------ TODO: make newtypes-type Mean = Double-type Weight = Double-type Centroid = (Mean, Weight)-type Size = Int---- | 'TDigest' is a tree of centroids.------ @compression@ is a @1/δ@. The greater the value of @compression@ the less--- likely value merging will happen.-data TDigest (compression :: Nat)- -- | Tree node- = Node- {-# UNPACK #-} !Size -- size of this tree/centroid- {-# UNPACK #-} !Mean -- mean of the centroid- {-# UNPACK #-} !Weight -- weight of the centrod- {-# UNPACK #-} !Weight -- total weight of the tree- !(TDigest compression) -- left subtree- !(TDigest compression) -- right subtree- -- | Empty tree- | Nil- deriving (Show)---- [Note: keep min & max in the tree]------ We tried it, but it seems the alloc/update cost is bigger than--- re-calculating them on need (it's O(log n) - calculation!)---- [Note: singleton node]--- We tried to add one, but haven't seen change in performance---- [Note: inlining balanceR and balanceL]--- We probably can squueze some performance by making--- 'balanceL' and 'balanceR' check arguments only once (like @containers@ do)--- and not use 'node' function.--- *But*, the benefit vs. code explosion is not yet worth.--instance KnownNat comp => Semigroup (TDigest comp) where- (<>) = combineDigest---- | Both 'cons' and 'snoc' are 'insert'-instance KnownNat comp => Reducer Double (TDigest comp) where- cons = insert- snoc = flip insert- unit = singleton--instance KnownNat comp => Monoid (TDigest comp) where- mempty = emptyTDigest- mappend = combineDigest---- | 'TDigest' has only strict fields.-instance NFData (TDigest comp) where- rnf x = x `seq` ()---- | 'TDigest' isn't compressed after de-serialisation,--- but it can be still smaller.-instance KnownNat comp => Binary (TDigest comp) where- put = put . getCentroids- get = foldl' (flip insertCentroid) emptyTDigest . lc <$> get- where- lc :: [Centroid] -> [Centroid]- lc = id--getCentroids :: TDigest comp -> [Centroid]-getCentroids = ($ []) . go- where- go Nil = id- go (Node _ x w _ l r) = go l . ((x,w) : ) . go r---- | Total count of samples.------ >>> totalWeight (tdigest [1..100] :: TDigest 3)--- 100.0----totalWeight :: TDigest comp -> Double-totalWeight Nil = 0-totalWeight (Node _ _ _ tw _ _) = tw--size :: TDigest comp -> Int-size Nil = 0-size (Node s _ _ _ _ _) = s---- | Center of left-most centroid. Note: may be different than min element inserted.------ >>> minimumValue (tdigest [1..100] :: TDigest 3)--- 1.0----minimumValue :: TDigest comp -> Mean-minimumValue = go posInf- where- go acc Nil = acc- go _acc (Node _ x _ _ l _) = go x l---- | Center of right-most centroid. Note: may be different than max element inserted.------ >>> maximumValue (tdigest [1..100] :: TDigest 3)--- 99.0----maximumValue :: TDigest comp -> Mean-maximumValue = go negInf- where- go acc Nil = acc- go _acc (Node _ x _ _ _ r) = go x r------------------------------------------------------------------------------------ Implementation----------------------------------------------------------------------------------emptyTDigest :: TDigest comp-emptyTDigest = Nil--combineDigest- :: KnownNat comp- => TDigest comp- -> TDigest comp- -> TDigest comp-combineDigest a Nil = a-combineDigest Nil b = b-combineDigest a@(Node n _ _ _ _ _) b@(Node m _ _ _ _ _)- -- TODO: merge first, then shuffle and insert (part of compress)- | n < m = compress $ foldl' (flip insertCentroid) b (getCentroids a)- | otherwise = compress $ foldl' (flip insertCentroid) a (getCentroids b)--insertCentroid- :: forall comp. KnownNat comp- => Centroid- -> TDigest comp- -> TDigest comp-insertCentroid (x, w) Nil = singNode x w-insertCentroid (mean, weight) td = go 0 mean weight False td- where- -- New weight of the tree- n :: Weight- n = totalWeight td + weight-- -- 1/delta- compression :: Double- compression = fromInteger $ natVal (Proxy :: Proxy comp)-- go- :: Weight -- weight to the left of this tree- -> Mean -- mean to insert- -> Weight -- weight to insert- -> Bool -- should insert everything.- -- if we merged somewhere on top, rest is inserted as is- -> TDigest comp -- subtree to insert/merge centroid into- -> TDigest comp- go _ newX newW _ Nil = singNode newX newW- go cum newX newW e (Node s x w tw l r) = case compare newX x of- -- Exact match, insert here- EQ -> Node s x (w + newW) (tw + newW) l r -- node x (w + newW) l r-- -- there is *no* room to insert into this node- LT | thr <= w -> balanceL x w (go cum newX newW e l) r- GT | thr <= w -> balanceR x w l (go (cum + totalWeight l + w) newX newW e r)-- -- otherwise go left ... or later right- LT | e -> balanceL x w (go cum newX newW e l) r- LT -> case l of- -- always create a new node- Nil -> case mrw of- Nothing -> node' s nx nw (tw + newW) Nil r- Just rw -> balanceL nx nw (go cum newX rw True Nil) r- Node _ _ _ _ _ _- | lmax < newX && abs (newX - x) < abs (newX - lmax) {- && newX < x -} -> case mrw of- Nothing -> node' s nx nw (tw + nw - w) l r- -- in this two last LT cases, we have to recalculate size- Just rw -> balanceL nx nw (go cum newX rw True l) r- | otherwise -> balanceL x w (go cum newX newW e l) r- where- lmax = maximumValue l-- -- ... or right- GT | e -> balanceR x w l (go (cum + totalWeight l + w) newX newW True r)- GT -> case r of- Nil -> case mrw of- Nothing -> node' s nx nw (tw + newW) l Nil- Just rw -> balanceR nx nw l (go (cum + totalWeight l + nw) newX rw True Nil)- Node _ _ _ _ _ _- | rmin > newX && abs (newX - x) < abs (newX - rmin) {- && newX > x -} -> case mrw of- Nothing -> node' s nx nw (tw + newW) l r- -- in this two last GT cases, we have to recalculate size- Just rw -> balanceR nx nw l (go (cum + totalWeight l + nw) newX rw True r)- | otherwise -> balanceR x w l (go (cum + totalWeight l + w) newX newW e r)- where- rmin = minimumValue r- where- -- quantile approximation of current node- q = (w / 2 + cum) / n-- -- threshold, max size of current node/centroid- thr = {- traceShowId $ traceShow (n, q) $ -} threshold n q compression-- -- We later use nx, nw and mrw:-- -- max size of current node- dw :: Weight- mrw :: Maybe Weight- (dw, mrw) =- let diff = assert (thr > w) "threshold should be larger than current node weight"- $ w + newW - thr- in if diff < 0 -- i.e. there is room- then (newW, Nothing)- else (thr - w, Just $ diff)-- -- the change of current node- (nx, nw) = {- traceShowId $ traceShow (newX, newW, x, dw, mrw) $ -} combinedCentroid x w x dw---- | Constructor which calculates size and total weight.-node :: Mean -> Weight -> TDigest comp -> TDigest comp -> TDigest comp-node x w l r = Node- (1 + size l + size r)- x w- (w + totalWeight l + totalWeight r)- l r---- | Balance after right insertion.-balanceR :: Mean -> Weight -> TDigest comp -> TDigest comp -> TDigest comp-balanceR x w l r- | size l + size r <= 1 = node x w l r- | size r > balOmega * size l = case r of- Nil -> error "balanceR: impossible happened"- (Node _ rx rw _ Nil rr) ->- -- assert (0 < balAlpha * size rr) "balanceR" $- -- single left rotation- node rx rw (node x w l Nil) rr- (Node _ rx rw _ rl rr)- | size rl < balAlpha * size rr ->- -- single left rotation- node rx rw (node x w l rl) rr- (Node _ rx rw _ (Node _ rlx rlw _ rll rlr) rr) ->- -- double left rotation- node rlx rlw (node x w l rll) (node rx rw rlr rr)- | otherwise = node x w l r---- | Balance after left insertion.-balanceL :: Mean -> Weight -> TDigest comp -> TDigest comp -> TDigest comp-balanceL x w l r- | size l + size r <= 1 = node x w l r- | size l > balOmega * size r = case l of- Nil -> error "balanceL: impossible happened"- (Node _ lx lw _ ll Nil) ->- -- assert (0 < balAlpha * size ll) "balanceL" $- -- single right rotation- node lx lw ll (node x w Nil r)- (Node _ lx lw _ ll lr)- | size lr < balAlpha * size ll ->- -- single right rotation- node lx lw ll (node x w lr r)- (Node _ lx lw _ ll (Node _ lrx lrw _ lrl lrr)) ->- -- double left rotation- node lrx lrw (node lx lw ll lrl) (node x w lrr r)- | otherwise = node x w l r---- | Alias to 'Node'-node' :: Int -> Mean -> Weight -> Weight -> TDigest comp -> TDigest comp -> TDigest comp-node' = Node---- | Create singular node.-singNode :: Mean -> Weight -> TDigest comp-singNode x w = Node 1 x w w Nil Nil---- | Add two weighted means together.-combinedCentroid- :: Mean -> Weight- -> Mean -> Weight- -> Centroid-combinedCentroid x w x' w' =- ( (x * w + x' * w') / w'' -- this is probably not num. stable- , w''- )- where- w'' = w + w'---- | Calculate the threshold, i.e. maximum weight of centroid.-threshold- :: Double -- ^ total weight- -> Double -- ^ quantile- -> Double -- ^ compression (1/δ)- -> Double-threshold n q compression = 4 * n * q * (1 - q) / compression------------------------------------------------------------------------------------ Compression------------------------------------------------------------------------------------ | Compress 'TDigest'.------ Reinsert the centroids in "better" order (in original paper: in random)--- so they have opportunity to merge.------ Compression will happen only if size is both:--- bigger than @'relMaxSize' * comp@ and bigger than 'absMaxSize'.----compress :: forall comp. KnownNat comp => TDigest comp -> TDigest comp-compress Nil = Nil-compress td- | size td > relMaxSize * compression && size td > absMaxSize- = forceCompress td- | otherwise- = td- where- compression = fromInteger $ natVal (Proxy :: Proxy comp)---- | Perform compression, even if current size says it's not necessary.-forceCompress :: forall comp. KnownNat comp => TDigest comp -> TDigest comp-forceCompress Nil = Nil-forceCompress td =- foldl' (flip insertCentroid) emptyTDigest $ fmap fst $ VU.toList centroids- where- -- Centroids are shuffled based on space- centroids :: VU.Vector (Centroid, Double)- centroids = runST $ do- v <- toMVector td- -- sort by cumulative weight- VHeap.sortBy (comparing snd) v- f <- VU.unsafeFreeze v- pure f--toMVector- :: forall comp s. KnownNat comp- => TDigest comp -- ^ t-Digest- -> ST s (VU.MVector s (Centroid, Double)) -- ^ return also a "space left in the centroid" value for "shuffling"-toMVector td = do- v <- MVU.new (size td)- (i, cum) <- go v (0 :: Int) (0 :: Double) td- pure $ assert (i == size td && abs (cum - totalWeight td) < 1e-6) "traversal in toMVector:" v- where- go _ i cum Nil = pure (i, cum)- go v i cum (Node _ x w _ l r) = do- (i', cum') <- go v i cum l- MVU.unsafeWrite v i' ((x, w), space w cum')- go v (i' + 1) (cum' + w) r-- n = totalWeight td- compression = fromInteger $ natVal (Proxy :: Proxy comp)-- space w cum = thr - w- where- q = (w / 2 + cum) / n- thr = threshold n q compression------------------------------------------------------------------------------------ Params------------------------------------------------------------------------------------ | Relative size parameter. Hard-coded value: 25.-relMaxSize :: Int-relMaxSize = 25---- | Absolute size parameter. Hard-coded value: 1000.-absMaxSize :: Int-absMaxSize = 1000------------------------------------------------------------------------------------ Tree balance parameters----------------------------------------------------------------------------------balOmega :: Int-balOmega = 3--balAlpha :: Int-balAlpha = 2---- balDelta = 0------------------------------------------------------------------------------------ Debug------------------------------------------------------------------------------------ | @'isRight' . 'validate'@-valid :: TDigest comp -> Bool-valid = isRight . validate---- | Check various invariants in the 'TDigest' tree.-validate :: TDigest comp -> Either String (TDigest comp)-validate td- | not (all sizeValid centroids) = Left "invalid sizes"- | not (all weightValid centroids) = Left "invalid weights"- | not (all orderValid centroids) = Left "invalid ordering"- | not (all balanced centroids) = Left "tree is ill-balanced"- | otherwise = Right td- where- centroids = goc td-- goc Nil = []- goc n@(Node _ _ _ _ l r) = n : goc l ++ goc r-- sizeValid Nil = True- sizeValid (Node s _ _ _ l r) = s == size l + size r + 1-- weightValid Nil = True- weightValid (Node _ _ w tw l r) = eq tw $ w + totalWeight l + totalWeight r-- orderValid Nil = True- orderValid (Node _ _ _ _ Nil Nil) = True- orderValid (Node _ x _ _ (Node _ lx _ _ _ _) Nil) = lx < x- orderValid (Node _ x _ _ Nil (Node _ rx _ _ _ _)) = x < rx- orderValid (Node _ x _ _ (Node _ lx _ _ _ _) (Node _ rx _ _ _ _)) = lx < x && x < rx-- balanced Nil = True- balanced (Node _ _ _ _ l r) =- size l <= max 1 (balOmega * size r) &&- size r <= max 1 (balOmega * size l)------------------------------------------------------------------------------------ Double helpers----------------------------------------------------------------------------------eq :: Double -> Double -> Bool-eq a b = abs (a-b) < 1e-6--negInf :: Double-negInf = negate posInf--posInf :: Double-posInf = 1/0------------------------------------------------------------------------------------ Higher level helpers------------------------------------------------------------------------------------ | Insert single value into 'TDigest'.-insert- :: KnownNat comp- => Double -- ^ element- -> TDigest comp- -> TDigest comp-insert x = compress . insert' x---- | Insert single value, don't compress 'TDigest' even if needed.------ For sensibly bounded input, it makes sense to let 'TDigest' grow (it might--- grow linearly in size), and after that compress it once.-insert'- :: KnownNat comp- => Double -- ^ element- -> TDigest comp- -> TDigest comp-insert' x = insertCentroid (x, 1)---- | Make a 'TDigest' of a single data point.-singleton :: KnownNat comp => Double -> TDigest comp-singleton x = insert x emptyTDigest---- | Strict 'foldl'' over 'Foldable' structure.-tdigest :: (Foldable f, KnownNat comp) => f Double -> TDigest comp-tdigest = forceCompress . foldl' insertChunk emptyTDigest . chunks . toList- where- -- compress after each chunk, forceCompress at the very end.- insertChunk td xs =- compress (foldl' (flip insert') td xs)-- chunks [] = []- chunks xs =- let (a, b) = splitAt 1000 xs -- 1000 is totally arbitrary.- in a : chunks b---- $setup--- >>> :set -XDataKinds
+ src/Data/TDigest/NonEmpty.hs view
@@ -0,0 +1,169 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE KindSignatures #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE ScopedTypeVariables #-}+-- | This is non empty version of 'Data.TDigest.TDigest', i.e. this is not a 'Monoid',+-- but on the other hand, 'quantile' returns 'Double' not @'Maybe' 'Double'@.+--+-- See "Data.TDigest" for documentation. The exports should be similar,+-- sans non-'Maybe' results.+--+-- === Examples+--+-- >>> quantile 0.99 (tdigest (1 :| [2..1000]) :: TDigest 25)+-- 990.5+--+-- >>> quantile 0.99 (tdigest (1 :| [2..1000]) :: TDigest 3)+-- 989.0...+--+-- t-Digest is more precise in tails, especially median is imprecise:+--+-- >>> median (forceCompress $ tdigest (1 :| [2..1000]) :: TDigest 25)+-- 497.6...+--+module Data.TDigest.NonEmpty (+ -- * Construction+ TDigest,+ tdigest,++ -- ** Population+ singleton,+ insert,+ insert',++ -- * Compression+ compress,+ forceCompress,++ -- * Statistics+ totalWeight,+ minimumValue,+ maximumValue,+ -- ** Histogram+ histogram,+ T.HistBin (..),+ -- ** Percentile+ median,+ quantile,+ -- ** Mean & variance+ mean,+ variance,+ stddev,+ -- ** CDF+ cdf,+ icdf,+ ) where++import Prelude ()+import Prelude.Compat++import Control.DeepSeq (NFData (..))+import Control.Monad (when)+import Data.Binary (Binary (..))+import Data.List.NonEmpty (NonEmpty)+import Data.Maybe (fromMaybe)+import Data.Semigroup (Semigroup (..))+import Data.Semigroup.Foldable (Foldable1)+import Data.Semigroup.Reducer (Reducer (..))+import GHC.TypeLits (KnownNat)++import qualified Data.TDigest as T+import qualified Data.TDigest.Internal as T+import qualified Data.TDigest.Postprocess as T++newtype TDigest comp = TDigest { unEmpty :: T.TDigest comp }++-------------------------------------------------------------------------------+-- Instances+-------------------------------------------------------------------------------++instance NFData (TDigest comp) where+ rnf (TDigest t) = rnf t++instance Show (TDigest comp) where+ showsPrec d (TDigest t) = showsPrec d t++instance KnownNat comp => Semigroup (TDigest comp) where+ TDigest a <> TDigest b = TDigest (a <> b)++instance KnownNat comp => Reducer Double (TDigest comp) where+ cons = insert+ snoc = flip insert+ unit = singleton++instance KnownNat comp => Binary (TDigest comp) where+ get = do+ t <- get+ when (T.size t <= 0) $ fail "empty TDigest.NonEmpty"+ return (TDigest t)++ put (TDigest t) = put t++-------------------------------------------------------------------------------+-- Functions+-------------------------------------------------------------------------------++overTDigest :: (T.TDigest c -> T.TDigest c) -> TDigest c -> TDigest c+overTDigest f = TDigest . f . unEmpty++singleton :: KnownNat comp => Double -> TDigest comp+singleton = TDigest . T.singleton++insert :: KnownNat comp => Double -> TDigest comp -> TDigest comp+insert x = TDigest . T.insert x . unEmpty++insert' :: KnownNat comp => Double -> TDigest comp -> TDigest comp+insert' x = overTDigest $ T.insert' x++compress :: forall comp. KnownNat comp => TDigest comp -> TDigest comp+compress = overTDigest T.compress++forceCompress :: forall comp. KnownNat comp => TDigest comp -> TDigest comp+forceCompress = overTDigest T.forceCompress++minimumValue :: TDigest comp -> T.Mean+minimumValue = T.minimumValue . unEmpty++maximumValue :: TDigest comp -> T.Mean+maximumValue = T.maximumValue . unEmpty++totalWeight :: TDigest comp -> T.Weight+totalWeight = T.totalWeight . unEmpty++histogram :: TDigest comp -> NonEmpty T.HistBin+histogram = fromMaybe (error "NonEmpty.histogram") . T.histogram . unEmpty++median :: TDigest comp -> Double+median = quantile 0.5++quantile :: Double -> TDigest comp -> Double+quantile q td = T.quantile' q (totalWeight td) $ histogram td++mean :: TDigest comp -> Double+mean td = T.mean' $ histogram td++variance :: TDigest comp -> Double+variance td = T.variance' $ histogram td++stddev :: TDigest comp -> Double+stddev = sqrt . variance++-- | Alias of 'quantile'.+icdf :: Double -> TDigest comp -> Double+icdf = quantile++cdf :: Double -> TDigest comp -> Double+cdf x = T.cdf x . unEmpty++tdigest :: (Foldable1 f, KnownNat comp) => f Double -> TDigest comp+tdigest = TDigest . T.tdigest++-- $setup+-- >>> :set -XDataKinds+-- >>> import Prelude.Compat+-- >>> import Data.List.NonEmpty (NonEmpty (..))+-- >>> import Data.List.Compat (foldl')+--+-- >>> let merge [] ys = []; merge xs [] = xs; merge (x:xs) (y:ys) = x : y : merge xs ys+-- >>> let fairshuffle' xs = uncurry merge (splitAt (length xs `div` 2) xs)+-- >>> let fairshuffle xs = iterate fairshuffle' xs !! 5
src/Data/TDigest/Postprocess.hs view
@@ -9,15 +9,33 @@ -- * Quantiles median, quantile,+ -- * Mean & variance+ --+ -- | As we have "full" histogram, we can calculate other statistical+ -- variables.+ mean,+ mean',+ variance,+ variance', -- * CDF cdf, icdf,+ -- * NonEmpty+ histogram',+ quantile',+ -- * Debug+ validateHistogram, ) where import Prelude () import Prelude.Compat-import Data.TDigest.Internal.Tree+import Data.Foldable (toList)+import Data.List.NonEmpty (NonEmpty (..), nonEmpty)+import Data.Semigroup (Semigroup (..))+import Data.Semigroup.Foldable (foldMap1) +import Data.TDigest.Internal+ ------------------------------------------------------------------------------- -- Histogram -------------------------------------------------------------------------------@@ -26,29 +44,36 @@ data HistBin = HistBin { hbMin :: !Double -- ^ lower bound , hbMax :: !Double -- ^ upper bound+ , hbValue :: !Double -- ^ original value: @(mi + ma) / 2@ , hbWeight :: !Double -- ^ weight ("area" of the bar) , hbCumWeight :: !Double -- ^ weight from the right } deriving (Show) -- | Calculate histogram based on the 'TDigest'.-histogram :: TDigest comp -> [HistBin]-histogram = iter Nothing 0 . getCentroids+histogram :: TDigest comp -> Maybe (NonEmpty HistBin)+histogram = fmap histogram' . nonEmpty . getCentroids++-- | Histogram from centroids+histogram' :: NonEmpty (Mean,Weight) -> NonEmpty HistBin+histogram' = make where- -- zero- iter :: Maybe (Mean, Weight) -> Weight -> [(Mean, Weight)] -> [HistBin]- iter _ _ [] = []+ make :: NonEmpty (Mean, Weight) -> NonEmpty HistBin -- one- iter Nothing t [(x, w)] = [HistBin x x w t]+ make ((x, w) :| []) = HistBin x x x w 0 :| [] -- first- iter Nothing t (c1@(x1, w1) : rest@((x2, _) : _))- = HistBin x1 (mid x1 x2) w1 t : iter (Just c1) (t + w1) rest+ make (c1@(x1, w1) :| rest@((x2, _) : _))+ = HistBin x1 (mid x1 x2) x1 w1 0 :| iter c1 w1 rest++ -- zero+ iter :: (Mean, Weight) -> Weight -> [(Mean, Weight)] -> [HistBin]+ iter _ _ [] = [] -- middle- iter (Just (x0, _)) t (c1@(x1, w1) : rest@((x2, _) : _))- = HistBin (mid x0 x1) (mid x1 x2) w1 t: iter (Just c1) (t + w1) rest+ iter (x0, _) t (c1@(x1, w1) : rest@((x2, _) : _))+ = HistBin (mid x0 x1) (mid x1 x2) x1 w1 t: iter c1 (t + w1) rest -- last- iter (Just (x0, _)) t [(x1, w1)]- = [HistBin (mid x0 x1) x1 w1 t]+ iter (x0, _) t [(x1, w1)]+ = [HistBin (mid x0 x1) x1 x1 w1 t] mid a b = (a + b) / 2 @@ -62,15 +87,18 @@ -- | Calculate quantile of a specific value. quantile :: Double -> TDigest comp -> Maybe Double-quantile q td =- iter $ histogram td+quantile q td = quantile' q (totalWeight td) <$> histogram td++-- | Quantile from the histogram.+quantile' :: Double -> Weight -> NonEmpty HistBin -> Double+quantile' q tw = iter . toList where- q' = q * totalWeight td+ q' = q * tw - iter [] = Nothing- iter [HistBin a b w t] = Just $ a + (b - a) * (q' - t) / w- iter (HistBin a b w t : rest)- | {- t < q' && -} q' < t + w = Just $ a + (b - a) * (q' - t) / w+ iter [] = error "quantile: empty NonEmpty"+ iter [HistBin a b _ w t] = a + (b - a) * (q' - t) / w+ iter (HistBin a b _ w t : rest)+ | {- t < q' && -} q' < t + w = a + (b - a) * (q' - t) / w | otherwise = iter rest -- | Alias of 'quantile'.@@ -78,6 +106,62 @@ icdf = quantile -------------------------------------------------------------------------------+-- Mean+-------------------------------------------------------------------------------++-- | Mean.+--+-- >>> mean (tdigest [1..100] :: TDigest 10)+-- Just 50.5+--+-- /Note:/ if you only need the mean, calculate it directly.+--+mean :: TDigest comp -> Maybe Double+mean td = mean' <$> histogram td++-- | Mean from the histogram.+mean' :: NonEmpty HistBin -> Double+mean' = getMean . foldMap1 toMean+ where+ toMean (HistBin _ _ x w _) = Mean w x++data Mean' = Mean !Double !Double++getMean :: Mean' -> Double+getMean (Mean _ x) = x++instance Semigroup Mean' where+ Mean w1 x1 <> Mean w2 x2 = Mean w x+ where+ w = w1 + w2+ x = (x1 * w1 + x2 * w2) / w+++-- | Variance.+--+variance :: TDigest comp -> Maybe Double+variance td = variance' <$> histogram td++-- | Variance from the histogram.+variance' :: NonEmpty HistBin -> Double+variance' = getVariance . foldMap1 toVariance+ where+ toVariance (HistBin _ _ x w _) = Variance w x 0++data Variance = Variance !Double !Double !Double++getVariance :: Variance -> Double+getVariance (Variance w _ d) = d / (w - 1)++-- See: https://izbicki.me/blog/gausian-distributions-are-monoids+instance Semigroup Variance where+ Variance w1 x1 d1 <> Variance w2 x2 d2 = Variance w x d+ where+ w = w1 + w2+ x = (x1 * w1 + x2 * w2) / w+ d = d1 + d2 + w1 * (x1 * x1) + w2 * (x2 * x2) - w * x * x++------------------------------------------------------------------------------- -- CDF - cumulative distribution function ------------------------------------------------------------------------------- @@ -86,13 +170,31 @@ -- /Note:/ if this is the only thing you need, it's more efficient to count -- this directly. cdf :: Double -> TDigest comp -> Double-cdf x td = - iter $ histogram td+cdf x td =+ iter $ foldMap toList $ histogram td where n = totalWeight td iter [] = 1- iter (HistBin a b w t : rest)+ iter (HistBin a b _ w t : rest) | x < a = 0 | x < b = (t + w * (x - a) / (b - a)) / n | otherwise = iter rest++-------------------------------------------------------------------------------+-- Debug+-------------------------------------------------------------------------------++-- | Validate that list of 'HistBin' is a valid "histogram".+validateHistogram :: Foldable f => f HistBin -> Either String (f HistBin)+validateHistogram bs = traverse validPair (pairs $ toList bs) >> pure bs+ where+ validPair (lb@(HistBin _ lmax _ lwt lcw), rb@(HistBin rmin _ _ _ rcw)) = do+ check (lmax == rmin) "gap between bins"+ check (lcw + lwt == rcw) "mismatch in weight cumulation"+ where+ check False err = Left $ err ++ " " ++ show (lb, rb)+ check True _ = Right ()+ pairs xs = zip xs $ tail xs++
tdigest.cabal view
@@ -1,5 +1,5 @@ name: tdigest-version: 0+version: 0.1 synopsis: On-line accumulation of rank-based statistics description: A new data structure for accurate on-line accumulation of rank-based statistics such as quantiles and trimmed means. .@@ -18,6 +18,7 @@ extra-source-files: README.md+ CHANGELOG.md source-repository head type: git@@ -40,12 +41,14 @@ , binary >=0.7.1.0 && <0.9 , reducers >=3.12.1 && <3.13 , semigroups >=0.18.2 && <0.19+ , semigroupoids >=5.1 && <5.2 , vector >=0.11 && <0.13 , vector-algorithms >=0.7.0.1 && <0.8 exposed-modules: Data.TDigest+ Data.TDigest.NonEmpty Data.TDigest.Postprocess- Data.TDigest.Internal.Tree+ Data.TDigest.Internal default-language: Haskell2010 other-extensions: DataKinds@@ -81,32 +84,4 @@ build-depends: base,- bytes,- directory >=1.0,- doctest >=0.11.1 && <0.12,- filepath >=1.2--benchmark tdigest-simple- type: exitcode-stdio-1.0- main-is: Simple.hs- hs-source-dirs:- bench- ghc-options: -Wall -threaded- build-depends:- base- , tdigest- , base-compat- , deepseq- , binary- , semigroups- , vector- , vector-algorithms- , Chart >=1.8.1 && <1.9- , Chart-diagrams >=1.8.1 && <1.9- , machines >=0.6.1 && <0.7- , parallel >=3.2.0.6 && <3.3- , mwc-random >=0.13.4.0 && <0.14- , statistics >=0.13.3.0 && <0.14- , time >=1.4.2 && <1.8- , optparse-applicative >=0.12.1.0 && <0.14- default-language: Haskell2010+ doctest >=0.11.1 && <0.12
tests/Tests.hs view
@@ -10,12 +10,24 @@ tests :: TestTree tests = testGroup "properties"- [ testProperty "valid" propValid+ [ testProperty "tdigest validity" propTDigestIsValid+ , testProperty "histogram validity" propHistogramIsValid ] -propValid :: [Double] -> Property-propValid ds = case validate td of+propTDigestIsValid :: [Double] -> Property+propTDigestIsValid ds = case validate td of Right _ -> property True Left err -> counterexample (err ++ " " ++ show td) (valid td)+ where+ td = tdigest ds :: TDigest 2++propHistogramIsValid :: [Double] -> Property+propHistogramIsValid ds = case fmap validateHistogram $ histogram td of+ Nothing -> property True+ Just (Right _) -> property True+ Just (Left err) -> counterexample msg $ property False+ where+ msg = "Error: " ++ err ++ ", " +++ "TDigest: " ++ show td where td = tdigest ds :: TDigest 2