{-# LANGUAGE RebindableSyntax #-}

-- | Process quantile statistics.
module Circuit.Stats.Quantiles
  ( median,
    quantiles,
    digitize,
    signalize,
    OnlineTDigest (..),
    emptyOnlineTDigest,
    onlineInsert,
    onlineCompress,
    onlineForceCompress,
  )
where

import Circuit.Stats
import Data.Sketch.TDigest qualified as TD
import NumHask.Prelude hiding (fold)

-- | An online t-digest with exponential decay weighting.
--
-- Each inserted point receives weight @r ** (-(n+1))@ where @n@ counts
-- insertions since the last forced compression. Periodically the weights
-- are rescaled so that older points decay relative to newer ones.
data OnlineTDigest = OnlineTDigest
  { OnlineTDigest -> TDigest
td :: TD.TDigest,
    OnlineTDigest -> Int
tdN :: Int,
    OnlineTDigest -> Double
tdRate :: Double
  }
  deriving (Int -> OnlineTDigest -> ShowS
[OnlineTDigest] -> ShowS
OnlineTDigest -> String
(Int -> OnlineTDigest -> ShowS)
-> (OnlineTDigest -> String)
-> ([OnlineTDigest] -> ShowS)
-> Show OnlineTDigest
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> OnlineTDigest -> ShowS
showsPrec :: Int -> OnlineTDigest -> ShowS
$cshow :: OnlineTDigest -> String
show :: OnlineTDigest -> String
$cshowList :: [OnlineTDigest] -> ShowS
showList :: [OnlineTDigest] -> ShowS
Show)

-- | T-digest compression parameter. Delta=50 reproduces the accuracy profile
-- of the previous @tdigest@ package at its compression=25 setting.
tdigestDelta :: Double
tdigestDelta :: Double
tdigestDelta = Double
50

emptyOnlineTDigest :: Double -> OnlineTDigest
emptyOnlineTDigest :: Double -> OnlineTDigest
emptyOnlineTDigest = TDigest -> Int -> Double -> OnlineTDigest
OnlineTDigest (Double -> TDigest
TD.emptyWith Double
tdigestDelta) Int
0

-- | Process quantiles based on the t-digest library.
quantiles :: Double -> [Double] -> Process Double [Double]
quantiles :: Double -> [Double] -> Process Double [Double]
quantiles Double
r [Double]
qs = (Double -> OnlineTDigest)
-> (OnlineTDigest -> Double -> OnlineTDigest)
-> (OnlineTDigest -> [Double])
-> Process Double [Double]
forall s a b. (a -> s) -> (s -> a -> s) -> (s -> b) -> Process a b
Process Double -> OnlineTDigest
inject OnlineTDigest -> Double -> OnlineTDigest
step OnlineTDigest -> [Double]
extract
  where
    step :: OnlineTDigest -> Double -> OnlineTDigest
step OnlineTDigest
x Double
a = Double -> OnlineTDigest -> OnlineTDigest
onlineInsert Double
a OnlineTDigest
x
    inject :: Double -> OnlineTDigest
inject Double
a = Double -> OnlineTDigest -> OnlineTDigest
onlineInsert Double
a (Double -> OnlineTDigest
emptyOnlineTDigest Double
r)
    extract :: OnlineTDigest -> [Double]
extract OnlineTDigest
x = Double -> Maybe Double -> Double
forall a. a -> Maybe a -> a
fromMaybe (Double
0 Double -> Double -> Double
forall a. Divisive a => a -> a -> a
/ Double
0) (Maybe Double -> Double)
-> (Double -> Maybe Double) -> Double -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (Double -> TDigest -> Maybe Double
`TD.quantile` TDigest
t) (Double -> Double) -> [Double] -> [Double]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Double]
qs
      where
        (OnlineTDigest TDigest
t Int
_ Double
_) = OnlineTDigest -> OnlineTDigest
onlineForceCompress OnlineTDigest
x

-- | Process median using the t-digest algorithm.
--
-- The t-digest algorithm works best at extremes and can be unreliable in the centre.
median :: Double -> Process Double Double
median :: Double -> Process Double Double
median Double
r = (Double -> OnlineTDigest)
-> (OnlineTDigest -> Double -> OnlineTDigest)
-> (OnlineTDigest -> Double)
-> Process Double Double
forall s a b. (a -> s) -> (s -> a -> s) -> (s -> b) -> Process a b
Process Double -> OnlineTDigest
inject OnlineTDigest -> Double -> OnlineTDigest
step OnlineTDigest -> Double
extract
  where
    step :: OnlineTDigest -> Double -> OnlineTDigest
step OnlineTDigest
x Double
a = Double -> OnlineTDigest -> OnlineTDigest
onlineInsert Double
a OnlineTDigest
x
    inject :: Double -> OnlineTDigest
inject Double
a = Double -> OnlineTDigest -> OnlineTDigest
onlineInsert Double
a (Double -> OnlineTDigest
emptyOnlineTDigest Double
r)
    extract :: OnlineTDigest -> Double
extract OnlineTDigest
x = Double -> Maybe Double -> Double
forall a. a -> Maybe a -> a
fromMaybe (Double
0 Double -> Double -> Double
forall a. Divisive a => a -> a -> a
/ Double
0) (Double -> TDigest -> Maybe Double
TD.quantile Double
0.5 TDigest
t)
      where
        (OnlineTDigest TDigest
t Int
_ Double
_) = OnlineTDigest -> OnlineTDigest
onlineForceCompress OnlineTDigest
x

onlineInsert' :: Double -> OnlineTDigest -> OnlineTDigest
onlineInsert' :: Double -> OnlineTDigest -> OnlineTDigest
onlineInsert' Double
x (OnlineTDigest TDigest
td' Int
n Double
r) =
  TDigest -> Int -> Double -> OnlineTDigest
OnlineTDigest
    (Double -> Double -> TDigest -> TDigest
TD.addWeighted Double
x (Double
r Double -> Integer -> Double
forall b a.
(Ord b, Divisive a, Subtractive b, Integral b) =>
a -> b -> a
^^ (-(Int -> Integer
forall a b. FromIntegral a b => b -> a
fromIntegral (Int -> Integer) -> Int -> Integer
forall a b. (a -> b) -> a -> b
$ Int
n Int -> Int -> Int
forall a. Additive a => a -> a -> a
+ Int
1) :: Integer)) TDigest
td')
    (Int
n Int -> Int -> Int
forall a. Additive a => a -> a -> a
+ Int
1)
    Double
r

onlineInsert :: Double -> OnlineTDigest -> OnlineTDigest
onlineInsert :: Double -> OnlineTDigest -> OnlineTDigest
onlineInsert Double
x OnlineTDigest
otd = OnlineTDigest -> OnlineTDigest
onlineCompress (Double -> OnlineTDigest -> OnlineTDigest
onlineInsert' Double
x OnlineTDigest
otd)

-- | Force a compression pass when the digest has grown enough that the
-- decay weights are becoming numerically awkward. The threshold is chosen
-- to keep weights within a comfortable double range while preserving the
-- exponential-decay semantics.
onlineCompress :: OnlineTDigest -> OnlineTDigest
onlineCompress :: OnlineTDigest -> OnlineTDigest
onlineCompress otd :: OnlineTDigest
otd@(OnlineTDigest TDigest
_ Int
n Double
_)
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
maxInsertsSinceCompress = OnlineTDigest -> OnlineTDigest
onlineForceCompress OnlineTDigest
otd
  | Bool
otherwise = OnlineTDigest
otd

maxInsertsSinceCompress :: Int
maxInsertsSinceCompress :: Int
maxInsertsSinceCompress = Int
200

-- | Rescale all centroid weights by @r ** n@ and reset the insertion counter.
-- This is the exponential-decay equivalent of normalising weights.
onlineForceCompress :: OnlineTDigest -> OnlineTDigest
onlineForceCompress :: OnlineTDigest -> OnlineTDigest
onlineForceCompress (OnlineTDigest TDigest
t Int
n Double
r) = TDigest -> Int -> Double -> OnlineTDigest
OnlineTDigest TDigest
t' Int
0 Double
r
  where
    t' :: TDigest
t' = TDigest -> TDigest
TD.compress (TDigest -> TDigest) -> TDigest -> TDigest
forall a b. (a -> b) -> a -> b
$ (TDigest -> Centroid -> TDigest)
-> TDigest -> [Centroid] -> TDigest
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' TDigest -> Centroid -> TDigest
rebuild (Double -> TDigest
TD.emptyWith Double
tdigestDelta) (TDigest -> [Centroid]
TD.centroidList TDigest
t)
    rebuild :: TDigest -> Centroid -> TDigest
rebuild TDigest
a Centroid
c = Double -> Double -> TDigest -> TDigest
TD.addWeighted (Centroid -> Double
TD.cMean Centroid
c) (Centroid -> Double
TD.cWeight Centroid
c Double -> Double -> Double
forall a. Multiplicative a => a -> a -> a
* Double
r Double -> Int -> Double
forall b a.
(Ord b, Divisive a, Subtractive b, Integral b) =>
a -> b -> a
^^ Int
n) TDigest
a

-- | A process that computes the running quantile bucket. For example,
-- in a scan, @digitize 0.9 [0.5]@ returns:
--
-- * 0 if the current value is less than the current process median.
--
-- * 1 if the current value is greater than the current process median.
digitize :: Double -> [Double] -> Process Double Int
digitize :: Double -> [Double] -> Process Double Int
digitize Double
r [Double]
qs = (Double -> (OnlineTDigest, Double))
-> ((OnlineTDigest, Double) -> Double -> (OnlineTDigest, Double))
-> ((OnlineTDigest, Double) -> Int)
-> Process Double Int
forall s a b. (a -> s) -> (s -> a -> s) -> (s -> b) -> Process a b
Process Double -> (OnlineTDigest, Double)
inject (OnlineTDigest, Double) -> Double -> (OnlineTDigest, Double)
forall {b}. (OnlineTDigest, b) -> Double -> (OnlineTDigest, Double)
step (OnlineTDigest, Double) -> Int
extract
  where
    step :: (OnlineTDigest, b) -> Double -> (OnlineTDigest, Double)
step (OnlineTDigest
x, b
_) Double
a = (Double -> OnlineTDigest -> OnlineTDigest
onlineInsert Double
a OnlineTDigest
x, Double
a)
    inject :: Double -> (OnlineTDigest, Double)
inject Double
a = (Double -> OnlineTDigest -> OnlineTDigest
onlineInsert Double
a (Double -> OnlineTDigest
emptyOnlineTDigest Double
r), Double
a)
    extract :: (OnlineTDigest, Double) -> Int
extract (OnlineTDigest
x, Double
l) = [Double] -> Double -> Int
forall {a} {a}. (FromInteger a, Additive a, Ord a) => [a] -> a -> a
bucket' [Double]
qs' Double
l
      where
        qs' :: [Double]
qs' = Double -> Maybe Double -> Double
forall a. a -> Maybe a -> a
fromMaybe (Double
0 Double -> Double -> Double
forall a. Divisive a => a -> a -> a
/ Double
0) (Maybe Double -> Double)
-> (Double -> Maybe Double) -> Double -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (Double -> TDigest -> Maybe Double
`TD.quantile` TDigest
t) (Double -> Double) -> [Double] -> [Double]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Double]
qs
        (OnlineTDigest TDigest
t Int
_ Double
_) = OnlineTDigest -> OnlineTDigest
onlineForceCompress OnlineTDigest
x
        bucket' :: [a] -> a -> a
bucket' [a]
xs a
l' =
          a -> Maybe a -> a
forall a. a -> Maybe a -> a
fromMaybe a
0
            (Maybe a -> a) -> ([a] -> Maybe a) -> [a] -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Process a a -> [a] -> Maybe a
forall a b. Process a b -> [a] -> Maybe b
fold ((a -> a) -> (a -> a -> a) -> (a -> a) -> Process a a
forall s a b. (a -> s) -> (s -> a -> s) -> (s -> b) -> Process a b
Process a -> a
forall a. a -> a
forall {k} (cat :: k -> k -> *) (a :: k). Category cat => cat a a
id a -> a -> a
forall a. Additive a => a -> a -> a
(+) a -> a
forall a. a -> a
forall {k} (cat :: k -> k -> *) (a :: k). Category cat => cat a a
id)
            ([a] -> a) -> [a] -> a
forall a b. (a -> b) -> a -> b
$ ( \a
x' ->
                  if a
x' a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
l'
                    then a
0
                    else a
1
              )
              (a -> a) -> [a] -> [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [a]
xs

-- | transform an input to a [0,1] signal, via digitalization.
signalize :: Double -> [Double] -> Process Double Double
signalize :: Double -> [Double] -> Process Double Double
signalize Double
r [Double]
qs' =
  Process Double Int -> (Int -> Double) -> Process Double Double
forall a b c. Process a b -> (b -> c) -> Process a c
after (Double -> [Double] -> Process Double Int
digitize Double
r [Double]
qs') (\Int
x -> Int -> Double
forall a b. FromIntegral a b => b -> a
fromIntegral Int
x Double -> Double -> Double
forall a. Divisive a => a -> a -> a
/ Int -> Double
forall a b. FromIntegral a b => b -> a
fromIntegral ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
qs' Int -> Int -> Int
forall a. Additive a => a -> a -> a
+ Int
1))