diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,6 @@ +## 0.1 + +- Add `validateHistogram` and `debugPrint` +- Fix a pointy centroid bug. +- Add `Data.TDigest.NonEmpty` module +- Add `mean`, `variance`, `stddev` diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -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 ``` -![Example](https://raw.githubusercontent.com/futurice/haskell-tdigest/master/example.png) +![Example](https://raw.githubusercontent.com/futurice/haskell-tdigest/master/tdigest/example.png) diff --git a/bench/Simple.hs b/bench/Simple.hs deleted file mode 100644 --- a/bench/Simple.hs +++ /dev/null @@ -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 diff --git a/src/Data/TDigest.hs b/src/Data/TDigest.hs --- a/src/Data/TDigest.hs +++ b/src/Data/TDigest.hs @@ -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) diff --git a/src/Data/TDigest/Internal.hs b/src/Data/TDigest/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Data/TDigest/Internal.hs @@ -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 +-- +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 diff --git a/src/Data/TDigest/Internal/Tree.hs b/src/Data/TDigest/Internal/Tree.hs deleted file mode 100644 --- a/src/Data/TDigest/Internal/Tree.hs +++ /dev/null @@ -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 --- -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 diff --git a/src/Data/TDigest/NonEmpty.hs b/src/Data/TDigest/NonEmpty.hs new file mode 100644 --- /dev/null +++ b/src/Data/TDigest/NonEmpty.hs @@ -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 diff --git a/src/Data/TDigest/Postprocess.hs b/src/Data/TDigest/Postprocess.hs --- a/src/Data/TDigest/Postprocess.hs +++ b/src/Data/TDigest/Postprocess.hs @@ -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 + + diff --git a/tdigest.cabal b/tdigest.cabal --- a/tdigest.cabal +++ b/tdigest.cabal @@ -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 diff --git a/tests/Tests.hs b/tests/Tests.hs --- a/tests/Tests.hs +++ b/tests/Tests.hs @@ -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