{-# LANGUAGE DataKinds #-}
{-# LANGUAGE RebindableSyntax #-}
{-# OPTIONS_GHC -fno-warn-name-shadowing #-}

-- | A histogram, if you squint, is a series of contiguous 'Range's, annotated with values.
module NumHask.Space.Histogram
  ( Histogram (..),
    emptyHistogram,
    DealOvers (..),
    fill,
    cutI,
    regular,
    makeRects,
    regularQuantiles,
    quantileFold,
    freq,
    average,
    quantiles,
    quantile,
  )
where

import Data.Map qualified as Map
import Data.TDigest qualified as TD
import Data.Vector qualified as V
import NumHask.Prelude
import NumHask.Space.Range
import NumHask.Space.Rect
import NumHask.Space.Types

-- $setup
-- >>> :set -XRebindableSyntax
-- >>> import NumHask.Prelude
-- >>> import NumHask.Space

-- | This Histogram is a list of contiguous boundaries (a boundary being the lower edge of one bucket and the upper edge of another), and a value (usually a count) for each bucket, represented here as a map
--
-- Overs and Unders are contained in key = 0 and key = length cuts
-- Intervals are defined as (l,u]
data Histogram = Histogram
  { Histogram -> Vector Double
cuts :: V.Vector Double, -- bucket boundaries
    Histogram -> Map Int Double
values :: Map.Map Int Double -- bucket counts
  }
  deriving (Int -> Histogram -> ShowS
[Histogram] -> ShowS
Histogram -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Histogram] -> ShowS
$cshowList :: [Histogram] -> ShowS
show :: Histogram -> String
$cshow :: Histogram -> String
showsPrec :: Int -> Histogram -> ShowS
$cshowsPrec :: Int -> Histogram -> ShowS
Show, Histogram -> Histogram -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Histogram -> Histogram -> Bool
$c/= :: Histogram -> Histogram -> Bool
== :: Histogram -> Histogram -> Bool
$c== :: Histogram -> Histogram -> Bool
Eq)

-- | A histogram with no cuts nor data.
emptyHistogram :: Histogram
emptyHistogram :: Histogram
emptyHistogram = Vector Double -> Map Int Double -> Histogram
Histogram forall a. Vector a
V.empty forall k a. Map k a
Map.empty

-- | Whether or not to ignore unders and overs.  If overs and unders are dealt with, IncludeOvers supplies an assumed width for the outer buckets.
data DealOvers = IgnoreOvers | IncludeOvers Double

-- | Fill a Histogram using pre-specified cuts
--
-- >>> fill [0,50,100] [0..99]
-- Histogram {cuts = [0.0,50.0,100.0], values = fromList [(1,50.0),(2,50.0)]}
fill :: (Foldable f) => [Double] -> f Double -> Histogram
fill :: forall (f :: * -> *).
Foldable f =>
[Double] -> f Double -> Histogram
fill [Double]
cs f Double
xs =
  Vector Double -> Map Int Double -> Histogram
Histogram
    (forall a. [a] -> Vector a
V.fromList [Double]
cs)
    (forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Map Int Double
x Double
a -> forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith forall a. Additive a => a -> a -> a
(+) (forall a. Ord a => Vector a -> a -> Int
cutI (forall a. [a] -> Vector a
V.fromList [Double]
cs) Double
a) Double
1 Map Int Double
x) forall k a. Map k a
Map.empty f Double
xs)

-- | find the index of the bucket the value is contained in.
cutI :: (Ord a) => V.Vector a -> a -> Int
cutI :: forall a. Ord a => Vector a -> a -> Int
cutI Vector a
cs a
a = Range Int -> Int
go (forall a. a -> a -> Range a
Range forall a. Additive a => a
zero (forall a. Vector a -> Int
V.length Vector a
cs))
  where
    go :: Range Int -> Int
go (Range Int
l Int
u) =
      let k :: Int
k = (Int
u forall a. Additive a => a -> a -> a
+ Int
l) forall a. Integral a => a -> a -> a
`div` Int
2
       in case forall a. Ord a => a -> a -> Ordering
compare a
a (Vector a
cs forall a. Vector a -> Int -> a
V.! Int
k) of
            Ordering
EQ -> Int
k forall a. Additive a => a -> a -> a
+ Int
1
            Ordering
LT -> forall a. a -> a -> Bool -> a
bool (Range Int -> Int
go (forall a. a -> a -> Range a
Range Int
l Int
k)) Int
k (Int
l forall a. Eq a => a -> a -> Bool
== Int
k)
            Ordering
GT ->
              forall a. a -> a -> Bool -> a
bool
                ( case forall a. Ord a => a -> a -> Ordering
compare a
a (Vector a
cs forall a. Vector a -> Int -> a
V.! (Int
k forall a. Additive a => a -> a -> a
+ forall a. Multiplicative a => a
one)) of
                    Ordering
EQ -> Int
k forall a. Additive a => a -> a -> a
+ Int
2
                    Ordering
LT -> Int
k forall a. Additive a => a -> a -> a
+ Int
1
                    Ordering
GT -> Range Int -> Int
go (forall a. a -> a -> Range a
Range Int
k Int
u)
                )
                (Int
k forall a. Additive a => a -> a -> a
+ Int
1)
                (Int
k forall a. Ord a => a -> a -> Bool
>= Int
u forall a. Subtractive a => a -> a -> a
- forall a. Multiplicative a => a
one)

-- | Make a histogram using n equally spaced cuts over the entire range of the data
--
-- >>> regular 4 [0..100]
-- Histogram {cuts = [0.0,25.0,50.0,75.0,100.0], values = fromList [(1,25.0),(2,25.0),(3,25.0),(4,25.0),(5,1.0)]}
regular :: Int -> [Double] -> Histogram
regular :: Int -> [Double] -> Histogram
regular Int
_ [] = Histogram
emptyHistogram
regular Int
n [Double]
xs = forall (f :: * -> *).
Foldable f =>
[Double] -> f Double -> Histogram
fill [Double]
cs [Double]
xs
  where
    cs :: [Element (Range Double)]
cs = forall s. FieldSpace s => Pos -> s -> Grid s -> [Element s]
grid Pos
OuterPos (forall s (f :: * -> *).
(Space s, Traversable f) =>
f (Element s) -> s
unsafeSpace1 [Double]
xs :: Range Double) Int
n

-- | Transform a Histogram to Rects
--
-- >>> makeRects IgnoreOvers (regular 4 [0..100])
-- [Rect 0.0 25.0 0.0 0.25,Rect 25.0 50.0 0.0 0.25,Rect 50.0 75.0 0.0 0.25,Rect 75.0 100.0 0.0 0.25]
makeRects :: DealOvers -> Histogram -> [Rect Double]
makeRects :: DealOvers -> Histogram -> [Rect Double]
makeRects DealOvers
o (Histogram Vector Double
cs Map Int Double
counts) = forall a. Vector a -> [a]
V.toList forall a b. (a -> b) -> a -> b
$ forall a b c d.
(a -> b -> c -> d) -> Vector a -> Vector b -> Vector c -> Vector d
V.zipWith3 (\Double
x Double
z Double
w' -> forall a. a -> a -> a -> a -> Rect a
Rect Double
x Double
z forall a. Additive a => a
zero Double
w') Vector Double
x Vector Double
z Vector Double
w'
  where
    w :: Vector Double
w =
      forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
V.zipWith
        forall a. Divisive a => a -> a -> a
(/)
        ((\Int
x' -> forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault Double
0 Int
x' Map Int Double
counts) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Num a => a -> Int -> Vector a
V.enumFromN Int
f (Int
l forall a. Subtractive a => a -> a -> a
- Int
f forall a. Additive a => a -> a -> a
+ forall a. Multiplicative a => a
one))
        (forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
V.zipWith (-) Vector Double
z Vector Double
x)
    f :: Int
f = case DealOvers
o of
      DealOvers
IgnoreOvers -> forall a. Multiplicative a => a
one
      IncludeOvers Double
_ -> forall a. Additive a => a
zero
    l :: Int
l = case DealOvers
o of
      DealOvers
IgnoreOvers -> forall (t :: * -> *) a. Foldable t => t a -> Int
length Vector Double
cs forall a. Subtractive a => a -> a -> a
- forall a. Multiplicative a => a
one
      IncludeOvers Double
_ -> forall (t :: * -> *) a. Foldable t => t a -> Int
length Vector Double
cs
    w' :: Vector Double
w' = (forall a. Divisive a => a -> a -> a
/ forall a (f :: * -> *). (Additive a, Foldable f) => f a -> a
sum Vector Double
w) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Vector Double
w
    x :: Vector Double
x = case DealOvers
o of
      DealOvers
IgnoreOvers -> Vector Double
cs
      IncludeOvers Double
outw ->
        forall a. a -> Vector a
V.singleton (forall a. Vector a -> a
V.head Vector Double
cs forall a. Subtractive a => a -> a -> a
- Double
outw)
          forall a. Semigroup a => a -> a -> a
<> Vector Double
cs
          forall a. Semigroup a => a -> a -> a
<> forall a. a -> Vector a
V.singleton (forall a. Vector a -> a
V.last Vector Double
cs forall a. Additive a => a -> a -> a
+ Double
outw)
    z :: Vector Double
z = forall a. Int -> Vector a -> Vector a
V.drop forall a. Multiplicative a => a
one Vector Double
x

-- | approx regular n-quantiles
--
-- >>> regularQuantiles 4 [0..100]
-- [0.0,24.75,50.0,75.25,100.0]
regularQuantiles :: Double -> [Double] -> [Double]
regularQuantiles :: Double -> [Double] -> [Double]
regularQuantiles Double
n [Double]
xs = [Double] -> [Double] -> [Double]
quantileFold [Double]
qs [Double]
xs
  where
    qs :: [Double]
qs = ((Double
1 forall a. Divisive a => a -> a -> a
/ Double
n) *) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Double
0 .. Double
n]

-- | one-pass approximate quantiles fold
quantileFold :: [Double] -> [Double] -> [Double]
quantileFold :: [Double] -> [Double] -> [Double]
quantileFold [Double]
qs [Double]
xs = forall {comp :: Nat}. KnownNat comp => TDigest comp -> [Double]
done forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall {comp :: Nat}.
KnownNat comp =>
TDigest comp -> Double -> TDigest comp
step TDigest 25
begin [Double]
xs
  where
    step :: TDigest comp -> Double -> TDigest comp
step TDigest comp
x Double
a = forall (comp :: Nat).
KnownNat comp =>
Double -> TDigest comp -> TDigest comp
TD.insert Double
a TDigest comp
x
    begin :: TDigest 25
begin = forall (f :: * -> *) (comp :: Nat).
(Foldable f, KnownNat comp) =>
f Double -> TDigest comp
TD.tdigest ([] :: [Double]) :: TD.TDigest 25
    done :: TDigest comp -> [Double]
done TDigest comp
x = forall a. a -> Maybe a -> a
fromMaybe (Double
0 forall a. Divisive a => a -> a -> a
/ Double
0) forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (forall (comp :: Nat). Double -> TDigest comp -> Maybe Double
`TD.quantile` forall (comp :: Nat). KnownNat comp => TDigest comp -> TDigest comp
TD.compress TDigest comp
x) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Double]
qs

-- | normalize a histogram
--
-- > \h -> sum (values $ freq h) == one
--
-- >>> freq $ fill [0,50,100] [0..99]
-- Histogram {cuts = [0.0,50.0,100.0], values = fromList [(1,0.5),(2,0.5)]}
freq :: Histogram -> Histogram
freq :: Histogram -> Histogram
freq (Histogram Vector Double
cs Map Int Double
vs) = Vector Double -> Map Int Double -> Histogram
Histogram Vector Double
cs forall a b. (a -> b) -> a -> b
$ forall a b k. (a -> b) -> Map k a -> Map k b
Map.map (forall a. Multiplicative a => a -> a -> a
* forall a. Divisive a => a -> a
recip (forall a (f :: * -> *). (Additive a, Foldable f) => f a -> a
sum Map Int Double
vs)) Map Int Double
vs

-- | average
--
-- >>> average [0..1000]
-- 500.0
average :: (Foldable f) => f Double -> Double
average :: forall (f :: * -> *). Foldable f => f Double -> Double
average f Double
xs = forall a (f :: * -> *). (Additive a, Foldable f) => f a -> a
sum f Double
xs forall a. Divisive a => a -> a -> a
/ forall a b. FromIntegral a b => b -> a
fromIntegral (forall (t :: * -> *) a. Foldable t => t a -> Int
length f Double
xs)

-- | Regularly spaced (approx) quantiles
--
-- >>> quantiles 5 [1..1000]
-- [1.0,200.5,400.5,600.5000000000001,800.5,1000.0]
quantiles :: (Foldable f) => Int -> f Double -> [Double]
quantiles :: forall (f :: * -> *). Foldable f => Int -> f Double -> [Double]
quantiles Int
n f Double
xs =
  ( \Double
x ->
      forall a. a -> Maybe a -> a
fromMaybe Double
0 forall a b. (a -> b) -> a -> b
$
        forall (comp :: Nat). Double -> TDigest comp -> Maybe Double
TD.quantile Double
x (forall (f :: * -> *) (comp :: Nat).
(Foldable f, KnownNat comp) =>
f Double -> TDigest comp
TD.tdigest f Double
xs :: TD.TDigest 25)
  )
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((forall a. Divisive a => a -> a -> a
/ forall a b. FromIntegral a b => b -> a
fromIntegral Int
n) forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall a b. FromIntegral a b => b -> a
fromIntegral forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int
0 .. Int
n])

-- | single (approx) quantile
--
-- >>> quantile 0.1 [1..1000]
-- 100.5
quantile :: (Foldable f) => Double -> f Double -> Double
quantile :: forall (f :: * -> *). Foldable f => Double -> f Double -> Double
quantile Double
p f Double
xs = forall a. a -> Maybe a -> a
fromMaybe Double
0 forall a b. (a -> b) -> a -> b
$ forall (comp :: Nat). Double -> TDigest comp -> Maybe Double
TD.quantile Double
p (forall (f :: * -> *) (comp :: Nat).
(Foldable f, KnownNat comp) =>
f Double -> TDigest comp
TD.tdigest f Double
xs :: TD.TDigest 25)