{-# LANGUAGE BangPatterns #-}

-- |
-- Module      :  Prior
-- Description :  Convenience functions for computing priors
-- Copyright   :  (c) Dominik Schrempf, 2020
-- License     :  GPL-3.0-or-later
--
-- Maintainer  :  dominik.schrempf@gmail.com
-- Stability   :  unstable
-- Portability :  portable
--
-- Creation date: Thu Jul 23 13:26:14 2020.
module Mcmc.Prior
  ( -- * Improper priors
    largerThan,
    positive,
    lowerThan,
    negative,

    -- * Continuous priors
    exponential,
    gamma,
    normal,
    uniform,

    -- * Discrete priors
    poisson,

    -- * Auxiliary functions
    product',
  )
where

import Control.Monad
import Data.Maybe (fromMaybe)
import Numeric.Log
import qualified Statistics.Distribution as S
import qualified Statistics.Distribution.Exponential as S
import qualified Statistics.Distribution.Gamma as S
import qualified Statistics.Distribution.Normal as S
import qualified Statistics.Distribution.Poisson as S

-- | Improper uniform prior; strictly larger than a given value.
largerThan :: Double -> Double -> Log Double
largerThan :: Double -> Double -> Log Double
largerThan Double
a Double
x
  | Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
a = Log Double
0
  | Bool
otherwise = Log Double
1

-- | Improper uniform prior; strictly larger than zero.
positive :: Double -> Log Double
positive :: Double -> Log Double
positive = Double -> Double -> Log Double
largerThan Double
0

-- | Improper uniform prior; strictly lower than a given value.
lowerThan :: Double -> Double -> Log Double
lowerThan :: Double -> Double -> Log Double
lowerThan Double
b Double
x
  | Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
b = Log Double
0
  | Bool
otherwise = Log Double
1

-- | Improper uniform prior; strictly lower than zero.
negative :: Double -> Log Double
negative :: Double -> Log Double
negative = Double -> Double -> Log Double
lowerThan Double
0

-- | Exponential distributed prior.
exponential ::
  -- | Rate.
  Double ->
  Double ->
  Log Double
exponential :: Double -> Double -> Log Double
exponential Double
l Double
x = Double -> Log Double
forall a. a -> Log a
Exp (Double -> Log Double) -> Double -> Log Double
forall a b. (a -> b) -> a -> b
$ ExponentialDistribution -> Double -> Double
forall d. ContDistr d => d -> Double -> Double
S.logDensity ExponentialDistribution
d Double
x
  where
    d :: ExponentialDistribution
d = Double -> ExponentialDistribution
S.exponential Double
l

-- | Gamma distributed prior.
gamma ::
  -- | Shape.
  Double ->
  -- | Scale.
  Double ->
  Double ->
  Log Double
gamma :: Double -> Double -> Double -> Log Double
gamma Double
k Double
t Double
x = Double -> Log Double
forall a. a -> Log a
Exp (Double -> Log Double) -> Double -> Log Double
forall a b. (a -> b) -> a -> b
$ GammaDistribution -> Double -> Double
forall d. ContDistr d => d -> Double -> Double
S.logDensity GammaDistribution
d Double
x
  where
    d :: GammaDistribution
d = Double -> Double -> GammaDistribution
S.gammaDistr Double
k Double
t

-- | Normal distributed prior.
normal ::
  -- | Mean.
  Double ->
  -- | Standard deviation.
  Double ->
  Double ->
  Log Double
normal :: Double -> Double -> Double -> Log Double
normal Double
m Double
s Double
x = Double -> Log Double
forall a. a -> Log a
Exp (Double -> Log Double) -> Double -> Log Double
forall a b. (a -> b) -> a -> b
$ NormalDistribution -> Double -> Double
forall d. ContDistr d => d -> Double -> Double
S.logDensity NormalDistribution
d Double
x
  where
    d :: NormalDistribution
d = Double -> Double -> NormalDistribution
S.normalDistr Double
m Double
s

-- | Uniform prior on [a, b].
uniform ::
  -- | Lower bound a.
  Double ->
  -- | Upper bound b.
  Double ->
  Double ->
  Log Double
uniform :: Double -> Double -> Double -> Log Double
uniform Double
a Double
b Double
x
  | Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
a = Log Double
0
  | Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
b = Log Double
0
  | Bool
otherwise = Double -> Log Double
forall a. a -> Log a
Exp Double
0

-- | Poisson distributed prior.
poisson ::
  -- | Rate.
  Double ->
  Int ->
  Log Double
poisson :: Double -> Int -> Log Double
poisson Double
l Int
x = Double -> Log Double
forall a. a -> Log a
Exp (Double -> Log Double) -> Double -> Log Double
forall a b. (a -> b) -> a -> b
$ PoissonDistribution -> Int -> Double
forall d. DiscreteDistr d => d -> Int -> Double
S.logProbability PoissonDistribution
d Int
x
  where
    d :: PoissonDistribution
d = Double -> PoissonDistribution
S.poisson Double
l

-- | Intelligent product that stops when encountering a zero.
--
-- Use with care because the elements are checked for positiveness, and this can
-- take some time if the list is long and does not contain any zeroes.
product' :: [Log Double] -> Log Double
product' :: [Log Double] -> Log Double
product' = Log Double -> Maybe (Log Double) -> Log Double
forall a. a -> Maybe a -> a
fromMaybe Log Double
0 (Maybe (Log Double) -> Log Double)
-> ([Log Double] -> Maybe (Log Double))
-> [Log Double]
-> Log Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Log Double] -> Maybe (Log Double)
prodM

-- The type could be generalized to any MonadPlus Integer
prodM :: [Log Double] -> Maybe (Log Double)
prodM :: [Log Double] -> Maybe (Log Double)
prodM = (Log Double -> Log Double -> Maybe (Log Double))
-> Log Double -> [Log Double] -> Maybe (Log Double)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (\ !Log Double
acc Log Double
x -> (Log Double
acc Log Double -> Log Double -> Log Double
forall a. Num a => a -> a -> a
* Log Double
x) Log Double -> Maybe () -> Maybe (Log Double)
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Log Double
acc Log Double -> Log Double -> Bool
forall a. Eq a => a -> a -> Bool
/= Log Double
0)) Log Double
1