{- | Handy distributions for the `LazyPPL` library, based on the `uniform` distribution. Mostly defined using the `Statistics.Distribution` module and family. 

Sometimes both a distribution (type @Prob a@) and pdf (type @a -> Double@) are given. Distributions are useful for sampling, densities are used for scoring. 

For more distributions, see the Statistics.Distribution in the statistics package. -}


module LazyPPL.Distributions (
       -- * Continuous distributions
       normal,normalPdf,exponential,expPdf,gamma, beta, dirichlet, uniformbounded,
       -- * Discrete distributions
       bernoulli, uniformdiscrete, categorical, poisson, poissonPdf,
       -- * Streams
       iid)
       where

import LazyPPL (Prob,uniform)
import Data.List (findIndex)
import Numeric.SpecFunctions
import Numeric.MathFunctions.Constants

{-|
  [Normal distribution](https://en.wikipedia.org/wiki/Normal_distribution)
-}
normal :: Double -- ^ mu, mean
       -> Double -- ^ sigma, standard deviation
       -> Prob Double
normal :: Double -> Double -> Prob Double
normal Double
m Double
s = do 
  Double
x <- Prob Double
uniform
  Double -> Prob Double
forall a. a -> Prob a
forall (m :: * -> *) a. Monad m => a -> m a
return (Double -> Prob Double) -> Double -> Prob Double
forall a b. (a -> b) -> a -> b
$ (- Double -> Double
invErfc (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x)) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
m_sqrt_2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
s) Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
m

normalPdf :: Double -> Double -> Double -> Double
normalPdf :: Double -> Double -> Double -> Double
normalPdf Double
m Double
s Double
x = Double -> Double
forall a. Floating a => a -> a
exp ((-(Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
m) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
-Double
m) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
s Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
s)) Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double -> Double
forall a. Floating a => a -> a
log (Double
m_sqrt_2_pi Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
s))


{-|
  [Exponential distribution](https://en.wikipedia.org/wiki/Exponential_distribution)
-}
exponential :: Double -- ^ lambda, rate
            -> Prob Double
exponential :: Double -> Prob Double
exponential Double
rate = do 
  Double
x <- Prob Double
uniform
  Double -> Prob Double
forall a. a -> Prob a
forall (m :: * -> *) a. Monad m => a -> m a
return (Double -> Prob Double) -> Double -> Prob Double
forall a b. (a -> b) -> a -> b
$ - (Double -> Double
forall a. Floating a => a -> a
log Double
x Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
rate)

expPdf :: Double -> Double -> Double
expPdf :: Double -> Double -> Double
expPdf Double
rate Double
x = Double -> Double
forall a. Floating a => a -> a
exp (-Double
rateDouble -> Double -> Double
forall a. Num a => a -> a -> a
*Double
x) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
rate

{-|
  [Gamma distribution](https://en.wikipedia.org/wiki/Gamma_distribution)
-}
gamma :: Double -- ^ k, shape
      -> Double -- ^ theta, scale
      -> Prob Double
gamma :: Double -> Double -> Prob Double
gamma Double
a Double
b = do
  Double
x <- Prob Double
uniform
  Double -> Prob Double
forall a. a -> Prob a
forall (m :: * -> *) a. Monad m => a -> m a
return (Double -> Prob Double) -> Double -> Prob Double
forall a b. (a -> b) -> a -> b
$ Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double -> Double
invIncompleteGamma Double
a Double
x

{-|
  [Beta distribution](https://en.wikipedia.org/wiki/Beta_distribution)
-}
beta :: Double -- ^ alpha
     -> Double -- ^ beta
     -> Prob Double
beta :: Double -> Double -> Prob Double
beta Double
a Double
b = do
  Double
x <- Prob Double
uniform
  Double -> Prob Double
forall a. a -> Prob a
forall (m :: * -> *) a. Monad m => a -> m a
return (Double -> Prob Double) -> Double -> Prob Double
forall a b. (a -> b) -> a -> b
$ Double -> Double -> Double -> Double
invIncompleteBeta Double
a Double
b Double
x

{-|
  [Poisson distribution](https://en.wikipedia.org/wiki/Poisson_distribution)
-}
poisson :: Double -- ^ lambda, rate
        -> Prob Integer
poisson :: Double -> Prob Integer
poisson Double
lambda = do
  Double
x <- Prob Double
uniform
  let cmf :: [Double]
cmf = (Integer -> Double) -> [Integer] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (\Integer
x -> Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double -> Double -> Double
incompleteGamma (Integer -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer
x Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1)) Double
lambda) [Integer
0,Integer
1..]
  let (Just Int
n) = (Double -> Bool) -> [Double] -> Maybe Int
forall a. (a -> Bool) -> [a] -> Maybe Int
findIndex (Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
x) [Double]
cmf
  Integer -> Prob Integer
forall a. a -> Prob a
forall (m :: * -> *) a. Monad m => a -> m a
return (Integer -> Prob Integer) -> Integer -> Prob Integer
forall a b. (a -> b) -> a -> b
$ Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n

poissonPdf :: Double -> Integer -> Double
poissonPdf :: Double -> Integer -> Double
poissonPdf Double
rate Integer
n = let result :: Double
result = Double -> Double
forall a. Floating a => a -> a
exp(-Double
rate) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
rate Double -> Integer -> Double
forall a b. (Fractional a, Integral b) => a -> b -> a
^^ (Integer -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
n) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Int -> Double
factorial (Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
n)) in 
  if (Double -> Bool
forall a. RealFloat a => a -> Bool
isInfinite Double
result) Bool -> Bool -> Bool
|| (Double -> Bool
forall a. RealFloat a => a -> Bool
isNaN Double
result) then Double -> Double
forall a. Floating a => a -> a
exp (-Double
rate Double -> Double -> Double
forall a. Num a => a -> a -> a
+ (Integer -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
n) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
log Double
rate Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double -> Double
logGamma (Integer -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
+Integer
1))) else Double
result


{-|
  [Dirichlet distribution](https://en.wikipedia.org/wiki/Dirichlet_distribution)
-}
dirichlet :: [Double] -- ^ vector of alphas; length is dimension
          -> Prob[Double]
dirichlet :: [Double] -> Prob [Double]
dirichlet [Double]
as = do
  [Double]
xs <- (Double -> Prob Double) -> [Double] -> Prob [Double]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (\Double
a -> Double -> Double -> Prob Double
gamma Double
a Double
1) [Double]
as
  let s :: Double
s = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
Prelude.sum [Double]
xs
  let ys :: [Double]
ys = (Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
s) [Double]
xs
  [Double] -> Prob [Double]
forall a. a -> Prob a
forall (m :: * -> *) a. Monad m => a -> m a
return [Double]
ys

-- | [Continuous uniform distribution on a bounded interval](https://en.wikipedia.org/wiki/Continuous_uniform_distribution)
uniformbounded :: Double -- ^ lower
               -> Double -- ^ upper
               -> Prob Double
uniformbounded :: Double -> Double -> Prob Double
uniformbounded Double
lower Double
upper = do
  Double
x <- Prob Double
uniform
  Double -> Prob Double
forall a. a -> Prob a
forall (m :: * -> *) a. Monad m => a -> m a
return (Double -> Prob Double) -> Double -> Prob Double
forall a b. (a -> b) -> a -> b
$ (Double
upper Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
lower) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
lower

-- | [Bernoulli distribution](https://en.wikipedia.org/wiki/Bernoulli_distribution)
bernoulli :: Double -- ^ bias
          -> Prob Bool
bernoulli :: Double -> Prob Bool
bernoulli Double
r = do
  Double
x <- Prob Double
uniform
  Bool -> Prob Bool
forall a. a -> Prob a
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool -> Prob Bool) -> Bool -> Prob Bool
forall a b. (a -> b) -> a -> b
$ Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
r

{-|
  [Discrete uniform distribution](https://en.wikipedia.org/wiki/Discrete_uniform_distribution) on [0, ..., n-1]
-}
uniformdiscrete :: Int -- ^ n
                -> Prob Int
uniformdiscrete :: Int -> Prob Int
uniformdiscrete Int
n =
  do
    let upper :: Double
upper = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n
    Double
r <- Double -> Double -> Prob Double
uniformbounded Double
0 Double
upper
    Int -> Prob Int
forall a. a -> Prob a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> Prob Int) -> Int -> Prob Int
forall a b. (a -> b) -> a -> b
$ Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor Double
r

{-| [Categorical distribution](https://www.google.com/search?client=safari&rls=en&q=categorical+distribution&ie=UTF-8&oe=UTF-8): Takes a list of k numbers that sum to 1, 
    and returns a random number between 0 and (k-1), weighted accordingly -}
categorical :: [Double] -> Prob Int
categorical :: [Double] -> Prob Int
categorical [Double]
xs = do 
  Double
r <- Prob Double
uniform
  case (Double -> Bool) -> [Double] -> Maybe Int
forall a. (a -> Bool) -> [a] -> Maybe Int
findIndex (Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>Double
r) ([Double] -> Maybe Int) -> [Double] -> Maybe Int
forall a b. (a -> b) -> a -> b
$ [Double] -> [Double]
forall a. HasCallStack => [a] -> [a]
tail ([Double] -> [Double]) -> [Double] -> [Double]
forall a b. (a -> b) -> a -> b
$ (Double -> Double -> Double) -> Double -> [Double] -> [Double]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl Double -> Double -> Double
forall a. Num a => a -> a -> a
(+) Double
0 [Double]
xs of
    Just Int
i -> Int -> Prob Int
forall a. a -> Prob a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
i
    Maybe Int
Nothing -> [Char] -> Prob Int
forall a. HasCallStack => [Char] -> a
error [Char]
"categorical: probabilities do not sum to 1"

{-| Returns an infinite stream of samples from the given distribution. --}
iid :: Prob a -> Prob [a]
iid :: forall a. Prob a -> Prob [a]
iid Prob a
p = do a
r <- Prob a
p; [a]
rs <- Prob a -> Prob [a]
forall a. Prob a -> Prob [a]
iid Prob a
p; [a] -> Prob [a]
forall a. a -> Prob a
forall (m :: * -> *) a. Monad m => a -> m a
return ([a] -> Prob [a]) -> [a] -> Prob [a]
forall a b. (a -> b) -> a -> b
$ a
r a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
rs