diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,10 @@ +# Revision history for dunning-t-digest + +## 0.1.0.0 -- 2025-06-01 + +* Initial release. +* Pure functional t-digest using finger trees (`Data.TDigest.Dunning`). +* Mutable t-digest using mutable vectors in ST (`Data.TDigest.Dunning.Mutable`). +* K1 (arcsine) scale function with O(log n) insertion and queries. +* O(δ log n) split-based compression. +* Freeze/thaw interop between pure and mutable variants. diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,27 @@ +Copyright (c) 2025, Nadia Yvette Chambers + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, + this list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +3. Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from this + software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +POSSIBILITY OF SUCH DAMAGE. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/app/Main.hs b/app/Main.hs new file mode 100644 --- /dev/null +++ b/app/Main.hs @@ -0,0 +1,112 @@ +module Main (main) where + +import Data.List (foldl') +import Data.Sketch.TDigest + +-- --------------------------------------------------------------------------- +-- Formatting helpers +-- --------------------------------------------------------------------------- + +showFFloat6 :: Double -> String +showFFloat6 = showFFloatN 6 + +showFFloat3 :: Double -> String +showFFloat3 = showFFloatN 3 + +showFFloatN :: Int -> Double -> String +showFFloatN n x + | isInfinite x = if x > 0 then "Inf" else "-Inf" + | isNaN x = "NaN" + | x < 0 = "-" ++ showFFloatN n (negate x) + | otherwise = + let factor = 10 ^ n :: Integer + scaled = round (x * fromIntegral factor) :: Integer + wholePart = scaled `div` factor + fracPart = scaled `mod` factor + fracStr = padLeftZ n (show fracPart) + in show wholePart ++ "." ++ fracStr + +padLeftZ :: Int -> String -> String +padLeftZ n s + | length s >= n = s + | otherwise = replicate (n - length s) '0' ++ s + +padRight :: Int -> String -> String +padRight n s + | length s >= n = s + | otherwise = s ++ replicate (n - length s) ' ' + +-- --------------------------------------------------------------------------- +-- Demo / self-test +-- --------------------------------------------------------------------------- + +main :: IO () +main = do + let numValues = 10000 :: Int + values = [fromIntegral i / fromIntegral numValues | i <- [0 .. numValues - 1]] + td = foldl' (flip add) empty values + + putStrLn $ "T-Digest demo: " ++ show numValues ++ " uniform values in [0, 1)" + putStrLn $ "Centroids: " ++ show (centroidCount td) + putStrLn "" + + putStrLn "Quantile estimates (expected ~ q for uniform):" + let qs = [0.001, 0.01, 0.1, 0.25, 0.5, 0.75, 0.9, 0.99, 0.999] :: [Double] + mapM_ + ( \q -> do + let Just est = quantile q td + err = abs (est - q) + putStrLn $ + " q=" + ++ padRight 6 (showFFloat3 q) + ++ " estimated=" + ++ showFFloat6 est + ++ " error=" + ++ showFFloat6 err + ) + qs + + putStrLn "" + + putStrLn "CDF estimates (expected ~ x for uniform):" + let xs = [0.001, 0.01, 0.1, 0.25, 0.5, 0.75, 0.9, 0.99, 0.999] :: [Double] + mapM_ + ( \x -> do + let Just est = cdf x td + err = abs (est - x) + putStrLn $ + " x=" + ++ padRight 6 (showFFloat3 x) + ++ " estimated=" + ++ showFFloat6 est + ++ " error=" + ++ showFFloat6 err + ) + xs + + putStrLn "" + + let vals1 = [fromIntegral i / fromIntegral numValues | i <- [0 .. 4999 :: Int]] + vals2 = [fromIntegral i / fromIntegral numValues | i <- [5000 .. 9999 :: Int]] + td1 = foldl' (flip add) empty vals1 + td2 = foldl' (flip add) empty vals2 + tdM = merge td1 td2 + + putStrLn "After merge of two 5000-element digests:" + case quantile 0.5 tdM of + Just m -> putStrLn $ " median=" ++ showFFloat6 m ++ " (expected ~0.5)" + Nothing -> putStrLn " median=N/A" + case quantile 0.99 tdM of + Just p -> putStrLn $ " p99 =" ++ showFFloat6 p ++ " (expected ~0.99)" + Nothing -> putStrLn " p99 =N/A" + putStrLn $ " centroids=" ++ show (centroidCount tdM) + + putStrLn "" + putStrLn $ + "Merge total weight: " + ++ show (totalWeight tdM) + ++ " (expected " + ++ show (totalWeight td1 + totalWeight td2) + ++ ")" + putStrLn "" + putStrLn "Done." diff --git a/benchmarks/Main.hs b/benchmarks/Main.hs new file mode 100644 --- /dev/null +++ b/benchmarks/Main.hs @@ -0,0 +1,290 @@ +-- Benchmark / asymptotic-behavior tests for the Haskell t-digest implementation. + +module Main where + +import Data.IORef +import Data.Maybe (fromMaybe) +import Data.Sketch.TDigest +import System.CPUTime +import Text.Printf + +-- --------------------------------------------------------------------------- +-- Helpers +-- --------------------------------------------------------------------------- + +getCPUTimeMs :: IO Double +getCPUTimeMs = do + t <- getCPUTime + return (fromIntegral t / 1e9) -- picoseconds -> milliseconds + +timeBlock :: IO a -> IO (Double, a) +timeBlock action = do + t0 <- getCPUTimeMs + result <- action + -- Force evaluation + t1 <- result `seq` getCPUTimeMs + return (t1 - t0, result) + +timeBlock_ :: IO () -> IO Double +timeBlock_ action = do + (ms, _) <- timeBlock action + return ms + +data TestState = TestState {passCount :: !Int, failCount :: !Int} + +newState :: TestState +newState = TestState 0 0 + +addPass :: IORef TestState -> String -> IO () +addPass ref label = do + s <- readIORef ref + writeIORef ref (s {passCount = passCount s + 1}) + printf " %s PASS\n" label + +addFail :: IORef TestState -> String -> IO () +addFail ref label = do + s <- readIORef ref + writeIORef ref (s {failCount = failCount s + 1}) + printf " %s FAIL\n" label + +check :: IORef TestState -> String -> Bool -> IO () +check ref label True = addPass ref label +check ref label False = addFail ref label + +ratioOk :: Double -> Double -> Bool +ratioOk ratio expected = ratio >= expected * 0.5 && ratio <= expected * 3.0 + +ratioOkWide :: Double -> Double -> Bool +ratioOkWide ratio expected = ratio >= expected * 0.2 && ratio <= expected * 5.0 + +-- Build a t-digest from n uniform values +buildDigest :: Double -> Int -> TDigest +buildDigest delta n = + let vals = map (\i -> fromIntegral i / fromIntegral n) [0 .. n - 1] + in foldl' (flip add) (emptyWith delta) vals + +main :: IO () +main = do + ref <- newIORef newState + + putStrLn "=== T-Digest Asymptotic Behavior Tests (Haskell) ===" + putStrLn "" + + -- ----------------------------------------------------------------------- + -- Test 1: add() is amortized O(1) + -- ----------------------------------------------------------------------- + putStrLn "--- Test 1: add() is amortized O(1) ---" + + let sizes = [1000, 10000, 100000, 1000000] :: [Int] + times <- + mapM + ( \n -> do + let go 0 td = td + go i td = go (i - 1) (add (fromIntegral i / fromIntegral n) td) + (ms, _) <- timeBlock (return $! go n (emptyWith 100)) + printf " N=%-9d time=%.1fms\n" n ms + return ms + ) + sizes + + mapM_ + ( \i -> do + let expected = fromIntegral (sizes !! i) / fromIntegral (sizes !! (i - 1)) :: Double + ratio = (times !! i) / (times !! (i - 1)) + check + ref + (printf "N=%d ratio=%.2f (expected ~%.1f)" (sizes !! i) ratio expected) + (ratioOk ratio expected) + ) + [1 .. length sizes - 1] + + putStrLn "" + + -- ----------------------------------------------------------------------- + -- Test 2: Centroid count bounded by O(delta) + -- ----------------------------------------------------------------------- + putStrLn "--- Test 2: Centroid count bounded by O(delta) ---" + + let delta = 100 :: Double + mapM_ + ( \n -> do + let td = buildDigest delta n + cc = centroidCount td + check + ref + (printf "N=%-9d centroids=%-4d (delta=%.0f, limit=%d)" n cc delta (5 * round delta :: Int)) + (cc <= 5 * round delta) + ) + sizes + + putStrLn "" + + -- ----------------------------------------------------------------------- + -- Test 3: Query time independent of N + -- ----------------------------------------------------------------------- + putStrLn "--- Test 3: Query time independent of N ---" + + let querySizes = [1000, 10000, 100000] :: [Int] + queryTimes <- + mapM + ( \n -> do + let td = compress (buildDigest 100 n) + iterations = 10000 :: Int + (ms, _) <- + timeBlock + ( return $! + foldl' + ( \acc _i -> + let q = fromMaybe 0 (quantile 0.5 td) + c = fromMaybe 0 (cdf 0.5 td) + in acc + q + c + ) + (0 :: Double) + [1 .. iterations] + ) + let usPerQuery = (ms * 1000.0) / fromIntegral iterations + printf " N=%-9d query_time=%.2fus\n" n usPerQuery + return usPerQuery + ) + querySizes + + mapM_ + ( \i -> do + let ratio = (queryTimes !! i) / (queryTimes !! (i - 1)) + check + ref + (printf "N=%d ratio=%.2f (expected ~1.0)" (querySizes !! i) ratio) + (ratioOkWide ratio 1.0) + ) + [1 .. length querySizes - 1] + + putStrLn "" + + -- ----------------------------------------------------------------------- + -- Test 4: Tail accuracy improves with delta + -- ----------------------------------------------------------------------- + putStrLn "--- Test 4: Tail accuracy improves with delta ---" + + let deltas = [50, 100, 200] :: [Double] + tailQs = [0.01, 0.001, 0.99, 0.999] :: [Double] + nAcc = 100000 :: Int + + mapM_ + ( \q -> do + errors <- + mapM + ( \d -> do + let td = buildDigest d nAcc + est = fromMaybe 0 (quantile q td) + err = abs (est - q) + printf " delta=%-5.0f q=%-6.3f error=%.6f\n" d q err + return err + ) + deltas + + mapM_ + ( \i -> do + let ok = (errors !! i) <= (errors !! (i - 1)) * 1.5 + 0.001 + check + ref + ( printf + "delta=%.0f q=%.3f error decreases (%.6f <= %.6f)" + (deltas !! i) + q + (errors !! i) + (errors !! (i - 1)) + ) + ok + ) + [1 .. length deltas - 1] + ) + tailQs + + putStrLn "" + + -- ----------------------------------------------------------------------- + -- Test 5: Merge preserves weight and accuracy + -- ----------------------------------------------------------------------- + putStrLn "--- Test 5: Merge preserves weight and accuracy ---" + + let nMerge = 10000 :: Int + td1_0 = + foldl' + (\td i -> add (fromIntegral i / fromIntegral nMerge) td) + (emptyWith 100) + [0 .. nMerge `div` 2 - 1] + td2_0 = + foldl' + (\td i -> add (fromIntegral i / fromIntegral nMerge) td) + (emptyWith 100) + [nMerge `div` 2 .. nMerge - 1] + wBefore = totalWeight td1_0 + totalWeight td2_0 + merged = merge td1_0 td2_0 + wAfter = totalWeight merged + + check + ref + (printf "weight_before=%.0f weight_after=%.0f (equal)" wBefore wAfter) + (abs (wBefore - wAfter) < 1e-9) + + let medianEst = fromMaybe 0 (quantile 0.5 merged) + medianErr = abs (medianEst - 0.5) + check + ref + (printf "median_error=%.6f (< 0.05)" medianErr) + (medianErr < 0.05) + + let p99Est = fromMaybe 0 (quantile 0.99 merged) + p99Err = abs (p99Est - 0.99) + check + ref + (printf "p99_error=%.6f (< 0.05)" p99Err) + (p99Err < 0.05) + + putStrLn "" + + -- ----------------------------------------------------------------------- + -- Test 6: compress is O(n log n) + -- ----------------------------------------------------------------------- + putStrLn "--- Test 6: compress is O(n log n) ---" + + let compressSizes = [500, 5000, 50000] :: [Int] + compressTimes <- + mapM + ( \bufN -> do + let buf = + map + ( \i -> + let v = fromIntegral i / fromIntegral bufN + in v + ) + [0 .. bufN - 1] + td0 = foldl' (flip add) (emptyWith 10000) buf + (ms, _) <- timeBlock (return $! centroidCount (compress td0)) + printf " buf_n=%-8d compress_time=%.2fms\n" bufN ms + return ms + ) + compressSizes + + mapM_ + ( \i -> do + let n0 = fromIntegral (compressSizes !! (i - 1)) :: Double + n1 = fromIntegral (compressSizes !! i) :: Double + expected = (n1 * logBase 2 n1) / (n0 * logBase 2 n0) + ratio = (compressTimes !! i) / (compressTimes !! (i - 1)) + ok = ratio >= expected * 0.3 && ratio <= expected * 4.0 + check + ref + (printf "buf_n=%d ratio=%.2f (expected ~%.1f)" (compressSizes !! i) ratio expected) + ok + ) + [1 .. length compressSizes - 1] + + putStrLn "" + + -- ----------------------------------------------------------------------- + -- Summary + -- ----------------------------------------------------------------------- + s <- readIORef ref + let total = passCount s + failCount s + printf "Summary: %d/%d tests passed\n" (passCount s) total diff --git a/dunning-t-digest.cabal b/dunning-t-digest.cabal new file mode 100644 --- /dev/null +++ b/dunning-t-digest.cabal @@ -0,0 +1,77 @@ +cabal-version: 3.0 +name: dunning-t-digest +version: 0.1.0.0 +synopsis: Dunning t-digest for online quantile estimation +description: + A pure functional implementation of the Dunning t-digest data structure + (merging digest variant, K1 arcsine scale function) using finger trees + with four-component monoidal measures for O(log n) insertion and queries. + . + Also provides a mutable variant backed by mutable vectors in the ST monad. + . + The t-digest provides streaming, mergeable, memory-bounded approximation + of quantile (percentile) queries with high accuracy in the tails. + . + Features: + . + * O(log n) insertion via split-by-mean (no buffering needed) + * O(log n) quantile queries via split-by-cumulative-weight + * O(log n) CDF queries via split-by-mean + * O(δ log n) compression via split-based greedy merge + * O(1) total weight, centroid count, and chunk mean computation + * Mutable variant with O(1) amortized insertion via buffering + +license: BSD-3-Clause +license-file: LICENSE +author: Nadia Yvette Chambers +maintainer: nadia.yvette.chambers@gmail.com +copyright: (c) 2025 Nadia Yvette Chambers +category: Data, Statistics +build-type: Simple +extra-doc-files: CHANGELOG.md +tested-with: GHC == 9.14.1 + +source-repository head + type: git + location: https://github.com/NadiaYvette/t-digest.git + subdir: haskell + +common warnings + ghc-options: -Wall + +library + import: warnings + exposed-modules: + Data.Sketch.TDigest + Data.Sketch.TDigest.Mutable + + build-depends: + , base >= 4.16 && < 5 + , fingertree >= 0.1 && < 0.2 + , vector >= 0.12 && < 0.14 + + hs-source-dirs: src + default-extensions: + MultiParamTypeClasses + RankNTypes + default-language: Haskell2010 + +executable dunning-t-digest-demo + import: warnings + main-is: Main.hs + build-depends: + , base >= 4.16 && < 5 + , dunning-t-digest + + hs-source-dirs: app + default-language: Haskell2010 + +executable dunning-t-digest-bench + import: warnings + main-is: Main.hs + build-depends: + , base >= 4.16 && < 5 + , dunning-t-digest + + hs-source-dirs: benchmarks + default-language: Haskell2010 diff --git a/src/Data/Sketch/TDigest.hs b/src/Data/Sketch/TDigest.hs new file mode 100644 --- /dev/null +++ b/src/Data/Sketch/TDigest.hs @@ -0,0 +1,943 @@ +-- | +-- Module : Data.Sketch.TDigest +-- Description : Dunning t-digest for online quantile estimation +-- Copyright : (c) Nadia Yvette Chambers, 2025 +-- License : BSD-3-Clause +-- Maintainer : nadia.yvette.chambers@gmail.com +-- Stability : experimental +-- +-- A pure functional implementation of the Dunning t-digest data structure, +-- using the merging digest variant with the \(K_1\) (arcsine) scale function. +-- The t-digest provides streaming, mergeable, memory-bounded approximation +-- of quantile (percentile) queries with high accuracy in the tails. +-- +-- == Background +-- +-- The /streaming quantile problem/ asks: given a (possibly unbounded) stream +-- of real-valued observations, answer queries of the form "what is the value +-- at the \(q\)-th quantile?" using bounded memory. +-- Munro & Paterson (1980) established that exact selection from a stream of +-- \(n\) elements requires \(\Omega(n)\) space in the comparison model +-- (), so any sub-linear space +-- algorithm must accept approximation. Greenwald & Khanna (2001) gave the +-- first \(\varepsilon\)-approximate streaming quantile summary with space +-- \(O\!\bigl(\frac{1}{\varepsilon}\log(\varepsilon n)\bigr)\) +-- (), guaranteeing uniform error across +-- all quantiles. The t-digest takes a different approach: it trades uniform +-- guarantees for much higher accuracy in the extreme tails (\(q \approx 0\) or +-- \(q \approx 1\)), which is the regime most relevant to SLA monitoring, +-- anomaly detection, and financial risk measurement. +-- +-- == The t-digest +-- +-- The t-digest, introduced by Ted Dunning +-- (; see also Dunning & Ertl, +-- ), represents an empirical distribution as +-- an ordered sequence of /centroids/ \((m_i, w_i)\), where \(m_i\) is a +-- weighted mean and \(w_i\) is a count of observations. Centroids are kept +-- sorted by mean. The key idea is to use a /scale function/ \(k(q, \delta)\) +-- that maps the quantile axis \([0, 1]\) to a "scale space" in which uniform +-- spacing corresponds to the desired non-uniform resolution in quantile space. +-- +-- This module implements the /merging digest/ variant with the \(K_1\) +-- (arcsine) scale function: +-- +-- \[ +-- k(q, \delta) \;=\; \frac{\delta}{2\pi}\,\arcsin(2q - 1) +-- \] +-- +-- The \(K_1\) function has infinite derivative at \(q = 0\) and \(q = 1\), +-- meaning it allocates proportionally more centroids near the tails. Its +-- inverse is: +-- +-- \[ +-- q(k, \delta) \;=\; \frac{1 + \sin\!\bigl(\frac{2\pi k}{\delta}\bigr)}{2} +-- \] +-- +-- A new observation may be merged into an existing centroid \(i\) only if the +-- resulting centroid would satisfy the /size constraint/: +-- +-- \[ +-- k\!\bigl(q_{\mathrm{upper}},\, \delta\bigr) \;-\; k\!\bigl(q_{\mathrm{lower}},\, \delta\bigr) \;\le\; 1 +-- \] +-- +-- where \(q_{\mathrm{lower}}\) and \(q_{\mathrm{upper}}\) are the quantile +-- boundaries of the (proposed) merged centroid. This constraint ensures that +-- centroids near \(q = 0\) and \(q = 1\) remain small (even singletons), +-- while centroids near the median may absorb many observations. +-- +-- == Space bounds +-- +-- The number of centroids in a t-digest is bounded by \(O(\delta)\) +-- /regardless/ of the number of observations \(n\). Specifically, the integer +-- range of the scale function is +-- \(\lceil k(0,\delta)\rceil \ldots \lfloor k(1,\delta)\rfloor = +-- \lceil -\delta/2\rceil \ldots \lfloor \delta/2\rfloor\), +-- giving at most \(\delta + 1\) unit intervals and therefore at most +-- \(\delta + 1\) centroids after compression. In practice the compression +-- threshold is set to \(3\delta\) centroids (before triggering a compress +-- pass), so the working-set size is at most \(3\delta\) centroids. With the +-- default \(\delta = 100\), this means at most 300 centroids regardless of +-- whether the stream contains \(10^3\) or \(10^{12}\) observations. +-- +-- == Implementation: finger trees with a four-component measure +-- +-- This module stores centroids in a @'Data.FingerTree.FingerTree'@ from the +-- @fingertree@ package, as described by Hinze & Paterson (2006) +-- (). Finger trees support +-- amortised \(O(\log n)\) split and concatenation, and \(O(1)\) access to +-- extremal elements, making them well suited for the sorted-centroid +-- representation. +-- +-- The monoidal measure carried by the tree has four components: +-- +-- 1. @mWeight@ \(= \sum w_i\): cumulative weight, enabling split-by-weight +-- for quantile queries. +-- 2. @mCount@ \(= |\{i\}|\): centroid count, enabling \(O(1)\) +-- 'centroidCount'. +-- 3. @mMaxMean@ \(= \max\{m_i\}\): maximum mean over the subtree, enabling +-- split-by-mean for insertion and CDF queries. Because centroids are +-- stored in sorted order, @mMaxMean@ is monotone over prefixes. +-- 4. @mMeanWeightSum@ \(= \sum m_i w_i\): the sum of products of mean and +-- weight. This enables \(O(1)\) computation of the merged mean of any +-- contiguous chunk: \(\bar{m} = \texttt{mMeanWeightSum} / +-- \texttt{mWeight}\). This is the key to achieving \(O(\delta \log n)\) +-- compression: each of the \(O(\delta)\) chunks produced by splitting at +-- scale-function unit boundaries can be collapsed into a single centroid +-- without traversing its elements. +-- +-- == Companion implementations: array-backed 2-3-4 trees +-- +-- Twenty-two mutable implementations in this project (in C, C++, Rust, Go, +-- Zig, Java, C#, and others) use array-backed 2-3-4 trees instead of +-- finger trees. The 2-3-4 tree is a B-tree of order 4 (Bayer & McCreight, +-- 1972; ), isomorphic to a red-black tree +-- via the correspondence established by Guibas & Sedgewick (1978) +-- (; see also Sedgewick, 2008, +-- ), +-- provides worst-case \(O(\log n)\) insertion, deletion, and search with +-- excellent cache locality when nodes are packed into a flat array. This is +-- particularly important for robustness at very fine-grained quantile queries +-- (e.g., \(q = 0.9999\)) where the tail centroids that determine accuracy +-- must be located quickly and updated with minimal overhead. The array-backed +-- layout avoids pointer-chasing and improves branch-prediction behaviour, +-- yielding 2--5\(\times\) speedups in practice over pointer-based trees. +-- +-- == Quick start +-- +-- @ +-- import Data.Sketch.TDigest +-- import Data.List ('Data.List.foldl\'') +-- +-- main :: IO () +-- main = do +-- let td = 'Data.List.foldl\'' (flip 'add') 'empty' [1.0 .. 10000.0] +-- print ('quantile' 0.99 td) -- Just ~9900.5 +-- print ('cdf' 5000.0 td) -- Just ~0.5 +-- @ +module Data.Sketch.TDigest + ( -- * Types + TDigest, + Centroid (..), + + -- * Construction + empty, + emptyWith, + + -- * Insertion + add, + addWeighted, + + -- * Compression + compress, + + -- * Queries + quantile, + cdf, + + -- * Merging + merge, + + -- * Accessors + totalWeight, + centroidCount, + centroidList, + getDelta, + getMin, + getMax, + + -- * Reconstruction + fromComponents, + ) +where + +import Data.FingerTree (FingerTree, Measured (..), ViewL (..), ViewR (..), (<|), (|>)) +import qualified Data.FingerTree as FT + +-- --------------------------------------------------------------------------- +-- Measure (monoidal annotation for the finger tree) +-- --------------------------------------------------------------------------- + +-- | Monoidal measure carried by every internal node of the finger tree. +-- +-- Following Hinze & Paterson (2006) +-- (), a finger tree is +-- parameterised by a monoid whose cached values enable efficient splitting. +-- The t-digest requires /four/ independent capabilities from the tree, so the +-- measure is a four-component product monoid: +-- +-- * @mWeight@ — cumulative weight \(\sum w_i\). Used by 'quantile' to +-- split the tree at a target cumulative weight in \(O(\log n)\). +-- +-- * @mCount@ — number of centroids \(|\{i\}|\). Provides \(O(1)\) +-- 'centroidCount' and is used during quantile interpolation to detect +-- boundary centroids. +-- +-- * @mMaxMean@ — maximum centroid mean \(\max\{m_i\}\) in the subtree. +-- Because centroids are sorted by mean, this value is monotone over +-- prefixes, enabling 'FT.split' by mean value for insertion ('addWeighted') +-- and CDF queries ('cdf'). +-- +-- * @mMeanWeightSum@ — the sum \(\sum m_i w_i\). Combined with @mWeight@, +-- this allows the weighted mean of any contiguous subtree to be computed in +-- \(O(1)\): \(\bar{m} = \texttt{mMeanWeightSum}\,/\,\texttt{mWeight}\). +-- This is the critical component that makes 'compress' run in +-- \(O(\delta \log n)\) rather than \(O(n)\): each chunk produced by +-- splitting at \(K_1\) unit boundaries is collapsed into a single centroid +-- without iterating over its elements. +data Measure = Measure + { mWeight :: {-# UNPACK #-} !Double, + mCount :: {-# UNPACK #-} !Int, + mMaxMean :: {-# UNPACK #-} !Double, + mMeanWeightSum :: {-# UNPACK #-} !Double + } + deriving (Show) + +instance Semigroup Measure where + (Measure w1 c1 mm1 mws1) <> (Measure w2 c2 mm2 mws2) = + Measure (w1 + w2) (c1 + c2) (max mm1 mm2) (mws1 + mws2) + +instance Monoid Measure where + mempty = Measure 0 0 (-(1 / 0)) 0 + +-- --------------------------------------------------------------------------- +-- Types +-- --------------------------------------------------------------------------- + +-- | A single centroid in the t-digest, representing a cluster of nearby +-- values by their weighted mean and total weight. +-- +-- In the t-digest framework (Dunning, 2021; +-- ), the empirical distribution +-- is approximated by an ordered sequence of centroids \((m_i, w_i)\). When a +-- new observation \(x\) with weight \(w\) is merged into an existing centroid +-- \((m_i, w_i)\), the weighted mean update rule is applied: +-- +-- \[ +-- m_i' \;=\; \frac{m_i \, w_i \;+\; x \, w}{w_i + w}, +-- \qquad +-- w_i' \;=\; w_i + w +-- \] +-- +-- This is the standard incremental weighted mean, which is exact in +-- floating-point arithmetic up to the usual rounding. Note that the centroid +-- does /not/ store individual observations — only the summary statistics +-- \((m_i, w_i)\) are retained, which is what gives the t-digest its bounded +-- space. +data Centroid = Centroid + { -- | Weighted mean of all values merged into this centroid. + cMean :: {-# UNPACK #-} !Double, + -- | Total weight (count) of values in this centroid. For unweighted + -- streams, this is simply the number of observations that have been + -- merged into this centroid. + cWeight :: {-# UNPACK #-} !Double + } + deriving (Show) + +instance Measured Measure Centroid where + measure c = Measure (cWeight c) 1 (cMean c) (cMean c * cWeight c) + +-- | The t-digest data structure for online quantile estimation. +-- +-- Internally, a t'TDigest' consists of: +-- +-- * A 'FingerTree' of t'Centroid's, sorted by mean. The tree carries the +-- four-component @Measure@ described above, enabling \(O(\log n)\) split +-- operations by both mean and cumulative weight. +-- +-- * Cached metadata: the total weight \(N = \sum w_i\), the global minimum +-- and maximum of all observed values, the compression parameter \(\delta\), +-- and the compression threshold \(3\delta\). +-- +-- __Invariants:__ +-- +-- 1. Centroids are sorted in non-decreasing order of 'cMean'. +-- 2. @tdTotalWeight@ equals @mWeight (measure tdCentroids)@ and equals the +-- sum of all 'cWeight' values. +-- 3. @tdMin@ \(\le m_1\) and @tdMax@ \(\ge m_k\) (where \(k\) is the number +-- of centroids), with equality in the singleton case. +-- 4. After 'compress', every centroid satisfies the \(K_1\) size constraint: +-- \(k(q_{\mathrm{upper}}, \delta) - k(q_{\mathrm{lower}}, \delta) \le 1\), +-- where \(q_{\mathrm{lower}}\) and \(q_{\mathrm{upper}}\) are the +-- normalised cumulative weight boundaries of the centroid. +-- 5. The centroid count never exceeds \(3\delta\) for sustained periods; +-- insertions that push the count above this threshold trigger an automatic +-- 'compress' pass. +data TDigest = TDigest + { tdCentroids :: !(FingerTree Measure Centroid), + tdTotalWeight :: !Double, + tdMin :: !Double, + tdMax :: !Double, + tdDelta :: !Double, + tdMaxCentroids :: {-# UNPACK #-} !Int + } + deriving (Show) + +-- --------------------------------------------------------------------------- +-- Construction +-- --------------------------------------------------------------------------- + +-- | Create an empty t-digest with the default compression parameter +-- \(\delta = 100\). +-- +-- This is a good starting point for most applications. With \(\delta = 100\), +-- the digest will use at most 300 centroids (the compression threshold is +-- \(3\delta\)), occupying roughly 4.8 KB of centroid data. Empirically, this +-- yields quantile errors below \(10^{-4}\) at the median and below +-- \(10^{-6}\) for \(q < 0.01\) or \(q > 0.99\) +-- (Dunning & Ertl, 2019; ). +empty :: TDigest +empty = emptyWith 100 + +-- | Create an empty t-digest with a given compression parameter \(\delta\). +-- +-- The compression parameter controls the trade-off between accuracy and space: +-- +-- * __Larger \(\delta\)__ (e.g., 200–500) means more centroids are retained, +-- giving higher accuracy — especially at extreme quantiles — at the cost of +-- more memory and slower queries. +-- * __Smaller \(\delta\)__ (e.g., 20–50) means fewer centroids, saving +-- memory but increasing quantile estimation error. +-- +-- The maximum number of centroids after compression is \(\delta + 1\) +-- (one per integer unit in the range of \(K_1\)), and the compression +-- threshold (the point at which automatic compression is triggered during +-- insertion) is set to \(\lceil 3\delta \rceil\). Typical values used in +-- production systems are \(\delta \in [50, 300]\). +-- +-- Setting \(\delta \le 0\) is not meaningful and will result in a digest that +-- compresses aggressively to zero or one centroid. +emptyWith :: Double -> TDigest +emptyWith delta = + TDigest + { tdCentroids = FT.empty, + tdTotalWeight = 0, + tdMin = 1 / 0, + tdMax = -(1 / 0), + tdDelta = delta, + tdMaxCentroids = ceiling (delta * 3) + } + +-- --------------------------------------------------------------------------- +-- Scale function K_1 +-- --------------------------------------------------------------------------- + +-- | The \(K_1\) (arcsine) scale function: +-- +-- \[ +-- k(q, \delta) \;=\; \frac{\delta}{2\pi}\,\arcsin(2q - 1) +-- \] +-- +-- This function maps the quantile domain \([0, 1]\) to the "scale space" +-- \([-\delta/2,\; \delta/2]\). Its derivative +-- \(k'(q) = \frac{\delta}{\pi\sqrt{q(1-q)}}\) diverges at \(q = 0\) and +-- \(q = 1\), causing centroids near the tails to be allocated much more +-- finely than centroids near the median — which is the defining feature of +-- the t-digest's accuracy profile. +kScale :: Double -> Double -> Double +kScale delta q = (delta / (2 * pi)) * asin (2 * q - 1) + +-- | Inverse of the \(K_1\) scale function: +-- +-- \[ +-- q(k, \delta) \;=\; \frac{1 + \sin\!\bigl(\frac{2\pi k}{\delta}\bigr)}{2} +-- \] +-- +-- Used during 'compress' to compute the quantile boundaries corresponding to +-- integer scale-function values, i.e., the boundaries of the unit intervals +-- in scale space. +kScaleInv :: Double -> Double -> Double +kScaleInv delta k = (1 + sin (2 * pi * k / delta)) / 2 + +-- --------------------------------------------------------------------------- +-- FingerTree helpers +-- --------------------------------------------------------------------------- + +ftToList :: FingerTree Measure Centroid -> [Centroid] +ftToList ft = case FT.viewl ft of + EmptyL -> [] + x :< rest -> x : ftToList rest + +splitByMean :: Double -> FingerTree Measure Centroid -> (FingerTree Measure Centroid, FingerTree Measure Centroid) +splitByMean x = FT.split (\m -> mMaxMean m >= x) + +-- --------------------------------------------------------------------------- +-- Adding values +-- --------------------------------------------------------------------------- + +-- | Add a single value with weight 1 to the digest. +-- +-- \(O(\log n)\) amortised, where \(n\) is the number of centroids. +-- Equivalent to @'addWeighted' x 1@. +add :: Double -> TDigest -> TDigest +add x = addWeighted x 1 + +-- | Add a value \(x\) with a given weight \(w\) to the digest. +-- +-- The algorithm proceeds as follows: +-- +-- 1. __Split__ the finger tree at the insertion point using +-- @'FT.split' (\m -> mMaxMean m >= x)@, yielding a left subtree (all +-- centroids with mean \(< x\)) and a right subtree (mean \(\ge x\)). +-- This is \(O(\log n)\) by the finger tree split theorem +-- (Hinze & Paterson, 2006; ). +-- +-- 2. __Find nearest neighbour:__ examine the rightmost centroid of the left +-- subtree and the leftmost centroid of the right subtree. For each +-- candidate neighbour \((m_i, w_i)\), compute the proposed merged weight +-- \(w_i + w\) and check the \(K_1\) scale-function constraint: +-- +-- \[ +-- k\!\bigl(q_{\mathrm{upper}},\, \delta\bigr) +-- \;-\; k\!\bigl(q_{\mathrm{lower}},\, \delta\bigr) +-- \;\le\; 1 +-- \] +-- +-- where \(q_{\mathrm{lower}}\) and \(q_{\mathrm{upper}}\) are the +-- normalised cumulative weight boundaries of the proposed merged centroid. +-- +-- 3. __Merge or insert:__ if one or both neighbours can absorb the new value, +-- merge with the /closer/ one (by distance \(|m_i - x|\)) using the +-- weighted mean update rule. If neither can absorb it (because doing so +-- would violate the size constraint), insert a new singleton centroid +-- \((x, w)\) into the tree. +-- +-- 4. __Auto-compress:__ if the centroid count exceeds the threshold +-- \(3\delta\), trigger a 'compress' pass. +-- +-- The overall amortised cost is \(O(\log n)\), dominated by the finger tree +-- split and concatenation. +addWeighted :: Double -> Double -> TDigest -> TDigest +addWeighted x w td = + let n = tdTotalWeight td + w + newMin = min x (tdMin td) + newMax = max x (tdMax td) + delta = tdDelta td + cs = tdCentroids td + newC = Centroid x w + td' = + if FT.null cs + then + td + { tdCentroids = FT.singleton newC, + tdTotalWeight = n, + tdMin = newMin, + tdMax = newMax + } + else + let (left, right) = splitByMean x cs + leftWeight = mWeight (FT.measure left) + result = tryMergeNeighbor delta n leftWeight left right newC + in td + { tdCentroids = result, + tdTotalWeight = n, + tdMin = newMin, + tdMax = newMax + } + in if mCount (FT.measure (tdCentroids td')) > tdMaxCentroids td' + then compress td' + else td' + +-- | Try to merge with nearest neighbor; insert if neither allows merging. +tryMergeNeighbor :: + Double -> + Double -> + Double -> + FingerTree Measure Centroid -> + FingerTree Measure Centroid -> + Centroid -> + FingerTree Measure Centroid +tryMergeNeighbor delta n leftWeight left right newC = + let x = cMean newC + k = kScale delta + + leftNeighbor = case FT.viewr left of + EmptyR -> Nothing + leftRest :> lc -> + let cumBefore = mWeight (FT.measure leftRest) + proposed = cWeight lc + cWeight newC + q0 = cumBefore / n + q1 = (cumBefore + proposed) / n + canMerge = k q1 - k q0 <= 1.0 + dist = abs (cMean lc - x) + in if canMerge then Just (leftRest, lc, dist) else Nothing + + rightNeighbor = case FT.viewl right of + EmptyL -> Nothing + rc :< rightRest -> + let proposed = cWeight rc + cWeight newC + q0 = leftWeight / n + q1 = (leftWeight + proposed) / n + canMerge = k q1 - k q0 <= 1.0 + dist = abs (cMean rc - x) + in if canMerge then Just (rightRest, rc, dist) else Nothing + in case (leftNeighbor, rightNeighbor) of + (Just (leftRest, lc, ldist), Just (rightRest, rc, rdist)) + | ldist <= rdist -> + (leftRest |> mergeCentroid lc newC) FT.>< right + | otherwise -> + left FT.>< (mergeCentroid rc newC <| rightRest) + (Just (leftRest, lc, _), Nothing) -> + (leftRest |> mergeCentroid lc newC) FT.>< right + (Nothing, Just (rightRest, rc, _)) -> + left FT.>< (mergeCentroid rc newC <| rightRest) + (Nothing, Nothing) -> + left FT.>< (newC <| right) + +-- | Merge two centroids using the weighted mean update rule: +-- +-- \[ +-- m' = \frac{m_a \, w_a + m_b \, w_b}{w_a + w_b}, +-- \qquad +-- w' = w_a + w_b +-- \] +mergeCentroid :: Centroid -> Centroid -> Centroid +mergeCentroid a b = + let w = cWeight a + cWeight b + m = (cMean a * cWeight a + cMean b * cWeight b) / w + in Centroid m w + +-- --------------------------------------------------------------------------- +-- Compression (split-based greedy merge) +-- --------------------------------------------------------------------------- + +-- | Compress the digest by merging centroids that fall within the same +-- \(K_1\) scale-function unit interval. +-- +-- The compression algorithm works as follows: +-- +-- 1. Compute the integer range of the \(K_1\) scale function: +-- \(j_{\min} = \lceil k(0, \delta) \rceil = \lceil -\delta/2 \rceil\) and +-- \(j_{\max} = \lfloor k(1, \delta) \rfloor = \lfloor \delta/2 \rfloor\). +-- +-- 2. For each integer \(j \in \{j_{\min}+1, \ldots, j_{\max}\}\), compute the +-- cumulative weight boundary \(b_j = k^{-1}(j, \delta) \cdot N\), where +-- \(N\) is the total weight. +-- +-- 3. Split the finger tree at each boundary \(b_j\) by cumulative weight +-- (using @'FT.split' (\m -> mWeight m > b_j)@), yielding \(O(\delta)\) +-- contiguous chunks. +-- +-- 4. Collapse each chunk into a single centroid using the @mMeanWeightSum@ +-- and @mWeight@ components of the monoidal measure: +-- \(\bar{m} = \texttt{mMeanWeightSum}\,/\,\texttt{mWeight}\). This is +-- \(O(1)\) per chunk — no traversal of individual centroids is needed. +-- +-- __Complexity:__ \(O(\delta \log n)\), because there are \(O(\delta)\) split +-- operations, each costing \(O(\log n)\) where \(n\) is the pre-compression +-- centroid count. After compression, the centroid count is at most +-- \(\delta + 1\). +compress :: TDigest -> TDigest +compress td + | cnt <= 1 = td + | otherwise = + let n = tdTotalWeight td + delta = tdDelta td + cs = tdCentroids td + -- K1 range: k(0) = -delta/2, k(1) = delta/2 + -- Integer unit boundaries from ceil(k(0)) to floor(k(1)) + kMin = kScale delta 0 -- = -delta/2 + kMax = kScale delta 1 -- = +delta/2 + jMin = ceiling kMin :: Int + jMax = floor kMax :: Int + -- Build boundaries: q values at each integer k-value + boundaries = [kScaleInv delta (fromIntegral j) * n | j <- [jMin + 1 .. jMax]] + -- Split-and-merge at each boundary + merged = splitMerge boundaries cs + in td {tdCentroids = merged} + where + cnt = mCount (FT.measure (tdCentroids td)) + +-- | Split a finger tree at cumulative weight boundaries and merge each +-- chunk into a single centroid. This is the inner loop of 'compress'. +-- +-- The function walks through the list of weight boundaries, performing +-- an @'FT.split'@ at each one. Each resulting chunk (a contiguous sub-tree +-- of centroids whose combined weight falls within a single \(K_1\) unit +-- interval) is collapsed via 'mergeChunk' into a single centroid and appended +-- to the accumulator. +splitMerge :: [Double] -> FingerTree Measure Centroid -> FingerTree Measure Centroid +splitMerge boundaries tree = go boundaries tree FT.empty + where + go [] remaining acc = + -- Last chunk: everything remaining + case mergeChunk remaining of + Nothing -> acc + Just c -> acc |> c + go (b : bs) remaining acc = + let (chunk, rest) = FT.split (\m -> mWeight m > b) remaining + in case mergeChunk chunk of + Nothing -> go bs rest acc + Just c -> go bs rest (acc |> c) + +-- | Merge all centroids in a finger tree chunk into a single centroid +-- using the monoidal measure. Runs in \(O(1)\) — no traversal of +-- individual centroids is needed, because the measure already caches +-- \(\sum w_i\) and \(\sum m_i w_i\). +mergeChunk :: FingerTree Measure Centroid -> Maybe Centroid +mergeChunk ft + | w == 0 = Nothing + | otherwise = Just (Centroid (mws / w) w) + where + m = FT.measure ft + w = mWeight m + mws = mMeanWeightSum m + +-- --------------------------------------------------------------------------- +-- Quantile estimation +-- --------------------------------------------------------------------------- + +-- | Estimate the value at quantile \(q\) (\(0 \le q \le 1\)). +-- +-- The algorithm uses an interpolation scheme that treats each centroid as +-- representing a point mass at its mean, spread uniformly over a weight +-- interval centred at the centroid's cumulative midpoint. Between +-- consecutive centroid midpoints, the estimated quantile function is linearly +-- interpolated: +-- +-- \[ +-- \hat{x}(q) \;=\; m_i + \frac{q \cdot N - \mathrm{mid}_i} +-- {\mathrm{mid}_{i+1} - \mathrm{mid}_i} \cdot (m_{i+1} - m_i) +-- \] +-- +-- where \(\mathrm{mid}_i = \sum_{j TDigest -> Maybe Double +quantile q td + | numCentroids == 0 = Nothing + | numCentroids == 1 = + case FT.viewl cs of + c :< _ -> Just (cMean c) + EmptyL -> Nothing + | otherwise = Just (findQuantile (clamp 0 1 q)) + where + cs = tdCentroids td + n = tdTotalWeight td + mn = tdMin td + mx = tdMax td + numCentroids = mCount (FT.measure cs) + + findQuantile :: Double -> Double + findQuantile q' = + let target = q' * n + (left, right) = FT.split (\m -> mWeight m > target) cs + leftWeight = mWeight (FT.measure left) + leftCount = mCount (FT.measure left) + in case FT.viewl right of + EmptyL -> + case FT.viewr left of + _ :> lastC -> interpolateRight lastC (leftWeight - cWeight lastC) target + EmptyR -> mx + cur :< rightRest -> + interpolateAt leftCount leftWeight cur left rightRest target + + interpolateAt :: Int -> Double -> Centroid -> FingerTree Measure Centroid -> FingerTree Measure Centroid -> Double -> Double + interpolateAt i cumulative c left rest target + | i == 0 && target < cWeight c / 2 = + if cWeight c == 1 + then mn + else mn + (cMean c - mn) * (target / (cWeight c / 2)) + | i == numCentroids - 1 = + if target > n - cWeight c / 2 + then + if cWeight c == 1 + then mx + else + let remaining = n - cWeight c / 2 + in cMean c + (mx - cMean c) * ((target - remaining) / (cWeight c / 2)) + else cMean c + | otherwise = + let mid = cumulative + cWeight c / 2 + in case FT.viewl rest of + nextC :< _ -> + let nextMid = cumulative + cWeight c + cWeight nextC / 2 + in if target <= nextMid + then + let frac = + if nextMid == mid + then 0.5 + else (target - mid) / (nextMid - mid) + in cMean c + frac * (cMean nextC - cMean c) + else + let newLeft = left FT.>< FT.singleton c + in interpolateAt (i + 1) (cumulative + cWeight c) nextC newLeft (ftTail rest) target + EmptyL -> cMean c + + interpolateRight :: Centroid -> Double -> Double -> Double + interpolateRight c _cumulative target = + if target > n - cWeight c / 2 + then + if cWeight c == 1 + then mx + else + let remaining = n - cWeight c / 2 + in cMean c + (mx - cMean c) * ((target - remaining) / (cWeight c / 2)) + else cMean c + + ftTail :: FingerTree Measure Centroid -> FingerTree Measure Centroid + ftTail ft = case FT.viewl ft of + EmptyL -> FT.empty + _ :< r -> r + +-- --------------------------------------------------------------------------- +-- CDF estimation +-- --------------------------------------------------------------------------- + +-- | Estimate the cumulative distribution function (CDF) at value \(x\), +-- i.e., the fraction of the distribution that lies at or below \(x\). +-- +-- The CDF is estimated by piecewise-linear interpolation between centroid +-- midpoints. For a query point \(x\) falling between the means of +-- consecutive centroids \(m_i\) and \(m_{i+1}\), the estimated CDF is: +-- +-- \[ +-- \hat{F}(x) \;=\; \frac{1}{N}\left( +-- \mathrm{mid}_i + \frac{x - m_i}{m_{i+1} - m_i} +-- \cdot (\mathrm{mid}_{i+1} - \mathrm{mid}_i) +-- \right) +-- \] +-- +-- where \(\mathrm{mid}_i = \sum_{j TDigest -> Maybe Double +cdf x td + | numCentroids == 0 = Nothing + | x <= mn = Just 0 + | x >= mx = Just 1 + | otherwise = Just (findCdf x) + where + cs = tdCentroids td + n = tdTotalWeight td + mn = tdMin td + mx = tdMax td + numCentroids = mCount (FT.measure cs) + + findCdf :: Double -> Double + findCdf x' = + let (left, right) = splitByMean x' cs + in case (FT.viewr left, FT.viewl right) of + (EmptyR, rc :< _) -> + cdfAtFirst rc x' + (_, EmptyL) -> + case FT.viewr left of + lRest :> lc -> + cdfAtLast lc (mWeight (FT.measure lRest)) x' + EmptyR -> 1.0 + (lRest :> lc, rc :< _) -> + let lcCum = mWeight (FT.measure lRest) + lcIdx = mCount (FT.measure lRest) + rcIdx = mCount (FT.measure left) + in if x' <= cMean lc + then + if lcIdx == 0 + then cdfAtFirst lc x' + else case FT.viewr lRest of + llRest :> llc -> + cdfBetween llc (mWeight (FT.measure llRest)) lc lcCum x' + EmptyR -> cdfAtFirst lc x' + else + if rcIdx == numCentroids - 1 && x' > cMean rc + then cdfAtLast rc (mWeight (FT.measure left)) x' + else cdfBetween lc lcCum rc (mWeight (FT.measure left)) x' + + cdfAtFirst :: Centroid -> Double -> Double + cdfAtFirst c x' + | x' < cMean c = + let innerW = cWeight c / 2 + frac = + if cMean c == mn + then 1.0 + else (x' - mn) / (cMean c - mn) + in (innerW * frac) / n + | otherwise = (cWeight c / 2) / n + + cdfAtLast :: Centroid -> Double -> Double -> Double + cdfAtLast c cumBefore x' + | x' > cMean c = + let halfW = cWeight c / 2 + rightW = n - cumBefore - halfW + frac = + if mx == cMean c + then 0.0 + else (x' - cMean c) / (mx - cMean c) + in (cumBefore + halfW + rightW * frac) / n + | otherwise = (cumBefore + cWeight c / 2) / n + + cdfBetween :: Centroid -> Double -> Centroid -> Double -> Double -> Double + cdfBetween lc lcCum rc rcCum x' + | x' <= cMean lc = (lcCum + cWeight lc / 2) / n + | x' >= cMean rc = (rcCum + cWeight rc / 2) / n + | otherwise = + let lMid = lcCum + cWeight lc / 2 + rMid = rcCum + cWeight rc / 2 + frac = + if cMean lc == cMean rc + then 0.5 + else (x' - cMean lc) / (cMean rc - cMean lc) + in (lMid + frac * (rMid - lMid)) / n + +-- --------------------------------------------------------------------------- +-- Merge +-- --------------------------------------------------------------------------- + +-- | Merge two t-digests into one, preserving accuracy. +-- +-- The merge operation inserts every centroid of the second digest into the +-- first (using 'addWeighted' with the centroid's mean and weight), then +-- applies 'compress' to restore the \(K_1\) size invariant. +-- +-- This is the standard approach for combining digests computed on +-- disjoint data partitions, enabling distributed and parallel quantile +-- estimation. In a MapReduce-style pipeline, each mapper builds a local +-- t'TDigest' and the reducer merges them with 'merge'. Because 'compress' +-- enforces the same \(O(\delta)\) centroid bound, the merged result has +-- the same space footprint as a single-stream digest. +-- +-- See Dunning (2021), Section 4.3 () +-- for a discussion of mergeability and its applications. +merge :: TDigest -> TDigest -> TDigest +merge td other = + let otherCs = ftToList (tdCentroids other) + combined = foldl' (\d c -> addWeighted (cMean c) (cWeight c) d) td otherCs + in compress combined + +-- --------------------------------------------------------------------------- +-- Queries +-- --------------------------------------------------------------------------- + +-- | Return the total weight of all values added to the digest. +-- +-- This is \(O(1)\), as the total weight is cached in the t'TDigest' record. +-- For an unweighted stream, this equals the number of observations. +totalWeight :: TDigest -> Double +totalWeight = tdTotalWeight + +-- | Return the number of centroids currently stored in the digest. +-- +-- This is \(O(1)\) via the @mCount@ component of the finger tree's monoidal +-- measure. The count is always at most \(3\delta\) (and at most +-- \(\delta + 1\) immediately after 'compress'). +centroidCount :: TDigest -> Int +centroidCount = mCount . FT.measure . tdCentroids + +-- --------------------------------------------------------------------------- +-- Utility +-- --------------------------------------------------------------------------- + +clamp :: Double -> Double -> Double -> Double +clamp lo hi x + | x < lo = lo + | x > hi = hi + | otherwise = x + +-- --------------------------------------------------------------------------- +-- Additional accessors (for Mutable interop) +-- --------------------------------------------------------------------------- + +-- | Return the list of centroids in sorted order (by mean). +-- +-- Useful for serialisation, interoperability with mutable implementations, +-- debugging, and visualisation of the digest's internal distribution. The +-- list is produced by an in-order traversal of the finger tree in +-- \(O(n)\). +centroidList :: TDigest -> [Centroid] +centroidList = ftToList . tdCentroids + +-- | Return the compression parameter \(\delta\). +-- +-- This is needed for serialisation and for reconstructing a digest with +-- 'fromComponents'. +getDelta :: TDigest -> Double +getDelta = tdDelta + +-- | Return the minimum observed value. +-- +-- The global minimum is tracked separately from the centroids because the +-- first centroid's mean may be larger than the minimum (if multiple values +-- have been merged into it). The minimum is used for boundary interpolation +-- in 'quantile' and 'cdf' at \(q \to 0\) and \(x \to \min\). +getMin :: TDigest -> Double +getMin = tdMin + +-- | Return the maximum observed value. +-- +-- Symmetric to 'getMin': the global maximum is used for boundary +-- interpolation in 'quantile' and 'cdf' at \(q \to 1\) and \(x \to \max\). +getMax :: TDigest -> Double +getMax = tdMax + +-- | Reconstruct a t-digest from its serialised components: a list of +-- centroids (which /must/ be in non-decreasing order of mean), the total +-- weight, the global minimum and maximum, and the compression parameter +-- \(\delta\). +-- +-- This function trusts the caller to provide correctly sorted centroids and +-- consistent metadata. It is intended for deserialisation and for +-- transferring digests between this pure implementation and the mutable +-- array-backed implementations in other languages. No validation or +-- re-compression is performed. +-- +-- __Usage example:__ +-- +-- @ +-- let cs = 'centroidList' td +-- tw = 'totalWeight' td +-- mn = 'getMin' td +-- mx = 'getMax' td +-- d = 'getDelta' td +-- td' = 'fromComponents' cs tw mn mx d +-- -- td' is equivalent to td +-- @ +fromComponents :: [Centroid] -> Double -> Double -> Double -> Double -> TDigest +fromComponents cs tw mn mx delta = + TDigest + { tdCentroids = FT.fromList cs, + tdTotalWeight = tw, + tdMin = mn, + tdMax = mx, + tdDelta = delta, + tdMaxCentroids = ceiling (delta * 3) + } diff --git a/src/Data/Sketch/TDigest/Mutable.hs b/src/Data/Sketch/TDigest/Mutable.hs new file mode 100644 --- /dev/null +++ b/src/Data/Sketch/TDigest/Mutable.hs @@ -0,0 +1,1022 @@ +-- | +-- Module : Data.Sketch.TDigest.Mutable +-- Description : Mutable t-digest via buffer-and-flush with greedy merge in the ST monad +-- Copyright : (c) Nadia Yvette Chambers, 2025 +-- License : BSD-3-Clause +-- Maintainer : nadia.yvette.chambers@gmail.com +-- Stability : experimental +-- +-- A mutable t-digest implementation backed by mutable vectors from the +-- @vector@ package, operating entirely within the 'Control.Monad.ST.ST' +-- monad. Centroids are stored in a mutable unboxed-style vector of +-- @(mean, weight)@ pairs kept sorted by mean. Prefix sums of weights +-- are maintained for \(O(\log n)\) quantile and CDF queries via binary +-- search. +-- +-- == Background +-- +-- The /t-digest/ is a streaming, mergeable sketch for approximate quantile +-- estimation, introduced by Dunning (2021) +-- (). It belongs to the +-- family of quantile summaries that trade bounded space for approximate +-- answers, a line of work originating with Munro & Paterson (1980) +-- () and continued by +-- Greenwald & Khanna (2001) +-- (). The key innovation of the +-- t-digest is the use of a /scale function/ to allow larger centroids in +-- the interior of the distribution while keeping centroids near the tails +-- small, yielding high relative accuracy at extreme quantiles (e.g., +-- \(q = 0.99\) or \(q = 0.001\)). +-- +-- This module provides the /mutable/ variant, which follows a +-- /buffer-and-flush/ strategy: incoming data points are appended to an +-- unsorted buffer in amortised \(O(1)\) time; when the buffer reaches +-- capacity, the entire buffer is flushed into the sorted centroid array +-- via insertion sort followed by a single-pass greedy merge. This +-- amortised design is the approach recommended by Dunning & Ertl (2019) +-- () for high-throughput ingestion. +-- +-- == The ST monad approach +-- +-- This module uses 'Control.Monad.ST.ST' rather than 'IO' for in-place +-- mutation. The 'ST' monad provides: +-- +-- * /True in-place mutation/ — centroid vectors, prefix-sum arrays, and +-- the pending-addition buffer are modified destructively, avoiding the +-- allocation overhead of persistent data structures. +-- * /Rank-2 type safety/ — the universally quantified state token @s@ in +-- 'runTDigest' (equivalently 'Control.Monad.ST.runST') guarantees that +-- no mutable reference can escape the computation. This is enforced +-- statically by the type system, with no runtime cost. +-- * /No IO escape/ — unlike @IORef@ or @IOVector@, 'STRef' and +-- 'Data.Vector.Mutable.MVector' in 'ST' cannot perform arbitrary +-- side-effects. The result of 'runTDigest' is a pure value. +-- +-- For a purely functional alternative that avoids mutable state entirely, +-- see "Data.Sketch.TDigest", which stores centroids in a finger tree +-- (Hinze & Paterson, 2006; +-- ) with a four-component +-- monoidal measure, providing \(O(\log n)\) insertion without buffering +-- and \(O(\delta \log n)\) compression via split-based merge. +-- +-- == Space bounds +-- +-- The t-digest maintains at most \(O(\delta)\) centroids after each +-- compression pass, where \(\delta\) is the compression parameter +-- (default 100). Between compressions the buffer may hold up to +-- \(5\delta\) pending additions, so peak memory usage is bounded by +-- \(O(\delta)\) centroid slots plus \(O(\delta)\) buffer slots, for a +-- total working set of \(O(\delta)\). The initial centroid vector is +-- allocated with capacity \(10\delta\) to accommodate the merge of the +-- buffer contents with the existing centroids without reallocation in +-- steady state. +-- +-- Because \(\delta\) is a user-chosen constant (typically 100–300), space +-- usage is /independent of the number of data points/ ingested — +-- precisely the guarantee required for streaming applications. +-- +-- == Algorithm +-- +-- The core algorithm is /buffer-and-flush with greedy merge/: +-- +-- 1. __Buffer phase.__ Each call to 'addWeighted' appends the +-- @(mean, weight)@ pair to the end of the buffer in \(O(1)\) +-- amortised time (the buffer is doubled if it overflows). When the +-- buffer length reaches the capacity \(5\delta\), 'compress' is +-- triggered automatically. +-- +-- 2. __Sort phase.__ On compress, all existing centroids and buffered +-- points are collected into a single temporary array and sorted by +-- mean using insertion sort. Insertion sort is chosen because the +-- existing centroids are already sorted, so the merge of two sorted +-- runs is nearly linear; in practice the buffer is small relative to +-- the total. +-- +-- 3. __Greedy merge phase.__ The sorted array is traversed left to +-- right. A running centroid accumulates incoming points as long as +-- the K1 scale function constraint is satisfied: +-- +-- \[ +-- k_1(q, \delta) = \frac{\delta}{2\pi} \arcsin(2q - 1) +-- \] +-- +-- Two adjacent quantile positions \(q_0\) and \(q_1\) may share a +-- centroid if and only if \(k_1(q_1) - k_1(q_0) \le 1\). When +-- the constraint would be violated, the accumulated centroid is +-- emitted and a new accumulation begins. The merged centroid's mean +-- is the standard weighted mean: +-- +-- \[ +-- \mu_{\text{new}} = \frac{\mu_a \, w_a + \mu_b \, w_b}{w_a + w_b} +-- \] +-- +-- 4. __Prefix-sum rebuild.__ After merging, the prefix-sum array is +-- rebuilt in a single linear pass so that @prefixSum[i]@ equals the +-- cumulative weight of centroids \(0, 1, \ldots, i{-}1\). This +-- array enables \(O(\log n)\) quantile and CDF queries via binary +-- search. +-- +-- == Companion implementations +-- +-- This project contains 28 language implementations of the merging +-- t-digest. While this Haskell module uses flat mutable vectors for +-- simplicity, 22 of the other mutable implementations store centroids +-- in /array-backed 2-3-4 trees/. The 2-3-4 tree is a B-tree of order 4 +-- (Bayer & McCreight, 1972; ), +-- equivalent via the well-known isomorphism to a red-black tree (Guibas +-- & Sedgewick, 1978; +-- ; see also Sedgewick, 2008; +-- +-- for the left-leaning specialisation). +-- +-- The 2-3-4 tree representation offers several advantages for +-- fine-grained quantile workloads: +-- +-- * /Cache locality/ — storing nodes in a contiguous array rather than +-- heap-allocated pointers improves spatial locality and reduces cache +-- misses, which matters when the centroid count \(\delta\) is in the +-- hundreds. +-- * /Worst-case \(O(\log n)\) insertion and deletion/ — unlike the +-- amortised buffer-and-flush approach here, the tree-based variants +-- can absorb each data point immediately with a guaranteed logarithmic +-- bound, which is useful in latency-sensitive contexts. +-- * /Robustness for fine-grained queries/ — maintaining a balanced tree +-- of centroids at all times (rather than deferring organisation to +-- periodic compressions) ensures that quantile and CDF queries always +-- see a fully up-to-date structure. +-- +-- == Quick start +-- +-- @ +-- import Data.Sketch.TDigest.Mutable +-- import Control.Monad (forM_) +-- +-- example :: Maybe Double +-- example = 'runTDigest' $ do +-- td <- 'new' +-- forM_ [1.0 .. 10000.0] $ \\v -> 'add' v td +-- 'quantile' 0.99 td +-- @ +module Data.Sketch.TDigest.Mutable + ( -- * Type + MDigest, + + -- * Construction + new, + newWith, + + -- * Insertion + add, + addWeighted, + + -- * Compression + compress, + + -- * Queries + quantile, + cdf, + + -- * Merging + merge, + + -- * Conversion + freeze, + thaw, + + -- * Accessors + totalWeight, + centroidCount, + + -- * Runner + runTDigest, + ) +where + +import Control.Monad (when) +import Control.Monad.ST (ST, runST) +import Data.STRef + ( STRef, + modifySTRef', + newSTRef, + readSTRef, + writeSTRef, + ) +import qualified Data.Sketch.TDigest as TD +import qualified Data.Vector.Mutable as MV + +-- --------------------------------------------------------------------------- +-- Type +-- --------------------------------------------------------------------------- + +-- | A truly mutable t-digest operating within the 'ST' monad, using +-- mutable vectors for centroids, prefix sums, and a pending-additions +-- buffer. +-- +-- The internal state comprises: +-- +-- * __Centroid vector__ (@mdCentroids@) — a mutable vector of +-- @(mean, weight)@ pairs maintained in sorted order by mean. After +-- each call to 'compress', this vector contains at most \(O(\delta)\) +-- entries. +-- +-- * __Prefix-sum vector__ (@mdPrefixSums@) — a mutable vector of length +-- \(n_c + 1\) (where \(n_c\) is the centroid count) satisfying +-- @prefixSum[0] = 0@ and @prefixSum[i] = \sum_{j=0}^{i-1} w_j@. +-- This enables \(O(\log n_c)\) quantile and CDF queries via binary +-- search without a linear scan. +-- +-- * __Buffer__ (@mdBuffer@) — an unsorted staging area for incoming +-- data points. Points are appended in \(O(1)\) amortised time. +-- When the buffer length reaches the capacity \(5\delta\), a +-- compress cycle is triggered automatically, flushing the buffer +-- into the centroid vector. +-- +-- * __Scalar accumulators__ — @mdTotalWeight@, @mdMin@, and @mdMax@ +-- track the running total weight and extrema across all points ever +-- ingested (including buffered ones not yet compressed). +-- +-- __Invariants.__ Between calls to exported functions: +-- +-- 1. The centroid vector is sorted by mean. +-- 2. The prefix-sum vector is consistent with the centroid vector. +-- 3. The buffer length is in \([0, 5\delta)\). +-- 4. @totalWeight@ equals the sum of all centroid weights plus all +-- buffered point weights. +-- +-- Invariants (1) and (2) may be temporarily violated while the buffer +-- is non-empty; they are restored by 'compress'. +data MDigest s = MDigest + { -- | Mutable vector of (mean, weight) pairs, sorted by mean. + mdCentroids :: !(STRef s (MV.MVector s (Double, Double))), + -- | Prefix sums: prefixSum[0] = 0, prefixSum[i] = sum of weights 0..i-1. + mdPrefixSums :: !(STRef s (MV.MVector s Double)), + -- | Buffer for pending additions. + mdBuffer :: !(STRef s (MV.MVector s (Double, Double))), + mdTotalWeight :: !(STRef s Double), + mdMin :: !(STRef s Double), + mdMax :: !(STRef s Double), + mdBufferLen :: !(STRef s Int), + mdCentroidCount :: !(STRef s Int), + mdDelta :: !(STRef s Double), + mdBufferCap :: !(STRef s Int) + } + +-- --------------------------------------------------------------------------- +-- Construction +-- --------------------------------------------------------------------------- + +-- | Create a new, empty mutable t-digest with the default compression +-- parameter \(\delta = 100\). +-- +-- This is equivalent to @'newWith' 100@. A \(\delta\) of 100 yields +-- roughly 100 centroids after compression and provides relative accuracy +-- on the order of \(10^{-3}\) at extreme quantiles — sufficient for most +-- monitoring and analytics workloads. See Dunning & Ertl (2019) +-- () for empirical accuracy tables. +new :: ST s (MDigest s) +new = newWith 100 + +-- | Create a new, empty mutable t-digest with a given compression +-- parameter \(\delta\). +-- +-- The compression parameter controls the trade-off between accuracy and +-- space. Larger values of \(\delta\) produce more centroids (up to +-- \(O(\delta)\)) and therefore higher accuracy, at the cost of increased +-- memory and compression time. Typical values range from 50 (coarse) to +-- 300 (very accurate). +-- +-- __Buffer capacity.__ The internal buffer is sized to hold +-- \(\lceil 5\delta \rceil\) pending additions. This factor of 5 is an +-- empirical choice: it amortises the cost of compression (which is +-- \(O(\delta)\) per flush) over enough insertions to make the per-insert +-- cost effectively \(O(1)\). +-- +-- __Initial centroid allocation.__ The centroid vector is pre-allocated +-- with capacity \(10\delta\) — enough to hold the existing centroids +-- (at most \(\sim\delta\) after the previous compression) plus a full +-- buffer of \(5\delta\) points, without reallocation during the merge +-- phase. +newWith :: Double -> ST s (MDigest s) +newWith delta = do + let bufCap = ceiling (delta * 5) :: Int + initCentroidCap = bufCap * 2 + centroids <- MV.new initCentroidCap + prefix <- MV.new 1 + MV.write prefix 0 0.0 + buf <- MV.new bufCap + cRef <- newSTRef centroids + pRef <- newSTRef prefix + bRef <- newSTRef buf + twRef <- newSTRef 0.0 + mnRef <- newSTRef (1 / 0) + mxRef <- newSTRef (-(1 / 0)) + blRef <- newSTRef 0 + ccRef <- newSTRef 0 + dRef <- newSTRef delta + bcRef <- newSTRef bufCap + return + MDigest + { mdCentroids = cRef, + mdPrefixSums = pRef, + mdBuffer = bRef, + mdTotalWeight = twRef, + mdMin = mnRef, + mdMax = mxRef, + mdBufferLen = blRef, + mdCentroidCount = ccRef, + mdDelta = dRef, + mdBufferCap = bcRef + } + +-- --------------------------------------------------------------------------- +-- Insertion +-- --------------------------------------------------------------------------- + +-- | Add a single value with unit weight to the digest. +-- +-- @'add' x md = 'addWeighted' x 1 md@ +-- +-- This is the common case for unweighted data streams. The value is +-- appended to the internal buffer in \(O(1)\) amortised time; +-- compression is triggered automatically when the buffer is full. +add :: Double -> MDigest s -> ST s () +add x = addWeighted x 1 + +-- | Add a value with a given weight to the digest. +-- +-- __Complexity.__ Amortised \(O(1)\). The value is appended to the +-- tail of the unsorted buffer; no sorting or merging occurs at this +-- stage. The running minimum, maximum, and total weight are updated +-- eagerly so that they are always available without a compress cycle. +-- +-- __Auto-compress.__ When the buffer length reaches the buffer capacity +-- \(\lceil 5\delta \rceil\), 'compress' is called automatically. This +-- ensures that memory usage never exceeds \(O(\delta)\) beyond the +-- allocated capacity. +-- +-- __Buffer growth.__ If the buffer's underlying vector is full (which +-- can happen if the buffer capacity has been reached but 'compress' has +-- not yet been triggered by a prior code path), the vector is doubled in +-- size via 'Data.Vector.Mutable.grow'. In steady-state operation this +-- branch is not taken because auto-compress fires at the capacity +-- threshold. +addWeighted :: Double -> Double -> MDigest s -> ST s () +addWeighted x w md = do + -- Update min/max + mn <- readSTRef (mdMin md) + when (x < mn) $ writeSTRef (mdMin md) x + mx <- readSTRef (mdMax md) + when (x > mx) $ writeSTRef (mdMax md) x + -- Update total weight + modifySTRef' (mdTotalWeight md) (+ w) + -- Append to buffer + bl <- readSTRef (mdBufferLen md) + buf <- readSTRef (mdBuffer md) + let bufLen = MV.length buf + -- Grow buffer if needed + buf' <- + if bl >= bufLen + then do + newBuf <- MV.grow buf bufLen + writeSTRef (mdBuffer md) newBuf + return newBuf + else return buf + MV.write buf' bl (x, w) + let bl' = bl + 1 + writeSTRef (mdBufferLen md) bl' + -- Compress if buffer is full + bc <- readSTRef (mdBufferCap md) + when (bl' >= bc) $ compress md + +-- --------------------------------------------------------------------------- +-- Compression +-- --------------------------------------------------------------------------- + +-- | Force compression of the buffer into the centroid list. +-- +-- Compression implements the /buffer-and-flush/ strategy described by +-- Dunning & Ertl (2019) (). The +-- algorithm proceeds in four stages: +-- +-- 1. __Collect.__ All existing centroids and buffered points are copied +-- into a single temporary array of length \(n_c + n_b\). +-- +-- 2. __Sort.__ The temporary array is sorted by centroid mean using +-- insertion sort. Because the first \(n_c\) entries are already in +-- sorted order (they come from the centroid vector), the sort is +-- adaptive: it performs at most \(O(n_b \cdot (n_c + n_b))\) +-- comparisons, which is efficient when \(n_b \ll n_c\). +-- +-- 3. __Greedy merge.__ The sorted array is traversed left to right. A +-- running centroid accumulates successive entries as long as the K1 +-- scale function constraint is satisfied. The K1 scale function is +-- defined as: +-- +-- \[ +-- k_1(q, \delta) \;=\; \frac{\delta}{2\pi}\,\arcsin(2q - 1) +-- \] +-- +-- Given a running accumulated weight \(W_{\text{so far}}\) and a +-- total digest weight \(N\), the quantile interval of the proposed +-- merged centroid spans \([q_0, q_1]\) where +-- \(q_0 = W_{\text{so far}} / N\) and +-- \(q_1 = (W_{\text{so far}} + w_{\text{proposed}}) / N\). +-- The merge is permitted if: +-- +-- \[ +-- k_1(q_1, \delta) - k_1(q_0, \delta) \;\le\; 1 +-- \] +-- +-- When this constraint would be violated, the accumulated centroid is +-- emitted and a fresh accumulation begins. Singletons (weight \(\le 1\)) +-- are always merged with their neighbour when not at the boundary, to +-- prevent centroid count blow-up from unit-weight insertions. +-- +-- 4. __Rebuild prefix sums.__ A single linear pass rebuilds the +-- prefix-sum array for subsequent \(O(\log n)\) queries. +-- +-- __Complexity.__ \(O((n_c + n_b)^2)\) worst-case due to insertion sort, +-- but \(O(n_c + n_b)\) in the common case when the buffer is small +-- relative to the sorted centroid array. The output centroid count is +-- bounded by \(O(\delta)\). +compress :: MDigest s -> ST s () +compress md = do + bl <- readSTRef (mdBufferLen md) + cc <- readSTRef (mdCentroidCount md) + when (bl > 0 || cc > 1) $ do + -- Collect all items: existing centroids + buffer + let totalItems = cc + bl + allItems <- MV.new totalItems + -- Copy centroids + centroids <- readSTRef (mdCentroids md) + copyN centroids allItems cc 0 0 + -- Copy buffer + buf <- readSTRef (mdBuffer md) + copyN buf allItems bl 0 cc + -- Sort all items by mean (insertion sort is fine for small arrays) + insertionSort allItems totalItems + -- Greedy merge + delta <- readSTRef (mdDelta md) + n <- readSTRef (mdTotalWeight md) + if totalItems == 0 + then do + writeSTRef (mdCentroidCount md) 0 + writeSTRef (mdBufferLen md) 0 + rebuildPrefixSums md + else do + -- Merge in-place into a result vector + merged <- MV.new totalItems + (m0, w0) <- MV.read allItems 0 + -- Walk and merge + newCount <- greedyMergeVec delta n allItems totalItems merged m0 w0 + -- Write back + writeSTRef (mdCentroids md) merged + writeSTRef (mdCentroidCount md) newCount + writeSTRef (mdBufferLen md) 0 + rebuildPrefixSums md + +-- Copy n elements from src starting at srcOff to dst starting at dstOff +copyN :: MV.MVector s (Double, Double) -> MV.MVector s (Double, Double) -> Int -> Int -> Int -> ST s () +copyN src dst n srcOff dstOff = go 0 + where + go i + | i >= n = return () + | otherwise = do + v <- MV.read src (srcOff + i) + MV.write dst (dstOff + i) v + go (i + 1) + +-- Insertion sort by first element of pair +insertionSort :: MV.MVector s (Double, Double) -> Int -> ST s () +insertionSort vec n = go 1 + where + go i + | i >= n = return () + | otherwise = do + val@(key, _) <- MV.read vec i + j <- findInsertPos vec key (i - 1) + -- Shift elements right + shiftRight vec (j + 1) i + MV.write vec (j + 1) val + go (i + 1) + + findInsertPos :: MV.MVector s (Double, Double) -> Double -> Int -> ST s Int + findInsertPos _ _ (-1) = return (-1) + findInsertPos v key j = do + (jKey, _) <- MV.read v j + if jKey > key + then findInsertPos v key (j - 1) + else return j + + shiftRight :: MV.MVector s (Double, Double) -> Int -> Int -> ST s () + shiftRight v from to + | from >= to = return () + | otherwise = go' (to - 1) + where + go' j + | j < from = return () + | otherwise = do + val <- MV.read v j + MV.write v (j + 1) val + go' (j - 1) + +-- Greedy merge: walk sorted items, merge adjacent when scale function allows. +-- Returns the number of merged centroids written to 'out'. +greedyMergeVec :: + Double -> + Double -> + MV.MVector s (Double, Double) -> + Int -> + MV.MVector s (Double, Double) -> + Double -> + Double -> + ST s Int +greedyMergeVec delta n items totalItems out initMean initWeight = go 1 0 initMean initWeight 0 + where + k q = (delta / (2 * pi)) * asin (2 * q - 1) + + go idx weightSoFar curMean curWeight outIdx + | idx >= totalItems = do + -- Emit final centroid + MV.write out outIdx (curMean, curWeight) + return (outIdx + 1) + | otherwise = do + (itemMean, itemWeight) <- MV.read items idx + let proposed = curWeight + itemWeight + q0 = weightSoFar / n + q1 = (weightSoFar + proposed) / n + canMerge = + (proposed <= 1 && idx < totalItems - 1) + || (k q1 - k q0 <= 1.0) + if canMerge + then do + -- Merge: weighted mean + let newW = curWeight + itemWeight + newM = (curMean * curWeight + itemMean * itemWeight) / newW + go (idx + 1) weightSoFar newM newW outIdx + else do + -- Emit current centroid, start new one + MV.write out outIdx (curMean, curWeight) + go (idx + 1) (weightSoFar + curWeight) itemMean itemWeight (outIdx + 1) + +-- Rebuild prefix sums from current centroids. +-- prefixSum has (centroidCount + 1) entries: +-- prefixSum[0] = 0 +-- prefixSum[i] = sum of weights of centroids 0..i-1 +rebuildPrefixSums :: MDigest s -> ST s () +rebuildPrefixSums md = do + cc <- readSTRef (mdCentroidCount md) + prefix <- MV.new (cc + 1) + MV.write prefix 0 0.0 + centroids <- readSTRef (mdCentroids md) + buildPS centroids prefix cc 0 0.0 + writeSTRef (mdPrefixSums md) prefix + where + buildPS _ _ n i _ + | i >= n = return () + buildPS cs ps n i acc = do + (_, w) <- MV.read cs i + let acc' = acc + w + MV.write ps (i + 1) acc' + buildPS cs ps n (i + 1) acc' + +-- --------------------------------------------------------------------------- +-- Queries +-- --------------------------------------------------------------------------- + +-- | Estimate the value at quantile \(q\) where \(0 \le q \le 1\). +-- +-- Returns 'Nothing' if the digest is empty (no data points have been +-- added). +-- +-- __Algorithm.__ The digest is first compressed (flushing any buffered +-- points) to ensure the centroid vector and prefix sums are up to date. +-- A binary search on the prefix-sum array locates the centroid \(c_i\) +-- whose cumulative weight interval contains the target rank +-- \(t = q \cdot N\). The returned value is then computed by linear +-- interpolation between adjacent centroid midpoints: +-- +-- * For the /leftmost/ centroid (\(i = 0\)), the target rank +-- \(t < w_0 / 2\) triggers interpolation between the observed minimum +-- and \(\mu_0\). +-- * For the /rightmost/ centroid (\(i = n_c - 1\)), the target rank +-- \(t > N - w_{n_c - 1} / 2\) triggers interpolation between +-- \(\mu_{n_c - 1}\) and the observed maximum. +-- * For /interior/ centroids, the result is linearly interpolated +-- between the midpoints of \(c_i\) and \(c_{i+1}\): +-- +-- \[ +-- \hat{x} = \mu_i + \frac{t - m_i}{m_{i+1} - m_i} \cdot (\mu_{i+1} - \mu_i) +-- \] +-- +-- where \(m_i = \text{cumBefore}_i + w_i / 2\) is the midpoint rank +-- of centroid \(i\). +-- +-- __Complexity.__ \(O(\delta)\) due to the initial compress (if the +-- buffer is non-empty), then \(O(\log \delta)\) for the binary search. +-- If the buffer is already empty, the cost is \(O(\log \delta)\). +quantile :: Double -> MDigest s -> ST s (Maybe Double) +quantile q md = do + compress md + cc <- readSTRef (mdCentroidCount md) + if cc == 0 + then return Nothing + else + if cc == 1 + then do + centroids <- readSTRef (mdCentroids md) + (m, _) <- MV.read centroids 0 + return (Just m) + else do + n <- readSTRef (mdTotalWeight md) + mn <- readSTRef (mdMin md) + mx <- readSTRef (mdMax md) + let q' = clamp 0 1 q + target = q' * n + centroids <- readSTRef (mdCentroids md) + prefix <- readSTRef (mdPrefixSums md) + -- Binary search: find largest i such that prefixSum[i] <= target + -- i is in [0, cc], and represents the centroid index boundary + i <- bsearchPrefix prefix (cc + 1) target + -- i is the index into prefix sums; the centroid index is (i - 1) + -- but we need to handle boundary cases + let ci = max 0 (min (cc - 1) (i - 1)) + -- Now interpolate + (cMean, cWeight) <- MV.read centroids ci + cumBefore <- MV.read prefix ci + let mid = cumBefore + cWeight / 2.0 + if ci == 0 && target < cWeight / 2.0 + then do + -- Left boundary: interpolate between min and first centroid + let result = + if cWeight == 1 + then mn + else mn + (cMean - mn) * (target / (cWeight / 2.0)) + return (Just result) + else + if ci == cc - 1 + then do + -- Right boundary + let remaining = n - cWeight / 2.0 + if target > n - cWeight / 2.0 + then do + let result = + if cWeight == 1 + then mx + else cMean + (mx - cMean) * ((target - remaining) / (cWeight / 2.0)) + return (Just result) + else return (Just cMean) + else do + -- Middle: interpolate between adjacent centroid midpoints + (nextMean, nextWeight) <- MV.read centroids (ci + 1) + cumNext <- MV.read prefix (ci + 1) + let nextMid = cumNext + nextWeight / 2.0 + if target <= nextMid + then do + let frac = + if nextMid == mid + then 0.5 + else (target - mid) / (nextMid - mid) + return (Just (cMean + frac * (nextMean - cMean))) + else do + -- Walk forward from ci+1 + walkQuantile centroids prefix cc n mn mx target (ci + 1) + +-- Walk forward to find the right centroid for the target +walkQuantile :: + MV.MVector s (Double, Double) -> + MV.MVector s Double -> + Int -> + Double -> + Double -> + Double -> + Double -> + Int -> + ST s (Maybe Double) +walkQuantile centroids prefix cc n mn mx target = go + where + go i + | i >= cc = return (Just mx) + | otherwise = do + (cMean, cWeight) <- MV.read centroids i + cumBefore <- MV.read prefix i + let mid = cumBefore + cWeight / 2.0 + if i == 0 && target < cWeight / 2.0 + then do + let result = + if cWeight == 1 + then mn + else mn + (cMean - mn) * (target / (cWeight / 2.0)) + return (Just result) + else + if i == cc - 1 + then do + let remaining = n - cWeight / 2.0 + if target > remaining + then do + let result = + if cWeight == 1 + then mx + else cMean + (mx - cMean) * ((target - remaining) / (cWeight / 2.0)) + return (Just result) + else return (Just cMean) + else do + (nextMean, nextWeight) <- MV.read centroids (i + 1) + cumNext <- MV.read prefix (i + 1) + let nextMid = cumNext + nextWeight / 2.0 + if target <= nextMid + then do + let frac = + if nextMid == mid + then 0.5 + else (target - mid) / (nextMid - mid) + return (Just (cMean + frac * (nextMean - cMean))) + else go (i + 1) + +-- Binary search on prefix sums: find largest i in [0, len-1] such that +-- prefix[i] <= target. +bsearchPrefix :: MV.MVector s Double -> Int -> Double -> ST s Int +bsearchPrefix prefix len target = go 0 (len - 1) + where + go lo hi + | lo >= hi = return lo + | otherwise = do + let mid = (lo + hi + 1) `div` 2 + v <- MV.read prefix mid + if v <= target + then go mid hi + else go lo (mid - 1) + +-- | Estimate the cumulative distribution function (CDF) at value \(x\), +-- i.e., the fraction of the distribution that lies at or below \(x\). +-- +-- Returns 'Nothing' if the digest is empty. +-- +-- __Algorithm.__ Like 'quantile', this function first compresses any +-- buffered points. It then performs a linear walk over the centroid +-- vector to locate the pair of centroids straddling \(x\), and +-- interpolates: +-- +-- * If \(x \le x_{\min}\), the result is 0. +-- * If \(x \ge x_{\max}\), the result is 1. +-- * If \(x\) falls in the half-weight region of the first centroid +-- (i.e., \(x < \mu_0\)), the result is interpolated between 0 and +-- \(w_0 / (2N)\). +-- * If \(x\) falls in the half-weight region of the last centroid, +-- the result is interpolated between +-- \((\sum w - w_{n-1}/2) / N\) and 1. +-- * Otherwise, the result is linearly interpolated between the midpoint +-- ranks of the two bracketing centroids, yielding: +-- +-- \[ +-- \widehat{F}(x) = \frac{m_i + \frac{x - \mu_i}{\mu_{i+1} - \mu_i} \cdot (m_{i+1} - m_i)}{N} +-- \] +-- +-- __Complexity.__ \(O(\delta)\) due to compression plus a linear walk +-- over centroids. +cdf :: Double -> MDigest s -> ST s (Maybe Double) +cdf x md = do + compress md + cc <- readSTRef (mdCentroidCount md) + if cc == 0 + then return Nothing + else do + mn <- readSTRef (mdMin md) + mx <- readSTRef (mdMax md) + if x <= mn + then return (Just 0) + else + if x >= mx + then return (Just 1) + else do + n <- readSTRef (mdTotalWeight md) + centroids <- readSTRef (mdCentroids md) + prefix <- readSTRef (mdPrefixSums md) + walkCdf centroids prefix cc n mn mx x + +walkCdf :: + MV.MVector s (Double, Double) -> + MV.MVector s Double -> + Int -> + Double -> + Double -> + Double -> + Double -> + ST s (Maybe Double) +walkCdf centroids prefix cc n mn mx x = go 0 + where + lastIdx = cc - 1 + + go i + | i >= cc = return (Just 1.0) + | otherwise = do + (cMean, cWeight) <- MV.read centroids i + cumBefore <- MV.read prefix i + if i == 0 && x < cMean + then do + let innerW = cWeight / 2.0 + frac = + if cMean == mn + then 1.0 + else (x - mn) / (cMean - mn) + return (Just ((innerW * frac) / n)) + else + if i == 0 && x == cMean + then return (Just ((cWeight / 2.0) / n)) + else + if i == lastIdx && x > cMean + then do + let halfW = cWeight / 2.0 + rightW = n - cumBefore - halfW + frac = + if mx == cMean + then 0.0 + else (x - cMean) / (mx - cMean) + return (Just ((cumBefore + halfW + rightW * frac) / n)) + else + if i == lastIdx + then return (Just ((cumBefore + cWeight / 2.0) / n)) + else do + let mid = cumBefore + cWeight / 2.0 + (nextMean, nextWeight) <- MV.read centroids (i + 1) + cumNext <- MV.read prefix (i + 1) + let nextMid = cumNext + nextWeight / 2.0 + if x < nextMean + then do + let frac = + if cMean == nextMean + then 0.5 + else (x - cMean) / (nextMean - cMean) + return (Just ((mid + frac * (nextMid - mid)) / n)) + else go (i + 1) + +-- --------------------------------------------------------------------------- +-- Accessors +-- --------------------------------------------------------------------------- + +-- | Return the total weight of all values added to the digest. +-- +-- This includes both compressed centroids and pending buffer entries. +-- The value is maintained eagerly (updated on every 'addWeighted' call), +-- so this accessor is \(O(1)\) and does not trigger compression. +totalWeight :: MDigest s -> ST s Double +totalWeight md = readSTRef (mdTotalWeight md) + +-- | Return the number of centroids, compressing any pending buffer first. +-- +-- Because the true centroid count is only well-defined after all buffered +-- points have been merged, this function calls 'compress' before reading +-- the count. If no buffer entries are pending, the compress is a no-op +-- (the guard @bl > 0 || cc > 1@ shortcuts immediately). +-- +-- __Complexity.__ \(O(\delta)\) if compression is needed, \(O(1)\) +-- otherwise. +centroidCount :: MDigest s -> ST s Int +centroidCount md = do + compress md + readSTRef (mdCentroidCount md) + +-- --------------------------------------------------------------------------- +-- Merge +-- --------------------------------------------------------------------------- + +-- | Merge a pure 'TD.TDigest' into the mutable digest. +-- +-- The pure digest is first compressed, then its centroids are extracted +-- as a list and fed one by one into 'addWeighted'. This triggers the +-- standard buffer-and-flush lifecycle: centroids accumulate in the +-- buffer and are flushed when the buffer fills. +-- +-- This operation is useful in /parallel and distributed/ settings: each +-- worker thread can build a local pure 'TD.TDigest' (or a local +-- t'MDigest' frozen via 'freeze'), and a coordinator can merge all +-- partial digests into a single mutable accumulator. Because the +-- t-digest is a mergeable sketch (Dunning, 2021; +-- ), the merged result has +-- accuracy comparable to a single-pass digest over the combined data. +-- +-- __Complexity.__ \(O(m)\) insertions where \(m\) is the centroid count +-- of the source digest, plus any triggered compressions. +merge :: TD.TDigest -> MDigest s -> ST s () +merge other md = do + let otherCompressed = TD.compress other + otherCs = TD.centroidList otherCompressed + mapM_ (\c -> addWeighted (TD.cMean c) (TD.cWeight c) md) otherCs + +-- --------------------------------------------------------------------------- +-- Freeze / Thaw +-- --------------------------------------------------------------------------- + +-- | Snapshot the mutable digest into a pure 'TD.TDigest'. +-- +-- The mutable digest is compressed first (flushing any buffered points), +-- then its centroids, total weight, extrema, and compression parameter +-- are read out and packaged into a pure 'TD.TDigest' via +-- 'TD.fromComponents'. +-- +-- The resulting pure digest is backed by a finger tree (Hinze & +-- Paterson, 2006; ) and +-- supports \(O(\log n)\) queries and further pure insertions. +-- +-- __Use case.__ 'freeze' is the primary exit path from a mutable +-- computation when the result must be returned to pure code or +-- serialised. It is also the mechanism for snapshotting a running +-- digest — the mutable digest remains usable after 'freeze'. +-- +-- __Complexity.__ \(O(\delta)\) for the compress plus a linear +-- traversal to extract centroids. +freeze :: MDigest s -> ST s TD.TDigest +freeze md = do + compress md + cc <- readSTRef (mdCentroidCount md) + centroids <- readSTRef (mdCentroids md) + cs <- readCentroids centroids cc 0 [] + tw <- readSTRef (mdTotalWeight md) + mn <- readSTRef (mdMin md) + mx <- readSTRef (mdMax md) + delta <- readSTRef (mdDelta md) + return (TD.fromComponents cs tw mn mx delta) + where + readCentroids _ 0 _ acc = return (reverse acc) + readCentroids v n i acc = do + (m, w) <- MV.read v i + readCentroids v (n - 1) (i + 1) (TD.Centroid m w : acc) + +-- | Create a mutable digest from a pure 'TD.TDigest'. +-- +-- The pure digest is compressed, its centroids are written into a fresh +-- mutable vector, and the scalar accumulators (total weight, min, max, +-- delta) are initialised from the pure digest's fields. Prefix sums +-- are rebuilt immediately. +-- +-- __Use case.__ 'thaw' is the entry path for converting a pure digest +-- (e.g., received from another thread or deserialised from storage) into +-- a mutable digest for continued high-throughput ingestion. In a +-- parallel/distributed pipeline, each worker can 'thaw' a shared seed +-- digest, ingest a partition of the data mutably, 'freeze' the result, +-- and return it for merging. +-- +-- __Complexity.__ \(O(\delta)\) for the copy and prefix-sum rebuild. +thaw :: TD.TDigest -> ST s (MDigest s) +thaw td = do + let td' = TD.compress td + cs = TD.centroidList td' + delta = TD.getDelta td' + md <- newWith delta + writeSTRef (mdTotalWeight md) (TD.totalWeight td') + writeSTRef (mdMin md) (TD.getMin td') + writeSTRef (mdMax md) (TD.getMax td') + let n = length cs + writeSTRef (mdCentroidCount md) n + centroids <- MV.new (max n 1) + writeCentroids centroids cs 0 + writeSTRef (mdCentroids md) centroids + rebuildPrefixSums md + return md + where + writeCentroids _ [] _ = return () + writeCentroids v (c : rest) i = do + MV.write v i (TD.cMean c, TD.cWeight c) + writeCentroids v rest (i + 1) + +-- --------------------------------------------------------------------------- +-- Convenience runner +-- --------------------------------------------------------------------------- + +-- | Run an 'ST' computation that uses a mutable t-digest and return the +-- pure result. +-- +-- This is a thin wrapper around 'Control.Monad.ST.runST'. The rank-2 +-- type @(forall s. 'ST' s a) -> a@ ensures that no mutable reference +-- (including the t'MDigest' itself, its internal 'STRef's, and its +-- 'Data.Vector.Mutable.MVector's) can escape the scope of the +-- computation. This guarantee is enforced statically by the Haskell +-- type checker via the universally quantified state token @s@ — any +-- attempt to return or store a value whose type mentions @s@ is a type +-- error. See Launchbury & Peyton Jones (1994), /Lazy Functional State +-- Threads/, for the theoretical foundation. +-- +-- __Usage pattern.__ Typically, one creates a digest with 'new' or +-- 'newWith', performs insertions with 'add' or 'addWeighted', and +-- extracts a result with 'quantile', 'cdf', or 'freeze' — all within +-- the 'runTDigest' block: +-- +-- @ +-- result :: Maybe Double +-- result = 'runTDigest' $ do +-- td <- 'new' +-- 'add' 42.0 td +-- 'quantile' 0.5 td +-- @ +runTDigest :: (forall s. ST s a) -> a +runTDigest = runST + +-- --------------------------------------------------------------------------- +-- Utility +-- --------------------------------------------------------------------------- + +clamp :: Double -> Double -> Double -> Double +clamp lo hi x + | x < lo = lo + | x > hi = hi + | otherwise = x