{-|
Copyright   : Predictable Network Solutions Ltd., 2020-2024
License     : BSD-3-Clause
Description : Discrete, finite signed measures on the number line.
-}
module Numeric.Measure.Discrete
    ( -- * Type
      Discrete
    , fromMap
    , toMap
    , zero
    , dirac
    , distribution

    -- * Observations
    , total
    , integrate

    -- * Operations, numerical
    , add
    , scale
    , translate
    , convolve
    ) where

import Data.List
    ( scanl'
    )
import Data.Map
    ( Map
    )
import Numeric.Function.Piecewise
    ( Piecewise
    )
import Numeric.Polynomial.Simple
    ( Poly
    )

import qualified Data.Map.Strict as Map
import qualified Numeric.Function.Piecewise as Piecewise
import qualified Numeric.Polynomial.Simple as Poly

{-----------------------------------------------------------------------------
    Type
------------------------------------------------------------------------------}
-- | A discrete, finite
-- [signed measure](https://en.wikipedia.org/wiki/Signed_measure)
-- on the number line.
newtype Discrete a = Discrete (Map a a)
    -- INVARIANT: All values are non-zero.
    deriving (Int -> Discrete a -> ShowS
[Discrete a] -> ShowS
Discrete a -> String
(Int -> Discrete a -> ShowS)
-> (Discrete a -> String)
-> ([Discrete a] -> ShowS)
-> Show (Discrete a)
forall a. Show a => Int -> Discrete a -> ShowS
forall a. Show a => [Discrete a] -> ShowS
forall a. Show a => Discrete a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> Discrete a -> ShowS
showsPrec :: Int -> Discrete a -> ShowS
$cshow :: forall a. Show a => Discrete a -> String
show :: Discrete a -> String
$cshowList :: forall a. Show a => [Discrete a] -> ShowS
showList :: [Discrete a] -> ShowS
Show)

-- | Internal.
-- Remove all zero values.
trim :: (Ord a, Num a) => Map a a -> Map a a
trim :: forall a. (Ord a, Num a) => Map a a -> Map a a
trim Map a a
m = (a -> Bool) -> Map a a -> Map a a
forall a k. (a -> Bool) -> Map k a -> Map k a
Map.filter (a -> a -> Bool
forall a. Eq a => a -> a -> Bool
/= a
0) Map a a
m

-- | Two measures are equal if they yield the same measures on every set.
--
-- > mx == my
-- >   implies
-- >   forall t. eval (distribution mx) t = eval (distribution my) t
instance (Ord a, Num a) => Eq (Discrete a) where
    Discrete Map a a
mx == :: Discrete a -> Discrete a -> Bool
== Discrete Map a a
my = Map a a
mx Map a a -> Map a a -> Bool
forall a. Eq a => a -> a -> Bool
== Map a a
my

{-----------------------------------------------------------------------------
    Operations
------------------------------------------------------------------------------}
-- | The measure that assigns @0@ to every set.
zero :: Num a => Discrete a
zero :: forall a. Num a => Discrete a
zero = Map a a -> Discrete a
forall a. Map a a -> Discrete a
Discrete Map a a
forall k a. Map k a
Map.empty

-- | A
-- [Dirac measure](https://en.wikipedia.org/wiki/Dirac_measure)
-- at the given point @x@.
--
-- > total (dirac x) = 1
dirac :: (Ord a, Num a) => a -> Discrete a
dirac :: forall a. (Ord a, Num a) => a -> Discrete a
dirac a
x = Map a a -> Discrete a
forall a. Map a a -> Discrete a
Discrete (a -> a -> Map a a
forall k a. k -> a -> Map k a
Map.singleton a
x a
1)

-- | Construct a discrete measure
-- from a collection of points and their measures.
fromMap :: (Ord a, Num a) => Map a a -> Discrete a
fromMap :: forall a. (Ord a, Num a) => Map a a -> Discrete a
fromMap = Map a a -> Discrete a
forall a. Map a a -> Discrete a
Discrete (Map a a -> Discrete a)
-> (Map a a -> Map a a) -> Map a a -> Discrete a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map a a -> Map a a
forall a. (Ord a, Num a) => Map a a -> Map a a
trim

-- | Decompose the discrete measure into a collection of points
-- and their measures.
toMap :: Num a => Discrete a -> Map a a
toMap :: forall a. Num a => Discrete a -> Map a a
toMap (Discrete Map a a
m) = Map a a
m

-- | The total of the measure applied to the set of real numbers.
total :: Num a => Discrete a -> a
total :: forall a. Num a => Discrete a -> a
total (Discrete Map a a
m) = Map a a -> a
forall a. Num a => Map a a -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum Map a a
m

-- | Integrate a function @f@ with respect to the given measure @m@,
-- \( \int f(x) dm(x) \).
integrate :: (Ord a, Num a) => (a -> a) -> Discrete a -> a
integrate :: forall a. (Ord a, Num a) => (a -> a) -> Discrete a -> a
integrate a -> a
f (Discrete Map a a
m) = Map a a -> a
forall a. Num a => Map a a -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (Map a a -> a) -> Map a a -> a
forall a b. (a -> b) -> a -> b
$ (a -> a -> a) -> Map a a -> Map a a
forall k a b. (k -> a -> b) -> Map k a -> Map k b
Map.mapWithKey (\a
x a
w -> a -> a
f a
x a -> a -> a
forall a. Num a => a -> a -> a
* a
w) Map a a
m

-- | @eval (distribution m) x@ is the measure of the interval \( (-∞, x] \).
--
-- This is known as the [distribution function
-- ](https://en.wikipedia.org/wiki/Distribution_function_%28measure_theory%29).
distribution :: (Ord a, Num a) => Discrete a -> Piecewise (Poly a)
distribution :: forall a. (Ord a, Num a) => Discrete a -> Piecewise (Poly a)
distribution (Discrete Map a a
m) =
    [(Domain (Poly a), Poly a)] -> Piecewise (Poly a)
forall o. Ord (Domain o) => [(Domain o, o)] -> Piecewise o
Piecewise.fromAscPieces
    ([(Domain (Poly a), Poly a)] -> Piecewise (Poly a))
-> [(Domain (Poly a), Poly a)] -> Piecewise (Poly a)
forall a b. (a -> b) -> a -> b
$ ((a, a) -> a -> (a, Poly a)) -> [(a, a)] -> [a] -> [(a, Poly a)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\(a
x,a
_) a
s -> (a
x,a -> Poly a
forall a. a -> Poly a
Poly.constant a
s)) [(a, a)]
diracs [a]
steps
  where
    diracs :: [(a, a)]
diracs = Map a a -> [(a, a)]
forall k a. Map k a -> [(k, a)]
Map.toAscList Map a a
m
    steps :: [a]
steps = [a] -> [a]
forall a. HasCallStack => [a] -> [a]
tail ([a] -> [a]) -> [a] -> [a]
forall a b. (a -> b) -> a -> b
$ (a -> a -> a) -> a -> [a] -> [a]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl' a -> a -> a
forall a. Num a => a -> a -> a
(+) a
0 ([a] -> [a]) -> [a] -> [a]
forall a b. (a -> b) -> a -> b
$ ((a, a) -> a) -> [(a, a)] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (a, a) -> a
forall a b. (a, b) -> b
snd [(a, a)]
diracs

-- | Add two measures.
--
-- > total (add mx my) = total mx + total my
add :: (Ord a, Num a) => Discrete a -> Discrete a -> Discrete a
add :: forall a. (Ord a, Num a) => Discrete a -> Discrete a -> Discrete a
add (Discrete Map a a
mx) (Discrete Map a a
my) =
    Map a a -> Discrete a
forall a. Map a a -> Discrete a
Discrete (Map a a -> Discrete a) -> Map a a -> Discrete a
forall a b. (a -> b) -> a -> b
$ Map a a -> Map a a
forall a. (Ord a, Num a) => Map a a -> Map a a
trim (Map a a -> Map a a) -> Map a a -> Map a a
forall a b. (a -> b) -> a -> b
$ (a -> a -> a) -> Map a a -> Map a a -> Map a a
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith a -> a -> a
forall a. Num a => a -> a -> a
(+) Map a a
mx Map a a
my

-- | Scale a measure by a constant.
--
-- > total (scale a mx) = a * total mx
scale :: (Ord a, Num a) => a -> Discrete a -> Discrete a
scale :: forall a. (Ord a, Num a) => a -> Discrete a -> Discrete a
scale a
0 (Discrete Map a a
_) = Map a a -> Discrete a
forall a. Map a a -> Discrete a
Discrete Map a a
forall k a. Map k a
Map.empty
scale a
s (Discrete Map a a
m) = Map a a -> Discrete a
forall a. Map a a -> Discrete a
Discrete (Map a a -> Discrete a) -> Map a a -> Discrete a
forall a b. (a -> b) -> a -> b
$ (a -> a) -> Map a a -> Map a a
forall a b k. (a -> b) -> Map k a -> Map k b
Map.map (a
s a -> a -> a
forall a. Num a => a -> a -> a
*) Map a a
m

-- | Translate a measure along the number line.
--
-- > eval (distribution (translate y m)) x
-- >    = eval (distribution m) (x - y)
translate :: (Ord a, Num a) => a -> Discrete a -> Discrete a
translate :: forall a. (Ord a, Num a) => a -> Discrete a -> Discrete a
translate a
y (Discrete Map a a
m) = Map a a -> Discrete a
forall a. Map a a -> Discrete a
Discrete (Map a a -> Discrete a) -> Map a a -> Discrete a
forall a b. (a -> b) -> a -> b
$ (a -> a) -> Map a a -> Map a a
forall k2 k1 a. Ord k2 => (k1 -> k2) -> Map k1 a -> Map k2 a
Map.mapKeys (a
y a -> a -> a
forall a. Num a => a -> a -> a
+) Map a a
m

-- | Additive convolution of two measures.
--
-- Properties:
--
-- > convolve (dirac x) (dirac y) = dirac (x + y)
convolve :: (Ord a, Num a) => Discrete a -> Discrete a -> Discrete a
-- >
-- > convolve mx my               =  convolve my mx
-- > convolve (add mx my) mz      =  add (convolve mx mz) (convolve my mz)
-- > translate z (convolve mx my) =  convolve (translate z mx) my
-- > total (convolve mx my)       =  total mx * total myconvolve :: (Ord a, Num a) => Discrete a -> Discrete a -> Discrete a
convolve :: forall a. (Ord a, Num a) => Discrete a -> Discrete a -> Discrete a
convolve (Discrete Map a a
mx) (Discrete Map a a
my) =
    Map a a -> Discrete a
forall a. Map a a -> Discrete a
Discrete (Map a a -> Discrete a) -> Map a a -> Discrete a
forall a b. (a -> b) -> a -> b
$ Map a a -> Map a a
forall a. (Ord a, Num a) => Map a a -> Map a a
trim (Map a a -> Map a a) -> Map a a -> Map a a
forall a b. (a -> b) -> a -> b
$ (a -> a -> a) -> [(a, a)] -> Map a a
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith a -> a -> a
forall a. Num a => a -> a -> a
(+)
        [ (a
x a -> a -> a
forall a. Num a => a -> a -> a
+ a
y, a
wx a -> a -> a
forall a. Num a => a -> a -> a
* a
wy)
        | (a
x,a
wx) <- Map a a -> [(a, a)]
forall k a. Map k a -> [(k, a)]
Map.toList Map a a
mx
        , (a
y,a
wy) <- Map a a -> [(a, a)]
forall k a. Map k a -> [(k, a)]
Map.toList Map a a
my
        ]