diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,7 @@ +## 0.2 + +- Add `Data.TDigest.Vector` module. + ## 0.1 - Add `validateHistogram` and `debugPrint` diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,33 @@ +{-# LANGUAGE CPP #-} +{-# OPTIONS_GHC -Wall #-} +module Main (main) where + +#ifndef MIN_VERSION_cabal_doctest +#define MIN_VERSION_cabal_doctest(x,y,z) 0 +#endif + +#if MIN_VERSION_cabal_doctest(1,0,0) + +import Distribution.Extra.Doctest ( defaultMainWithDoctests ) +main :: IO () +main = defaultMainWithDoctests "doctests" + +#else + +#ifdef MIN_VERSION_Cabal +-- If the macro is defined, we have new cabal-install, +-- but for some reason we don't have cabal-doctest in package-db +-- +-- Probably we are running cabal sdist, when otherwise using new-build +-- workflow +#warning You are configuring this package without cabal-doctest installed. \ + The doctests test-suite will not work as a result. \ + To fix this, install cabal-doctest before configuring. +#endif + +import Distribution.Simple + +main :: IO () +main = defaultMain + +#endif diff --git a/Setup.lhs b/Setup.lhs deleted file mode 100644 --- a/Setup.lhs +++ /dev/null @@ -1,182 +0,0 @@ -\begin{code} -{-# LANGUAGE CPP #-} -{-# LANGUAGE OverloadedStrings #-} -module Main (main) where - -#ifndef MIN_VERSION_cabal_doctest -#define MIN_VERSION_cabal_doctest(x,y,z) 0 -#endif - - -#if MIN_VERSION_cabal_doctest(1,0,0) -import Distribution.Extra.Doctest ( defaultMainWithDoctests ) -#else - --- Otherwise we provide a shim - -#ifndef MIN_VERSION_Cabal -#define MIN_VERSION_Cabal(x,y,z) 0 -#endif -#ifndef MIN_VERSION_directory -#define MIN_VERSION_directory(x,y,z) 0 -#endif -#if MIN_VERSION_Cabal(1,24,0) -#define InstalledPackageId UnitId -#endif - -import Control.Monad ( when ) -import Data.List ( nub ) -import Data.String ( fromString ) -import Distribution.Package ( InstalledPackageId ) -import Distribution.Package ( PackageId, Package (..), packageVersion ) -import Distribution.PackageDescription ( PackageDescription(), TestSuite(..) , Library (..), BuildInfo (..)) -import Distribution.Simple ( defaultMainWithHooks, UserHooks(..), simpleUserHooks ) -import Distribution.Simple.Utils ( rewriteFile, createDirectoryIfMissingVerbose ) -import Distribution.Simple.BuildPaths ( autogenModulesDir ) -import Distribution.Simple.Setup ( BuildFlags(buildDistPref, buildVerbosity), fromFlag) -import Distribution.Simple.LocalBuildInfo ( withPackageDB, withLibLBI, withTestLBI, LocalBuildInfo(), ComponentLocalBuildInfo(componentPackageDeps), compiler ) -import Distribution.Simple.Compiler ( showCompilerId , PackageDB (..)) -import Distribution.Text ( display , simpleParse ) -import System.FilePath ( () ) - -#if MIN_VERSION_Cabal(1,25,0) -import Distribution.Simple.BuildPaths ( autogenComponentModulesDir ) -#endif - -#if MIN_VERSION_directory(1,2,2) -import System.Directory (makeAbsolute) -#else -import System.Directory (getCurrentDirectory) -import System.FilePath (isAbsolute) - -makeAbsolute :: FilePath -> IO FilePath -makeAbsolute p | isAbsolute p = return p - | otherwise = do - cwd <- getCurrentDirectory - return $ cwd p -#endif - -generateBuildModule :: String -> BuildFlags -> PackageDescription -> LocalBuildInfo -> IO () -generateBuildModule testsuiteName flags pkg lbi = do - let verbosity = fromFlag (buildVerbosity flags) - let distPref = fromFlag (buildDistPref flags) - - -- Package DBs - let dbStack = withPackageDB lbi ++ [ SpecificPackageDB $ distPref "package.conf.inplace" ] - let dbFlags = "-hide-all-packages" : packageDbArgs dbStack - - withLibLBI pkg lbi $ \lib libcfg -> do - let libBI = libBuildInfo lib - - -- modules - let modules = exposedModules lib ++ otherModules libBI - -- it seems that doctest is happy to take in module names, not actual files! - let module_sources = modules - - -- We need the directory with library's cabal_macros.h! -#if MIN_VERSION_Cabal(1,25,0) - let libAutogenDir = autogenComponentModulesDir lbi libcfg -#else - let libAutogenDir = autogenModulesDir lbi -#endif - - -- Lib sources and includes - iArgs <- mapM (fmap ("-i"++) . makeAbsolute) $ libAutogenDir : hsSourceDirs libBI - includeArgs <- mapM (fmap ("-I"++) . makeAbsolute) $ includeDirs libBI - - -- CPP includes, i.e. include cabal_macros.h - let cppFlags = map ("-optP"++) $ - [ "-include", libAutogenDir ++ "/cabal_macros.h" ] - ++ cppOptions libBI - - withTestLBI pkg lbi $ \suite suitecfg -> when (testName suite == fromString testsuiteName) $ do - - -- get and create autogen dir -#if MIN_VERSION_Cabal(1,25,0) - let testAutogenDir = autogenComponentModulesDir lbi suitecfg -#else - let testAutogenDir = autogenModulesDir lbi -#endif - createDirectoryIfMissingVerbose verbosity True testAutogenDir - - -- write autogen'd file - rewriteFile (testAutogenDir "Build_doctests.hs") $ unlines - [ "module Build_doctests where" - , "" - -- -package-id etc. flags - , "pkgs :: [String]" - , "pkgs = " ++ (show $ formatDeps $ testDeps libcfg suitecfg) - , "" - , "flags :: [String]" - , "flags = " ++ show (iArgs ++ includeArgs ++ dbFlags ++ cppFlags) - , "" - , "module_sources :: [String]" - , "module_sources = " ++ show (map display module_sources) - ] - where - -- we do this check in Setup, as then doctests don't need to depend on Cabal - isOldCompiler = maybe False id $ do - a <- simpleParse $ showCompilerId $ compiler lbi - b <- simpleParse "7.5" - return $ packageVersion (a :: PackageId) < b - - formatDeps = map formatOne - formatOne (installedPkgId, pkgId) - -- The problem is how different cabal executables handle package databases - -- when doctests depend on the library - | packageId pkg == pkgId = "-package=" ++ display pkgId - | otherwise = "-package-id=" ++ display installedPkgId - - -- From Distribution.Simple.Program.GHC - packageDbArgs :: [PackageDB] -> [String] - packageDbArgs | isOldCompiler = packageDbArgsConf - | otherwise = packageDbArgsDb - - -- GHC <7.6 uses '-package-conf' instead of '-package-db'. - packageDbArgsConf :: [PackageDB] -> [String] - packageDbArgsConf dbstack = case dbstack of - (GlobalPackageDB:UserPackageDB:dbs) -> concatMap specific dbs - (GlobalPackageDB:dbs) -> ("-no-user-package-conf") - : concatMap specific dbs - _ -> ierror - where - specific (SpecificPackageDB db) = [ "-package-conf=" ++ db ] - specific _ = ierror - ierror = error $ "internal error: unexpected package db stack: " - ++ show dbstack - - -- GHC >= 7.6 uses the '-package-db' flag. See - -- https://ghc.haskell.org/trac/ghc/ticket/5977. - packageDbArgsDb :: [PackageDB] -> [String] - -- special cases to make arguments prettier in common scenarios - packageDbArgsDb dbstack = case dbstack of - (GlobalPackageDB:UserPackageDB:dbs) - | all isSpecific dbs -> concatMap single dbs - (GlobalPackageDB:dbs) - | all isSpecific dbs -> "-no-user-package-db" - : concatMap single dbs - dbs -> "-clear-package-db" - : concatMap single dbs - where - single (SpecificPackageDB db) = [ "-package-db=" ++ db ] - single GlobalPackageDB = [ "-global-package-db" ] - single UserPackageDB = [ "-user-package-db" ] - isSpecific (SpecificPackageDB _) = True - isSpecific _ = False - -testDeps :: ComponentLocalBuildInfo -> ComponentLocalBuildInfo -> [(InstalledPackageId, PackageId)] -testDeps xs ys = nub $ componentPackageDeps xs ++ componentPackageDeps ys - -defaultMainWithDoctests :: String -> IO () -defaultMainWithDoctests testSuiteName = defaultMainWithHooks simpleUserHooks - { buildHook = \pkg lbi hooks flags -> do - generateBuildModule testSuiteName flags pkg lbi - buildHook simpleUserHooks pkg lbi hooks flags - } - -#endif - -main :: IO () -main = defaultMainWithDoctests "doctests" - -\end{code} diff --git a/src/Data/TDigest.hs b/src/Data/TDigest.hs --- a/src/Data/TDigest.hs +++ b/src/Data/TDigest.hs @@ -1,127 +1,2 @@ -{-# LANGUAGE ScopedTypeVariables #-} --- | --- A new data structure for accurate on-line accumulation of rank-based --- statistics such as quantiles and trimmed means. --- . --- See original paper: "Computing extremely accurate quantiles using t-digest" --- by Ted Dunning and Otmar Ertl for more details --- . --- --- === Examples --- --- >>> quantile 0.99 (tdigest [1..1000] :: TDigest 25) --- Just 990.5 --- --- >>> quantile 0.99 (tdigest [1..1000] :: TDigest 3) --- Just 989.0... --- --- t-Digest is more precise in tails, especially median is imprecise: --- --- >>> 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, - tdigest, - - -- ** Population - singleton, - insert, - insert', - - -- * Compression - -- - -- | - -- - -- >>> let digest = foldl' (flip insert') mempty [0..1000] :: TDigest 10 - -- >>> (size digest, size $ compress digest) - -- (1001,52) - -- - -- >>> (quantile 0.1 digest, quantile 0.1 $ compress digest) - -- (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) - -- (78,78,48) - -- - -- >>> quantile 0.1 digest - -- Just 98.9... - -- - compress, - forceCompress, - - -- * Statistics - totalWeight, - minimumValue, - maximumValue, - -- ** Histogram - histogram, - HistBin (..), - -- ** Percentile - median, - quantile, - -- ** Mean & Variance - mean, - variance, - stddev, - -- ** CDF - cdf, - icdf, - - -- * Debug - valid, - validate, - debugPrint, - validateHistogram, - ) where - -import Prelude () -import Prelude.Compat - -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) --- >>> let fairshuffle xs = iterate fairshuffle' xs !! 5 +module Data.TDigest (module Data.TDigest.Tree) where +import Data.TDigest.Tree diff --git a/src/Data/TDigest/Internal.hs b/src/Data/TDigest/Internal.hs --- a/src/Data/TDigest/Internal.hs +++ b/src/Data/TDigest/Internal.hs @@ -1,458 +1,22 @@ -{-# 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 +-- Assert ------------------------------------------------------------------------------- --- | 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) +{-# INLINE assert #-} +assert :: Bool -> String -> a -> a +assert _ _ = \x -> x +{- +assert False msg _ = error msg +assert True _ x = x +-} ------------------------------------------------------------------------------- -- Double helpers @@ -466,46 +30,3 @@ 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/NonEmpty.hs b/src/Data/TDigest/NonEmpty.hs --- a/src/Data/TDigest/NonEmpty.hs +++ b/src/Data/TDigest/NonEmpty.hs @@ -1,169 +1,2 @@ -{-# 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 +module Data.TDigest.NonEmpty (module Data.TDigest.Tree.NonEmpty) where +import Data.TDigest.Tree.NonEmpty 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 @@ -1,11 +1,7 @@ --- | 'TDigest' postprocessing functions. --- --- These are re-exported from "Data.TDigest" module. --- module Data.TDigest.Postprocess ( -- * Histogram - histogram, - HistBin (..), + I.HasHistogram (..), + I.HistBin (..), -- * Quantiles median, quantile, @@ -14,187 +10,59 @@ -- | As we have "full" histogram, we can calculate other statistical -- variables. mean, - mean', variance, - variance', + stddev, -- * CDF cdf, icdf, - -- * NonEmpty - histogram', - quantile', - -- * Debug - validateHistogram, + -- * Affine + I.Affine (..) ) where import Prelude () import Prelude.Compat -import Data.Foldable (toList) -import Data.List.NonEmpty (NonEmpty (..), nonEmpty) -import Data.Semigroup (Semigroup (..)) -import Data.Semigroup.Foldable (foldMap1) - -import Data.TDigest.Internal - -------------------------------------------------------------------------------- --- Histogram -------------------------------------------------------------------------------- - --- | Histogram bin -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 -> Maybe (NonEmpty HistBin) -histogram = fmap histogram' . nonEmpty . getCentroids - --- | Histogram from centroids -histogram' :: NonEmpty (Mean,Weight) -> NonEmpty HistBin -histogram' = make - where - make :: NonEmpty (Mean, Weight) -> NonEmpty HistBin - -- one - make ((x, w) :| []) = HistBin x x x w 0 :| [] - -- first - 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 (x0, _) t (c1@(x1, w1) : rest@((x2, _) : _)) - = HistBin (mid x0 x1) (mid x1 x2) x1 w1 t: iter c1 (t + w1) rest - -- last - iter (x0, _) t [(x1, w1)] - = [HistBin (mid x0 x1) x1 x1 w1 t] - - mid a b = (a + b) / 2 +import qualified Data.List.NonEmpty as NE -------------------------------------------------------------------------------- --- Quantile -------------------------------------------------------------------------------- +import qualified Data.TDigest.Postprocess.Internal as I -- | Median, i.e. @'quantile' 0.5@. -median :: TDigest comp -> Maybe Double +median :: I.HasHistogram a f => a -> f Double median = quantile 0.5 -- | Calculate quantile of a specific value. -quantile :: Double -> TDigest comp -> Maybe Double -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 * tw - - 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'. -icdf :: Double -> TDigest comp -> Maybe Double -icdf = quantile - -------------------------------------------------------------------------------- --- Mean -------------------------------------------------------------------------------- +quantile :: I.HasHistogram a f => Double -> a -> f Double +quantile q x = I.quantile q (I.totalWeight x) <$> I.histogram x -- | Mean. -- --- >>> mean (tdigest [1..100] :: TDigest 10) +-- >>> mean (Tree.tdigest [1..100] :: Tree.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 - +mean :: I.HasHistogram a f => a -> f Double +mean x = I.mean <$> I.histogram x -- | 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 +variance :: I.HasHistogram a f => a -> f Double +variance x = I.variance <$> I.histogram x -------------------------------------------------------------------------------- --- CDF - cumulative distribution function -------------------------------------------------------------------------------- +-- | Standard deviation, square root of variance. +stddev :: I.HasHistogram a f => a -> f Double +stddev = fmap sqrt . variance -- | Cumulative distribution function. -- -- /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 $ foldMap toList $ histogram td - where - n = totalWeight td - - iter [] = 1 - 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 +cdf :: I.HasHistogram a f => Double -> a -> Double +cdf q x = I.affine 1 (I.cdf q (I.totalWeight x) . NE.toList) $ I.histogram x +-- | An alias for 'quantile'. +icdf :: I.HasHistogram a f => Double -> a -> f Double +icdf = quantile +-- $setup +-- >>> :set -XDataKinds +-- >>> import qualified Data.TDigest.Tree as Tree diff --git a/src/Data/TDigest/Postprocess/Internal.hs b/src/Data/TDigest/Postprocess/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Data/TDigest/Postprocess/Internal.hs @@ -0,0 +1,209 @@ +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE FunctionalDependencies #-} +{-# LANGUAGE GADTs #-} +-- | 'TDigest' postprocessing functions. +-- +-- These are re-exported from "Data.TDigest" module. +-- +module Data.TDigest.Postprocess.Internal ( + -- * Histogram + HasHistogram (..), + HistBin (..), + histogramFromCentroids, + -- * Quantiles + quantile, + -- * Mean & variance + -- + -- | As we have "full" histogram, we can calculate other statistical + -- variables. + mean, + variance, + -- * CDF + cdf, + -- * Debug + validateHistogram, + -- * Affine - internal + Affine (..), + ) where + +import Data.Foldable (toList) +import Data.Functor.Compose (Compose (..)) +import Data.Functor.Identity (Identity (..)) +import Data.List.NonEmpty (NonEmpty (..), nonEmpty) +import Data.Proxy (Proxy (..)) +import Data.Semigroup (Semigroup (..)) +import Data.Semigroup.Foldable (foldMap1) +import Prelude () +import Prelude.Compat + +import qualified Data.List.NonEmpty as NE + +import Data.TDigest.Internal + +------------------------------------------------------------------------------- +-- Histogram +------------------------------------------------------------------------------- + +-- | Histogram bin +data HistBin = HistBin + { hbMin :: !Mean -- ^ lower bound + , hbMax :: !Mean -- ^ upper bound + , hbValue :: !Mean -- ^ original value: @(mi + ma) / 2@ + , hbWeight :: !Weight -- ^ weight ("area" of the bar) + , hbCumWeight :: !Weight -- ^ weight from the right, excludes this bin + } + deriving (Show) + +-- | Types from which we can extract histogram. +class Affine f => HasHistogram a f | a -> f where + histogram :: a -> f (NonEmpty HistBin) + totalWeight :: a -> Weight + +instance (HistBin ~ e) => HasHistogram (NonEmpty HistBin) Identity where + histogram = Identity + totalWeight = tw . NE.last where + tw hb = hbWeight hb + hbCumWeight hb + +instance (HistBin ~ e) => HasHistogram [HistBin] Maybe where + histogram = nonEmpty + totalWeight = affine 0 totalWeight . histogram + +-- | Histogram from centroids +histogramFromCentroids :: NonEmpty Centroid -> NonEmpty HistBin +histogramFromCentroids = make + where + make :: NonEmpty Centroid -> NonEmpty HistBin + -- one + make ((x, w) :| []) = HistBin x x x w 0 :| [] + -- first + 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 (x0, _) t (c1@(x1, w1) : rest@((x2, _) : _)) + = HistBin (mid x0 x1) (mid x1 x2) x1 w1 t: iter c1 (t + w1) rest + -- last + iter (x0, _) t [(x1, w1)] + = [HistBin (mid x0 x1) x1 x1 w1 t] + + mid a b = (a + b) / 2 + +------------------------------------------------------------------------------- +-- Quantile +------------------------------------------------------------------------------- + +-- | Quantile from the histogram. +quantile :: Double -> Weight -> NonEmpty HistBin -> Double +quantile q tw = iter . toList + where + q' = q * tw + + 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 + +------------------------------------------------------------------------------- +-- Mean +------------------------------------------------------------------------------- + +-- | 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 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 +------------------------------------------------------------------------------- + +-- | Cumulative distribution function. +cdf :: Double + -> Double -- ^ total weight + -> [HistBin] -> Double +cdf x n = iter + where + iter [] = 1 + 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 + +------------------------------------------------------------------------------- +-- Affine +------------------------------------------------------------------------------- + +-- | Affine containers, i.e. containing at most 1 element +-- +-- This class doesn't have 'traverse' analogie +-- as it would require using 'Pointed' which is disputed type class. +-- +-- > traverseAff :: Pointed f => (a -> f b) -> t a -> f (t b) +-- +class Traversable t => Affine t where + -- | Like `foldMap` + affine :: b -> (a -> b) -> t a -> b + affine x f = fromAffine x . fmap f + + fromAffine :: a -> t a -> a + fromAffine x = affine x id + + {-# MINIMAL fromAffine | affine #-} + +instance Affine Identity where fromAffine _ = runIdentity +instance Affine Maybe where affine = maybe +instance Affine Proxy where affine x _ _ = x + +-- | Composition of 'Affine' containers is 'Affine' +instance (Affine f, Affine g) => Affine (Compose f g) where + affine x f (Compose c) = affine x (affine x f) c diff --git a/src/Data/TDigest/Tree.hs b/src/Data/TDigest/Tree.hs new file mode 100644 --- /dev/null +++ b/src/Data/TDigest/Tree.hs @@ -0,0 +1,116 @@ +{-# LANGUAGE ScopedTypeVariables #-} +-- | +-- A new data structure for accurate on-line accumulation of rank-based +-- statistics such as quantiles and trimmed means. +-- . +-- See original paper: "Computing extremely accurate quantiles using t-digest" +-- by Ted Dunning and Otmar Ertl for more details +-- . +-- +-- === Examples +-- +-- >>> quantile 0.99 (tdigest [1..1000] :: TDigest 25) +-- Just 990.5 +-- +-- >>> quantile 0.99 (tdigest [1..1000] :: TDigest 3) +-- Just 989.0... +-- +-- t-Digest is more precise in tails, especially median is imprecise: +-- +-- >>> 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.Tree ( + -- * Construction + TDigest, + tdigest, + + -- ** Population + singleton, + insert, + insert', + + -- * Compression + -- + -- | + -- + -- >>> let digest = foldl' (flip insert') mempty [0..1000] :: TDigest 10 + -- >>> (size digest, size $ compress digest) + -- (1001,52) + -- + -- >>> (quantile 0.1 digest, quantile 0.1 $ compress digest) + -- (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) + -- (78,78,48) + -- + -- >>> quantile 0.1 digest + -- Just 98.9... + -- + compress, + forceCompress, + + -- * Statistics + minimumValue, + maximumValue, + -- ** Percentile + median, + quantile, + -- ** Mean & Variance + -- + -- | + -- -- >>> stddev (tdigest $ fairshuffle [0..100] :: TDigest 10) + -- Just 29.1... + mean, + variance, + stddev, + -- ** CDF + cdf, + icdf, + + -- * Debug + size, + valid, + validate, + debugPrint, + ) where + +import Data.TDigest.Tree.Internal +import Data.TDigest.Tree.Postprocess + +-- $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 (Prelude.Compat.length xs `div` 2) xs) +-- >>> let fairshuffle xs = iterate fairshuffle' xs !! 5 diff --git a/src/Data/TDigest/Tree/Internal.hs b/src/Data/TDigest/Tree/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Data/TDigest/Tree/Internal.hs @@ -0,0 +1,492 @@ +{-# 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.Tree.Internal where + +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.List.NonEmpty (nonEmpty) +import Data.Ord (comparing) +import Data.Proxy (Proxy (..)) +import Data.Semigroup (Semigroup (..)) +import Data.Semigroup.Reducer (Reducer (..)) +import GHC.TypeLits (KnownNat, Nat, natVal) +import Prelude () +import Prelude.Compat + +import qualified Data.Vector.Algorithms.Heap as VHeap +import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector.Unboxed.Mutable as MVU + +import Data.TDigest.Internal +import qualified Data.TDigest.Postprocess.Internal as PP + +------------------------------------------------------------------------------- +-- TDigest +------------------------------------------------------------------------------- + +-- | '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 + +instance PP.HasHistogram (TDigest comp) Maybe where + histogram = fmap PP.histogramFromCentroids . nonEmpty . getCentroids + totalWeight = totalWeight + +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) + +------------------------------------------------------------------------------- +-- 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/Tree/NonEmpty.hs b/src/Data/TDigest/Tree/NonEmpty.hs new file mode 100644 --- /dev/null +++ b/src/Data/TDigest/Tree/NonEmpty.hs @@ -0,0 +1,166 @@ +{-# 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.Tree.NonEmpty ( + -- * Construction + TDigest, + tdigest, + + -- ** Population + singleton, + insert, + insert', + + -- * Compression + compress, + forceCompress, + + -- * Statistics + totalWeight, + minimumValue, + maximumValue, + -- ** 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.Functor.Identity (Identity (..)) +import Data.Semigroup (Semigroup (..)) +import Data.Semigroup.Foldable (Foldable1) +import Data.Semigroup.Reducer (Reducer (..)) +import GHC.TypeLits (KnownNat) + +import Data.TDigest.Internal +import qualified Data.TDigest.Postprocess as PP +import qualified Data.TDigest.Tree.Internal 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 + +instance PP.HasHistogram (TDigest comp) Identity where + histogram = maybe (error "NonEmpty.histogram") Identity . PP.histogram . unEmpty + totalWeight = PP.totalWeight . unEmpty + +------------------------------------------------------------------------------- +-- 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 -> Mean +minimumValue = T.minimumValue . unEmpty + +maximumValue :: TDigest comp -> Mean +maximumValue = T.maximumValue . unEmpty + +totalWeight :: TDigest comp -> Weight +totalWeight = T.totalWeight . unEmpty + +median :: TDigest comp -> Double +median = runIdentity . PP.median + +quantile :: Double -> TDigest comp -> Double +quantile q = runIdentity . PP.quantile q + +mean :: TDigest comp -> Double +mean = runIdentity . PP.mean + +variance :: TDigest comp -> Double +variance = runIdentity . PP.variance + +stddev :: TDigest comp -> Double +stddev = runIdentity . PP.variance + +-- | Alias of 'quantile'. +icdf :: Double -> TDigest comp -> Double +icdf = quantile + +cdf :: Double -> TDigest comp -> Double +cdf = PP.cdf + +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/Tree/Postprocess.hs b/src/Data/TDigest/Tree/Postprocess.hs new file mode 100644 --- /dev/null +++ b/src/Data/TDigest/Tree/Postprocess.hs @@ -0,0 +1,76 @@ +-- | 'TDigest' postprocessing functions. +-- +-- These are re-exported from "Data.TDigest" module. +-- +module Data.TDigest.Tree.Postprocess ( + -- * Quantiles + median, + quantile, + -- * Mean & variance + -- + -- | As we have "full" histogram, we can calculate other statistical + -- variables. + mean, + variance, + stddev, + -- * CDF + cdf, + icdf, + ) where + +import Prelude () +import Prelude.Compat + +import Data.TDigest.Tree.Internal + +import qualified Data.TDigest.Postprocess as PP + +------------------------------------------------------------------------------- +-- Quantile +------------------------------------------------------------------------------- + +-- | Median, i.e. @'quantile' 0.5@. +median :: TDigest comp -> Maybe Double +median = PP.median + +-- | Calculate quantile of a specific value. +quantile :: Double -> TDigest comp -> Maybe Double +quantile = PP.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 = PP.mean + +-- | Variance. +-- +variance :: TDigest comp -> Maybe Double +variance = PP.variance + +-- | Standard deviation, square root of variance. +stddev :: TDigest comp -> Maybe Double +stddev = PP.stddev + +------------------------------------------------------------------------------- +-- CDF - cumulative distribution function +------------------------------------------------------------------------------- + +-- | Cumulative distribution function. +-- +-- /Note:/ if this is the only thing you need, it's more efficient to count +-- this directly. +cdf :: Double -> TDigest comp -> Double +cdf = PP.cdf + +-- | An alias for 'quantile' +icdf :: Double -> TDigest comp -> Maybe Double +icdf = quantile diff --git a/src/Data/TDigest/Vector.hs b/src/Data/TDigest/Vector.hs new file mode 100644 --- /dev/null +++ b/src/Data/TDigest/Vector.hs @@ -0,0 +1,106 @@ +{-# LANGUAGE ScopedTypeVariables #-} +-- | +-- A new data structure for accurate on-line accumulation of rank-based +-- statistics such as quantiles and trimmed means. +-- . +-- See original paper: "Computing extremely accurate quantiles using t-digest" +-- by Ted Dunning and Otmar Ertl for more details +-- . +-- +-- === Examples +-- +-- >>> quantile 0.99 (tdigest [1..1000] :: TDigest 25) +-- Just 990.5 +-- +-- >>> quantile 0.99 (tdigest [1..1000] :: TDigest 3) +-- Just 990.3... +-- +-- t-Digest is more precise in tails, especially median is imprecise: +-- +-- >>> median (forceCompress $ tdigest [1..1000] :: TDigest 10) +-- Just 500.5 +-- +-- === 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 750.5 +-- +-- >>> median ((td [1..500] <> td [501..1000]) <> td [1001..1500]) +-- Just 750.5 +-- +-- 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.5 +-- +-- >>> median ((td' [1..500] <> td' [501..1000]) <> td' [1001..1500]) +-- Just 750.5 +-- +module Data.TDigest.Vector ( + -- * Construction + TDigest, + tdigest, + + -- ** Population + singleton, + insert, + insert', + + -- * Compression + -- + -- | + -- + -- >>> let digest = foldl' (flip insert') mempty [0..1000] :: TDigest 5 + -- >>> (size digest, size $ compress digest) + -- (1001,173) + -- + -- >>> (quantile 0.1 digest, quantile 0.1 $ compress digest) + -- (Just 99.6...,Just 99.6...) + -- + compress, + forceCompress, + + -- * Statistics + minimumValue, + maximumValue, + -- ** Percentile + median, + quantile, + -- ** Mean & Variance + -- + -- | + -- + -- >>> stddev (tdigest $ fairshuffle [0..100] :: TDigest 10) + -- Just 29.0... + mean, + variance, + stddev, + -- ** CDF + cdf, + icdf, + + -- * Debug + size, + valid, + validate, + ) where + +import Data.TDigest.Vector.Internal +import Data.TDigest.Vector.Postprocess + +-- $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 (Prelude.Compat.length xs `div` 2) xs) +-- >>> let fairshuffle xs = iterate fairshuffle' xs !! 5 diff --git a/src/Data/TDigest/Vector/Internal.hs b/src/Data/TDigest/Vector/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Data/TDigest/Vector/Internal.hs @@ -0,0 +1,284 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +module Data.TDigest.Vector.Internal where + +import Control.DeepSeq (NFData (..)) +import Data.Either (isRight) +import Data.Foldable (toList) +import Data.List (foldl', sortBy) +import Data.List.NonEmpty (nonEmpty) +import Data.Ord (comparing) +import Data.Proxy (Proxy (..)) +import Data.Semigroup (Semigroup (..)) +import Data.Semigroup.Reducer (Reducer (..)) +import GHC.TypeLits (KnownNat, Nat, natVal) +import Prelude () +import Prelude.Compat + +import qualified Data.Vector.Unboxed as VU + +import Data.TDigest.Internal +import qualified Data.TDigest.Postprocess.Internal as PP + +-- import Debug.Trace +-- +-- | 'TDigest' is a vector of centroids plus not yet merged elements. +-- +-- The size of structure is dictated by @compression@, *𝛿*. And is *O(𝛿)*. +-- +data TDigest (compression :: Nat) = TDigest + { tdigestTotalWeight :: !Size -- ^ sum of vector and buffer size + , tdigestData :: !(VU.Vector Centroid) -- ^ actual data. *Invariants:* sorted by mean; length <= 2 𝛿 (soft) + , tdigestBufferSize :: !Size + , tdigestBuffer :: [Double] -- ^ addition buffer, elements with weight 1. *Invariants:* length 2 <= 𝛿 + , tdigestDirection :: !Bool -- ^ direction is a hack, so we merge from left and right. *TODO* remove? + } + deriving Show + +instance KnownNat comp => Semigroup (TDigest comp) where + (<>) = combineTDigest + +instance KnownNat comp => Monoid (TDigest comp) where + mempty = emptyTDigest + mappend = (<>) + +-- | Both 'cons' and 'snoc' are 'insert' +instance KnownNat comp => Reducer Double (TDigest comp) where + cons = insert + snoc = flip insert + unit = singleton + +instance NFData (TDigest comp) where + rnf (TDigest _ _ _ b _) = rnf b + +instance KnownNat comp => PP.HasHistogram (TDigest comp) Maybe where + histogram = fmap PP.histogramFromCentroids . nonEmpty . VU.toList . tdigestData . finalize + totalWeight = totalWeight + +-- | Size of structure +size :: TDigest comp -> Int +size td = VU.length (tdigestData td) + tdigestBufferSize td + +totalWeight :: TDigest comp -> Weight +totalWeight = fromIntegral . tdigestTotalWeight + +-- | Center of left-most centroid. Note: may be different than min element inserted. +-- +-- >>> minimumValue (tdigest [1..100] :: TDigest 3) +-- 1.0 +-- +minimumValue :: KnownNat comp => TDigest comp -> Mean +minimumValue td + | VU.null d = posInf + | otherwise = fst (VU.head d) + where + d = tdigestData (finalize td) + +-- | Center of right-most centroid. Note: may be different than max element inserted. +-- +-- >>> maximumValue (tdigest [1..100] :: TDigest 3) +-- 100.0 +-- +maximumValue :: KnownNat comp => TDigest comp -> Mean +maximumValue td + | VU.null d = posInf + | otherwise = fst (VU.last d) + where + d = tdigestData (finalize td) + +------------------------------------------------------------------------------- +-- Mapping function +------------------------------------------------------------------------------- + +-- | Mapping from quantile *q* to notional index *k* with compression parameter *𝛿*. +-- +-- >>> ksize 42 0 +-- 0.0 +-- +-- >>> ksize 42 1 +-- 42.0 +-- +-- *q@ is clamped.: +-- +-- >>> ksize 42 2 +-- 42.0 +-- +ksize + :: Double -- ^ compression parameter, 𝛿 + -> Double -- ^ quantile, q + -> Double -- ^ notional index, k +ksize comp q = comp * (asin (2 * clamp q - 1) / pi + 0.5) + +clamp :: Double -> Double +clamp x + | x < 0.0 = 0.0 + | x > 1.0 = 1.0 + | otherwise = x + +-- | Inverse of 'ksize'. +-- +-- >>> ksizeInv 42 0 +-- 0.0 +-- +-- >>> ksizeInv 42 42 +-- 1.0 +-- +-- >>> ksizeInv 42 (ksize 42 0.3) +-- 0.3 +-- +ksizeInv + :: Double -- ^ compression parameter, 𝛿 + -> Double -- ^ notional index, k + -> Double -- ^ quantile, q +ksizeInv comp k + | k > comp = 1 + | k < 0 = 0 + | otherwise = 0.5 * (sin ((k / comp - 0.5) * pi) + 1) + +------------------------------------------------------------------------------- +-- Merging +------------------------------------------------------------------------------- + +merge :: Int -> Double -> [(Mean, Weight)] -> [(Mean, Weight)] +merge _ _ [] = [] +merge tw' comp (y:ys) = go 0 (qLimit' 0) y ys + where + -- total weight + tw = fromIntegral tw' + + qLimit' :: Double -> Double + qLimit' q0 = ksizeInv comp (ksize comp q0 + 1) -- k⁻¹ (k (q₀, 𝛿) + 1, 𝛿) + + go :: Double -- q0 + -> Double -- qLimit + -> (Mean, Weight) -- sigma + -> [(Mean, Weight)] + -> [(Mean, Weight)] + go _q0 _qLimit sigma [] = [sigma] -- C'.append(σ) + go q0 qLimit sigma (x:xs) + | q <= qLimit = go q0 qLimit (plus sigma x) xs + | otherwise = sigma : go q0' (qLimit' q0') x xs +-- traceShow ("q", sigma, x, q, qLimit) $ + where + q = q0 + (snd sigma + snd x) / tw + q0' = q0 + snd sigma / tw + + plus :: Centroid -> Centroid -> Centroid + plus (m1,w1) (m2,w2) = ((m1 * w1 + m2 * w2) / w, w) where w = w1 + w2 + +------------------------------------------------------------------------------- +-- Implementation +------------------------------------------------------------------------------- + +emptyTDigest :: TDigest comp +emptyTDigest = TDigest 0 mempty 0 mempty True + +combineTDigest :: forall comp. KnownNat comp => TDigest comp -> TDigest comp -> TDigest comp +combineTDigest (TDigest tw d _ b dir) (TDigest tw' d' _ b' dir') = + TDigest (tw + tw') newD 0 [] (dir /= dir') + where + newD = VU.fromList + . merge tw comp + . sortBy (comparing fst) -- sort + $ VU.toList d ++ VU.toList d' ++ map (flip (,) 1) (b ++ b') + + comp = fromInteger (natVal (Proxy :: Proxy comp)) * sizeCoefficient + +-- | Flush insertion buffer +finalize :: forall comp. KnownNat comp => TDigest comp -> TDigest comp +finalize td + | null (tdigestBuffer td) = td + | otherwise = forceCompress td + +forceCompress :: forall comp. KnownNat comp => TDigest comp -> TDigest comp +forceCompress (TDigest tw d _bs b dir) = TDigest tw d' 0 [] (not dir) + where + d' = VU.fromList + . rev + . merge tw comp -- compress + . rev + . sortBy (comparing fst) -- sort + . (++ map (flip (,) 1) b) -- add buffer + . VU.toList + $ d + comp = fromInteger (natVal (Proxy :: Proxy comp)) * sizeCoefficient + rev | dir = id + | otherwise = reverse + +compress :: forall comp. KnownNat comp => TDigest comp -> TDigest comp +compress t@(TDigest _ _ bs _ _) + | bs > compInt * 2 = forceCompress t + | otherwise = t + where + compInt = fromInteger (natVal (Proxy :: Proxy comp)) * sizeCoefficient + +------------------------------------------------------------------------------- +-- Params +------------------------------------------------------------------------------- + +sizeCoefficient :: Num a => a +sizeCoefficient = 32 + +------------------------------------------------------------------------------- +-- Debug +------------------------------------------------------------------------------- + +-- | @'isRight' . 'validate'@ +valid :: TDigest comp -> Bool +valid = isRight . validate + +-- | Check various invariants in the 'TDigest' structure. +validate :: TDigest comp -> Either String (TDigest comp) +validate td@(TDigest tw d bs b _dir) + | not (bs == length b) = + Left $ "Buffer lenght don't match: " ++ show (bs, length b) + | not (tw == bs + round dw) = + Left $ "Total weight doesn't match" + | dl /= sortBy (comparing fst) dl = + Left $ "Data buffer isn't ordered" + | otherwise = Right td + where + dl :: [Centroid] + dl = VU.toList d + + -- total weight of @d@ + dw :: Double + dw = sum (map snd dl) + +------------------------------------------------------------------------------- +-- 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. +-- +-- This may violate the insertion buffer size invariant. +-- +-- 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 (TDigest s d sb b dir) = TDigest (s + 1) d (sb + 1) (x : b) dir + +-- | Make a 'TDigest' of a single data point. +singleton :: Double -> TDigest comp +singleton x = TDigest 1 (VU.singleton (x, 1)) 0 [] True + +-- | Strict 'foldl'' over 'Foldable' structure. +tdigest :: (Foldable f, KnownNat comp) => f Double -> TDigest comp +tdigest = foldl' (flip insert) mempty . toList + +-- $setup +-- >>> :set -XDataKinds diff --git a/src/Data/TDigest/Vector/NonEmpty.hs b/src/Data/TDigest/Vector/NonEmpty.hs new file mode 100644 --- /dev/null +++ b/src/Data/TDigest/Vector/NonEmpty.hs @@ -0,0 +1,1 @@ +module Data.TDigest.Vector.NonEmpty where diff --git a/src/Data/TDigest/Vector/Postprocess.hs b/src/Data/TDigest/Vector/Postprocess.hs new file mode 100644 --- /dev/null +++ b/src/Data/TDigest/Vector/Postprocess.hs @@ -0,0 +1,77 @@ +-- | 'TDigest' postprocessing functions. +-- +-- These are re-exported from "Data.TDigest" module. +-- +module Data.TDigest.Vector.Postprocess ( + -- * Quantiles + median, + quantile, + -- * Mean & variance + -- + -- | As we have "full" histogram, we can calculate other statistical + -- variables. + mean, + variance, + stddev, + -- * CDF + cdf, + icdf, + ) where + +import GHC.TypeLits (KnownNat) +import Prelude () +import Prelude.Compat + +import Data.TDigest.Vector.Internal + +import qualified Data.TDigest.Postprocess as PP + +------------------------------------------------------------------------------- +-- Quantile +------------------------------------------------------------------------------- + +-- | Median, i.e. @'quantile' 0.5@. +median :: KnownNat comp => TDigest comp -> Maybe Double +median = PP.median + +-- | Calculate quantile of a specific value. +quantile :: KnownNat comp => Double -> TDigest comp -> Maybe Double +quantile = PP.quantile + +------------------------------------------------------------------------------- +-- Mean +------------------------------------------------------------------------------- + +-- | Mean. +-- +-- >>> mean (tdigest [1..100] :: TDigest 10) +-- Just 50.5 +-- +-- /Note:/ if you only need the mean, calculate it directly. +-- +mean :: KnownNat comp => TDigest comp -> Maybe Double +mean = PP.mean + +-- | Variance. +-- +variance :: KnownNat comp => TDigest comp -> Maybe Double +variance = PP.variance + +-- | Standard deviation, square root of variance. +stddev :: KnownNat comp => TDigest comp -> Maybe Double +stddev = PP.stddev + +------------------------------------------------------------------------------- +-- CDF - cumulative distribution function +------------------------------------------------------------------------------- + +-- | Cumulative distribution function. +-- +-- /Note:/ if this is the only thing you need, it's more efficient to count +-- this directly. +cdf :: KnownNat comp => Double -> TDigest comp -> Double +cdf = PP.cdf + +-- | An alias for 'quantile' +icdf :: KnownNat comp => Double -> TDigest comp -> Maybe Double +icdf = quantile diff --git a/tdigest.cabal b/tdigest.cabal --- a/tdigest.cabal +++ b/tdigest.cabal @@ -1,20 +1,22 @@ +cabal-version: >= 1.10 name: tdigest -version: 0.1 +version: 0.2 + 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. . See original paper: "Computing extremely accurate quantiles using t-digest" by Ted Dunning and Otmar Ertl - for more details . + for more details . category: Numeric + homepage: https://github.com/futurice/haskell-tdigest#readme bug-reports: https://github.com/futurice/haskell-tdigest/issues author: Oleg Grenrus maintainer: Oleg Grenrus license: BSD3 license-file: LICENSE -tested-with: GHC==7.8.4, GHC==7.10.3, GHC==8.0.1, GHC==8.0.2 +tested-with: GHC==7.8.4, GHC==7.10.3, GHC==8.0.1, GHC==8.0.2, GHC==8.2.2, GHC==8.4.1 build-type: Custom -cabal-version: >= 1.10 extra-source-files: README.md @@ -26,29 +28,51 @@ custom-setup setup-depends: - base >=4.5 && <5, + base >=4.7 && <5, Cabal >=1.14, - cabal-doctest >=1 && <1.1 + cabal-doctest >=1.0.6 && <1.1 library hs-source-dirs: src ghc-options: -Wall + + -- GHC boot libraries build-depends: - base >=4.7 && <4.10 - , base-compat >=0.9.1 && <0.10 + base >=4.7 && <4.12 , deepseq >=1.3.0.2 && <1.5 - , 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 + , binary >=0.7.1.0 && <0.10 + , transformers >=0.3 && <0.6 + + if !impl(ghc >= 8.0) + build-depends: + semigroups >=0.18.4 && <0.19 + + -- other dependencies + build-depends: + base-compat >=0.10.1 && <0.11 + , reducers >=3.12.2 && <3.13 + , semigroupoids >=5.2.2 && <5.3 + , vector >=0.12.0.1 && <0.13 , vector-algorithms >=0.7.0.1 && <0.8 + exposed-modules: Data.TDigest Data.TDigest.NonEmpty Data.TDigest.Postprocess + Data.TDigest.Tree + Data.TDigest.Tree.NonEmpty + Data.TDigest.Tree.Postprocess + Data.TDigest.Vector + Data.TDigest.Vector.NonEmpty + Data.TDigest.Vector.Postprocess + + -- Internal modules are exposed, but aren't under PVP contract. + exposed-modules: Data.TDigest.Internal + Data.TDigest.Postprocess.Internal + Data.TDigest.Tree.Internal + Data.TDigest.Vector.Internal default-language: Haskell2010 other-extensions: DataKinds @@ -72,8 +96,8 @@ semigroups, vector, vector-algorithms, - tasty >=0.11.0.4 && <0.12, - tasty-quickcheck >=0.8.4 && <0.9 + tasty >=0.11.0.4 && <1.1, + tasty-quickcheck >=0.8.4 && <0.11 test-suite doctests default-language: Haskell2010 @@ -84,4 +108,4 @@ build-depends: base, - doctest >=0.11.1 && <0.12 + doctest >=0.11.1 && <0.16 diff --git a/tests/Tests.hs b/tests/Tests.hs --- a/tests/Tests.hs +++ b/tests/Tests.hs @@ -2,6 +2,8 @@ module Main (main) where import Data.TDigest +import Data.TDigest.Postprocess +import Data.TDigest.Postprocess.Internal (validateHistogram) import Test.Tasty import Test.Tasty.QuickCheck