{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE BangPatterns, ScopedTypeVariables #-}
-- |
-- Module    : Statistics.Distribution
-- Copyright : (c) 2009 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Type classes for probability distributions

module Statistics.Distribution
    (
      -- * Type classes
      Distribution(..)
    , DiscreteDistr(..)
    , ContDistr(..)
      -- ** Distribution statistics
    , MaybeMean(..)
    , Mean(..)
    , MaybeVariance(..)
    , Variance(..)
    , MaybeEntropy(..)
    , Entropy(..)
    , FromSample(..)
      -- ** Random number generation
    , ContGen(..)
    , DiscreteGen(..)
    , genContinuous
      -- * Helper functions
    , findRoot
    , sumProbabilities
    ) where

import Prelude hiding (sum)
import Statistics.Function        (square)
import Statistics.Sample.Internal (sum)
import System.Random.Stateful     (StatefulGen, uniformDouble01M)
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Generic as G


-- | Type class common to all distributions. Only c.d.f. could be
-- defined for both discrete and continuous distributions.
class Distribution d where
    -- | Cumulative distribution function.  The probability that a
    -- random variable /X/ is less or equal than /x/,
    -- i.e. P(/X/≤/x/). Cumulative should be defined for
    -- infinities as well:
    --
    -- > cumulative d +∞ = 1
    -- > cumulative d -∞ = 0
    cumulative :: d -> Double -> Double
    cumulative d
d Double
x = Double
1 forall a. Num a => a -> a -> a
- forall d. Distribution d => d -> Double -> Double
complCumulative d
d Double
x
    -- | One's complement of cumulative distribution:
    --
    -- > complCumulative d x = 1 - cumulative d x
    --
    -- It's useful when one is interested in P(/X/>/x/) and
    -- expression on the right side begin to lose precision. This
    -- function have default implementation but implementors are
    -- encouraged to provide more precise implementation.
    complCumulative :: d -> Double -> Double
    complCumulative d
d Double
x = Double
1 forall a. Num a => a -> a -> a
- forall d. Distribution d => d -> Double -> Double
cumulative d
d Double
x
    {-# MINIMAL (cumulative | complCumulative) #-}


-- | Discrete probability distribution.
class Distribution  d => DiscreteDistr d where
    -- | Probability of n-th outcome.
    probability :: d -> Int -> Double
    probability d
d = forall a. Floating a => a -> a
exp forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall d. DiscreteDistr d => d -> Int -> Double
logProbability d
d
    -- | Logarithm of probability of n-th outcome
    logProbability :: d -> Int -> Double
    logProbability d
d = forall a. Floating a => a -> a
log forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall d. DiscreteDistr d => d -> Int -> Double
probability d
d
    {-# MINIMAL (probability | logProbability) #-}

-- | Continuous probability distribution.
--
--   Minimal complete definition is 'quantile' and either 'density' or
--   'logDensity'.
class Distribution d => ContDistr d where
    -- | Probability density function. Probability that random
    -- variable /X/ lies in the infinitesimal interval
    -- [/x/,/x+/δ/x/) equal to /density(x)/⋅δ/x/
    density :: d -> Double -> Double
    density d
d = forall a. Floating a => a -> a
exp forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall d. ContDistr d => d -> Double -> Double
logDensity d
d
    -- | Natural logarithm of density.
    logDensity :: d -> Double -> Double
    logDensity d
d = forall a. Floating a => a -> a
log forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall d. ContDistr d => d -> Double -> Double
density d
d
    -- | Inverse of the cumulative distribution function. The value
    -- /x/ for which P(/X/≤/x/) = /p/. If probability is outside
    -- of [0,1] range function should call 'error'
    quantile :: d -> Double -> Double
    quantile d
d Double
x = forall d. ContDistr d => d -> Double -> Double
complQuantile d
d (Double
1 forall a. Num a => a -> a -> a
- Double
x)
    -- | 1-complement of @quantile@:
    --
    -- > complQuantile x ≡ quantile (1 - x)
    complQuantile :: d -> Double -> Double
    complQuantile d
d Double
x = forall d. ContDistr d => d -> Double -> Double
quantile d
d (Double
1 forall a. Num a => a -> a -> a
- Double
x)
    {-# MINIMAL (density | logDensity), (quantile | complQuantile) #-}

-- | Type class for distributions with mean. 'maybeMean' should return
--   'Nothing' if it's undefined for current value of data
class Distribution d => MaybeMean d where
    maybeMean :: d -> Maybe Double

-- | Type class for distributions with mean. If a distribution has
--   finite mean for all valid values of parameters it should be
--   instance of this type class.
class MaybeMean d => Mean d where
    mean :: d -> Double



-- | Type class for distributions with variance. If variance is
--   undefined for some parameter values both 'maybeVariance' and
--   'maybeStdDev' should return Nothing.
--
--   Minimal complete definition is 'maybeVariance' or 'maybeStdDev'
class MaybeMean d => MaybeVariance d where
    maybeVariance :: d -> Maybe Double
    maybeVariance = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Double -> Double
square forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall d. MaybeVariance d => d -> Maybe Double
maybeStdDev
    maybeStdDev   :: d -> Maybe Double
    maybeStdDev   = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. Floating a => a -> a
sqrt forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall d. MaybeVariance d => d -> Maybe Double
maybeVariance
    {-# MINIMAL (maybeVariance | maybeStdDev) #-}

-- | Type class for distributions with variance. If distribution have
--   finite variance for all valid parameter values it should be
--   instance of this type class.
--
--   Minimal complete definition is 'variance' or 'stdDev'
class (Mean d, MaybeVariance d) => Variance d where
    variance :: d -> Double
    variance d
d = Double -> Double
square (forall d. Variance d => d -> Double
stdDev d
d)
    stdDev   :: d -> Double
    stdDev = forall a. Floating a => a -> a
sqrt forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall d. Variance d => d -> Double
variance
    {-# MINIMAL (variance | stdDev) #-}


-- | Type class for distributions with entropy, meaning Shannon entropy
--   in the case of a discrete distribution, or differential entropy in the
--   case of a continuous one.  'maybeEntropy' should return 'Nothing' if
--   entropy is undefined for the chosen parameter values.
class (Distribution d) => MaybeEntropy d where
  -- | Returns the entropy of a distribution, in nats, if such is defined.
  maybeEntropy :: d -> Maybe Double

-- | Type class for distributions with entropy, meaning Shannon
--   entropy in the case of a discrete distribution, or differential
--   entropy in the case of a continuous one.  If the distribution has
--   well-defined entropy for all valid parameter values then it
--   should be an instance of this type class.
class (MaybeEntropy d) => Entropy d where
  -- | Returns the entropy of a distribution, in nats.
  entropy :: d -> Double

-- | Generate discrete random variates which have given
--   distribution.
class Distribution d => ContGen d where
  genContVar :: (StatefulGen g m) => d -> g -> m Double

-- | Generate discrete random variates which have given
--   distribution. 'ContGen' is superclass because it's always possible
--   to generate real-valued variates from integer values
class (DiscreteDistr d, ContGen d) => DiscreteGen d where
  genDiscreteVar :: (StatefulGen g m) => d -> g -> m Int

-- | Estimate distribution from sample. First parameter in sample is
--   distribution type and second is element type.
class FromSample d a where
  -- | Estimate distribution from sample. Returns 'Nothing' if there is
  --   not enough data, or if no usable fit results from the method
  --   used, e.g., the estimated distribution parameters would be
  --   invalid or inaccurate.
  fromSample :: G.Vector v a => v a -> Maybe d


-- | Generate variates from continuous distribution using inverse
--   transform rule.
genContinuous :: (ContDistr d, StatefulGen g m) => d -> g -> m Double
genContinuous :: forall d g (m :: * -> *).
(ContDistr d, StatefulGen g m) =>
d -> g -> m Double
genContinuous d
d g
gen = do
  Double
x <- forall g (m :: * -> *). StatefulGen g m => g -> m Double
uniformDouble01M g
gen
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! forall d. ContDistr d => d -> Double -> Double
quantile d
d Double
x

data P = P {-# UNPACK #-} !Double {-# UNPACK #-} !Double

-- | Approximate the value of /X/ for which P(/x/>/X/)=/p/.
--
-- This method uses a combination of Newton-Raphson iteration and
-- bisection with the given guess as a starting point.  The upper and
-- lower bounds specify the interval in which the probability
-- distribution reaches the value /p/.
findRoot :: ContDistr d =>
            d                   -- ^ Distribution
         -> Double              -- ^ Probability /p/
         -> Double              -- ^ Initial guess
         -> Double              -- ^ Lower bound on interval
         -> Double              -- ^ Upper bound on interval
         -> Double
findRoot :: forall d.
ContDistr d =>
d -> Double -> Double -> Double -> Double -> Double
findRoot d
d Double
prob = Int -> Double -> Double -> Double -> Double -> Double
loop Int
0 Double
1
  where
    loop :: Int -> Double -> Double -> Double -> Double -> Double
loop !(Int
i::Int) !Double
dx !Double
x !Double
lo !Double
hi
      | forall a. Num a => a -> a
abs Double
dx forall a. Ord a => a -> a -> Bool
<= Double
accuracy Bool -> Bool -> Bool
|| Int
i forall a. Ord a => a -> a -> Bool
>= Int
maxIters = Double
x
      | Bool
otherwise                           = Int -> Double -> Double -> Double -> Double -> Double
loop (Int
iforall a. Num a => a -> a -> a
+Int
1) Double
dx'' Double
x'' Double
lo' Double
hi'
      where
        err :: Double
err                   = forall d. Distribution d => d -> Double -> Double
cumulative d
d Double
x forall a. Num a => a -> a -> a
- Double
prob
        P Double
lo' Double
hi' | Double
err forall a. Ord a => a -> a -> Bool
< Double
0   = Double -> Double -> P
P Double
x Double
hi
                  | Bool
otherwise = Double -> Double -> P
P Double
lo Double
x
        pdf :: Double
pdf                   = forall d. ContDistr d => d -> Double -> Double
density d
d Double
x
        P Double
dx' Double
x' | Double
pdf forall a. Eq a => a -> a -> Bool
/= Double
0   = Double -> Double -> P
P (Double
err forall a. Fractional a => a -> a -> a
/ Double
pdf) (Double
x forall a. Num a => a -> a -> a
- Double
dx)
                 | Bool
otherwise  = Double -> Double -> P
P Double
dx Double
x
        P Double
dx'' Double
x''
            | Double
x' forall a. Ord a => a -> a -> Bool
< Double
lo' Bool -> Bool -> Bool
|| Double
x' forall a. Ord a => a -> a -> Bool
> Double
hi' Bool -> Bool -> Bool
|| Double
pdf forall a. Eq a => a -> a -> Bool
== Double
0 = let y :: Double
y = (Double
lo' forall a. Num a => a -> a -> a
+ Double
hi') forall a. Fractional a => a -> a -> a
/ Double
2
                                                 in  Double -> Double -> P
P (Double
yforall a. Num a => a -> a -> a
-Double
x) Double
y
            | Bool
otherwise                        = Double -> Double -> P
P Double
dx' Double
x'
    accuracy :: Double
accuracy = Double
1e-15
    maxIters :: Int
maxIters = Int
150

-- | Sum probabilities in inclusive interval.
sumProbabilities :: DiscreteDistr d => d -> Int -> Int -> Double
sumProbabilities :: forall d. DiscreteDistr d => d -> Int -> Int -> Double
sumProbabilities d
d Int
low Int
hi =
  -- Return value is forced to be less than 1 to guard against roundoff errors.
  -- ATTENTION! this check should be removed for testing or it could mask bugs.
  forall a. Ord a => a -> a -> a
min Double
1 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (v :: * -> *). Vector v Double => v Double -> Double
sum forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
U.map (forall d. DiscreteDistr d => d -> Int -> Double
probability d
d) forall a b. (a -> b) -> a -> b
$ forall a. (Unbox a, Enum a) => a -> a -> Vector a
U.enumFromTo Int
low Int
hi