{-# LANGUAGE OverloadedStrings, PatternGuards,
             DeriveDataTypeable, DeriveGeneric #-}
-- |
-- Module    : Statistics.Distribution.NegativeBinomial
-- Copyright : (c) 2022 Lorenz Minder
-- License   : BSD3
--
-- Maintainer  : lminder@gmx.net
-- Stability   : experimental
-- Portability : portable
--
-- The negative binomial distribution.  This is the discrete probability
-- distribution of the number of failures in a sequence of independent
-- yes\/no experiments before a specified number of successes /r/.  Each
-- Bernoulli trial has success probability /p/ in the range (0, 1].  The
-- parameter /r/ must be positive, but does not have to be integer.

module Statistics.Distribution.NegativeBinomial (
      NegativeBinomialDistribution
    -- * Constructors
    , negativeBinomial
    , negativeBinomialE
    -- * Accessors
    , nbdSuccesses
    , nbdProbability
) where

import Control.Applicative
import Data.Aeson                       (FromJSON(..), ToJSON, Value(..), (.:))
import Data.Binary                      (Binary(..))
import Data.Data                        (Data, Typeable)
import Data.Foldable                    (foldl')
import GHC.Generics                     (Generic)
import Numeric.SpecFunctions            (incompleteBeta, log1p)
import Numeric.SpecFunctions.Extra      (logChooseFast)
import Numeric.MathFunctions.Constants  (m_epsilon, m_tiny)

import qualified Statistics.Distribution as D
import Statistics.Internal

-- Math helper functions

-- | Generalized binomial coefficients.
--
--   These computes binomial coefficients with the small generalization
--   that the /n/ need not be integer, but can be real.
gChoose :: Double -> Int -> Double
gChoose :: Double -> Int -> Double
gChoose Double
n Int
k
    | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0             = Double
0
    | Double
k' Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
50          = Double -> Double
forall a. Floating a => a -> a
exp (Double -> Double) -> Double -> Double
forall a b. (a -> b) -> a -> b
$ Double -> Double -> Double
logChooseFast Double
n Double
k' 
    | Bool
otherwise         = (Double -> Double -> Double) -> Double -> [Double] -> Double
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) Double
1 [Double]
factors
    where   factors :: [Double]
factors = [ (Double
n Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
k' Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
j) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
j | Double
j <- [Double
1..Double
k'] ]
            k' :: Double
k' = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k


-- Implementation of Negative Binomial

-- | The negative binomial distribution.
data NegativeBinomialDistribution = NBD {
      NegativeBinomialDistribution -> Double
nbdSuccesses   :: {-# UNPACK #-} !Double
    -- ^ Number of successes until stop
    , NegativeBinomialDistribution -> Double
nbdProbability :: {-# UNPACK #-} !Double
    -- ^ Success probability.
    } deriving (NegativeBinomialDistribution
-> NegativeBinomialDistribution -> Bool
(NegativeBinomialDistribution
 -> NegativeBinomialDistribution -> Bool)
-> (NegativeBinomialDistribution
    -> NegativeBinomialDistribution -> Bool)
-> Eq NegativeBinomialDistribution
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: NegativeBinomialDistribution
-> NegativeBinomialDistribution -> Bool
== :: NegativeBinomialDistribution
-> NegativeBinomialDistribution -> Bool
$c/= :: NegativeBinomialDistribution
-> NegativeBinomialDistribution -> Bool
/= :: NegativeBinomialDistribution
-> NegativeBinomialDistribution -> Bool
Eq, Typeable, Typeable NegativeBinomialDistribution
Typeable NegativeBinomialDistribution =>
(forall (c :: * -> *).
 (forall d b. Data d => c (d -> b) -> d -> c b)
 -> (forall g. g -> c g)
 -> NegativeBinomialDistribution
 -> c NegativeBinomialDistribution)
-> (forall (c :: * -> *).
    (forall b r. Data b => c (b -> r) -> c r)
    -> (forall r. r -> c r)
    -> Constr
    -> c NegativeBinomialDistribution)
-> (NegativeBinomialDistribution -> Constr)
-> (NegativeBinomialDistribution -> DataType)
-> (forall (t :: * -> *) (c :: * -> *).
    Typeable t =>
    (forall d. Data d => c (t d))
    -> Maybe (c NegativeBinomialDistribution))
-> (forall (t :: * -> * -> *) (c :: * -> *).
    Typeable t =>
    (forall d e. (Data d, Data e) => c (t d e))
    -> Maybe (c NegativeBinomialDistribution))
-> ((forall b. Data b => b -> b)
    -> NegativeBinomialDistribution -> NegativeBinomialDistribution)
-> (forall r r'.
    (r -> r' -> r)
    -> r
    -> (forall d. Data d => d -> r')
    -> NegativeBinomialDistribution
    -> r)
-> (forall r r'.
    (r' -> r -> r)
    -> r
    -> (forall d. Data d => d -> r')
    -> NegativeBinomialDistribution
    -> r)
-> (forall u.
    (forall d. Data d => d -> u)
    -> NegativeBinomialDistribution -> [u])
-> (forall u.
    Int
    -> (forall d. Data d => d -> u)
    -> NegativeBinomialDistribution
    -> u)
-> (forall (m :: * -> *).
    Monad m =>
    (forall d. Data d => d -> m d)
    -> NegativeBinomialDistribution -> m NegativeBinomialDistribution)
-> (forall (m :: * -> *).
    MonadPlus m =>
    (forall d. Data d => d -> m d)
    -> NegativeBinomialDistribution -> m NegativeBinomialDistribution)
-> (forall (m :: * -> *).
    MonadPlus m =>
    (forall d. Data d => d -> m d)
    -> NegativeBinomialDistribution -> m NegativeBinomialDistribution)
-> Data NegativeBinomialDistribution
NegativeBinomialDistribution -> Constr
NegativeBinomialDistribution -> DataType
(forall b. Data b => b -> b)
-> NegativeBinomialDistribution -> NegativeBinomialDistribution
forall a.
Typeable a =>
(forall (c :: * -> *).
 (forall d b. Data d => c (d -> b) -> d -> c b)
 -> (forall g. g -> c g) -> a -> c a)
-> (forall (c :: * -> *).
    (forall b r. Data b => c (b -> r) -> c r)
    -> (forall r. r -> c r) -> Constr -> c a)
-> (a -> Constr)
-> (a -> DataType)
-> (forall (t :: * -> *) (c :: * -> *).
    Typeable t =>
    (forall d. Data d => c (t d)) -> Maybe (c a))
-> (forall (t :: * -> * -> *) (c :: * -> *).
    Typeable t =>
    (forall d e. (Data d, Data e) => c (t d e)) -> Maybe (c a))
-> ((forall b. Data b => b -> b) -> a -> a)
-> (forall r r'.
    (r -> r' -> r) -> r -> (forall d. Data d => d -> r') -> a -> r)
-> (forall r r'.
    (r' -> r -> r) -> r -> (forall d. Data d => d -> r') -> a -> r)
-> (forall u. (forall d. Data d => d -> u) -> a -> [u])
-> (forall u. Int -> (forall d. Data d => d -> u) -> a -> u)
-> (forall (m :: * -> *).
    Monad m =>
    (forall d. Data d => d -> m d) -> a -> m a)
-> (forall (m :: * -> *).
    MonadPlus m =>
    (forall d. Data d => d -> m d) -> a -> m a)
-> (forall (m :: * -> *).
    MonadPlus m =>
    (forall d. Data d => d -> m d) -> a -> m a)
-> Data a
forall u.
Int
-> (forall d. Data d => d -> u)
-> NegativeBinomialDistribution
-> u
forall u.
(forall d. Data d => d -> u) -> NegativeBinomialDistribution -> [u]
forall r r'.
(r -> r' -> r)
-> r
-> (forall d. Data d => d -> r')
-> NegativeBinomialDistribution
-> r
forall r r'.
(r' -> r -> r)
-> r
-> (forall d. Data d => d -> r')
-> NegativeBinomialDistribution
-> r
forall (m :: * -> *).
Monad m =>
(forall d. Data d => d -> m d)
-> NegativeBinomialDistribution -> m NegativeBinomialDistribution
forall (m :: * -> *).
MonadPlus m =>
(forall d. Data d => d -> m d)
-> NegativeBinomialDistribution -> m NegativeBinomialDistribution
forall (c :: * -> *).
(forall b r. Data b => c (b -> r) -> c r)
-> (forall r. r -> c r) -> Constr -> c NegativeBinomialDistribution
forall (c :: * -> *).
(forall d b. Data d => c (d -> b) -> d -> c b)
-> (forall g. g -> c g)
-> NegativeBinomialDistribution
-> c NegativeBinomialDistribution
forall (t :: * -> *) (c :: * -> *).
Typeable t =>
(forall d. Data d => c (t d))
-> Maybe (c NegativeBinomialDistribution)
forall (t :: * -> * -> *) (c :: * -> *).
Typeable t =>
(forall d e. (Data d, Data e) => c (t d e))
-> Maybe (c NegativeBinomialDistribution)
$cgfoldl :: forall (c :: * -> *).
(forall d b. Data d => c (d -> b) -> d -> c b)
-> (forall g. g -> c g)
-> NegativeBinomialDistribution
-> c NegativeBinomialDistribution
gfoldl :: forall (c :: * -> *).
(forall d b. Data d => c (d -> b) -> d -> c b)
-> (forall g. g -> c g)
-> NegativeBinomialDistribution
-> c NegativeBinomialDistribution
$cgunfold :: forall (c :: * -> *).
(forall b r. Data b => c (b -> r) -> c r)
-> (forall r. r -> c r) -> Constr -> c NegativeBinomialDistribution
gunfold :: forall (c :: * -> *).
(forall b r. Data b => c (b -> r) -> c r)
-> (forall r. r -> c r) -> Constr -> c NegativeBinomialDistribution
$ctoConstr :: NegativeBinomialDistribution -> Constr
toConstr :: NegativeBinomialDistribution -> Constr
$cdataTypeOf :: NegativeBinomialDistribution -> DataType
dataTypeOf :: NegativeBinomialDistribution -> DataType
$cdataCast1 :: forall (t :: * -> *) (c :: * -> *).
Typeable t =>
(forall d. Data d => c (t d))
-> Maybe (c NegativeBinomialDistribution)
dataCast1 :: forall (t :: * -> *) (c :: * -> *).
Typeable t =>
(forall d. Data d => c (t d))
-> Maybe (c NegativeBinomialDistribution)
$cdataCast2 :: forall (t :: * -> * -> *) (c :: * -> *).
Typeable t =>
(forall d e. (Data d, Data e) => c (t d e))
-> Maybe (c NegativeBinomialDistribution)
dataCast2 :: forall (t :: * -> * -> *) (c :: * -> *).
Typeable t =>
(forall d e. (Data d, Data e) => c (t d e))
-> Maybe (c NegativeBinomialDistribution)
$cgmapT :: (forall b. Data b => b -> b)
-> NegativeBinomialDistribution -> NegativeBinomialDistribution
gmapT :: (forall b. Data b => b -> b)
-> NegativeBinomialDistribution -> NegativeBinomialDistribution
$cgmapQl :: forall r r'.
(r -> r' -> r)
-> r
-> (forall d. Data d => d -> r')
-> NegativeBinomialDistribution
-> r
gmapQl :: forall r r'.
(r -> r' -> r)
-> r
-> (forall d. Data d => d -> r')
-> NegativeBinomialDistribution
-> r
$cgmapQr :: forall r r'.
(r' -> r -> r)
-> r
-> (forall d. Data d => d -> r')
-> NegativeBinomialDistribution
-> r
gmapQr :: forall r r'.
(r' -> r -> r)
-> r
-> (forall d. Data d => d -> r')
-> NegativeBinomialDistribution
-> r
$cgmapQ :: forall u.
(forall d. Data d => d -> u) -> NegativeBinomialDistribution -> [u]
gmapQ :: forall u.
(forall d. Data d => d -> u) -> NegativeBinomialDistribution -> [u]
$cgmapQi :: forall u.
Int
-> (forall d. Data d => d -> u)
-> NegativeBinomialDistribution
-> u
gmapQi :: forall u.
Int
-> (forall d. Data d => d -> u)
-> NegativeBinomialDistribution
-> u
$cgmapM :: forall (m :: * -> *).
Monad m =>
(forall d. Data d => d -> m d)
-> NegativeBinomialDistribution -> m NegativeBinomialDistribution
gmapM :: forall (m :: * -> *).
Monad m =>
(forall d. Data d => d -> m d)
-> NegativeBinomialDistribution -> m NegativeBinomialDistribution
$cgmapMp :: forall (m :: * -> *).
MonadPlus m =>
(forall d. Data d => d -> m d)
-> NegativeBinomialDistribution -> m NegativeBinomialDistribution
gmapMp :: forall (m :: * -> *).
MonadPlus m =>
(forall d. Data d => d -> m d)
-> NegativeBinomialDistribution -> m NegativeBinomialDistribution
$cgmapMo :: forall (m :: * -> *).
MonadPlus m =>
(forall d. Data d => d -> m d)
-> NegativeBinomialDistribution -> m NegativeBinomialDistribution
gmapMo :: forall (m :: * -> *).
MonadPlus m =>
(forall d. Data d => d -> m d)
-> NegativeBinomialDistribution -> m NegativeBinomialDistribution
Data, (forall x.
 NegativeBinomialDistribution -> Rep NegativeBinomialDistribution x)
-> (forall x.
    Rep NegativeBinomialDistribution x -> NegativeBinomialDistribution)
-> Generic NegativeBinomialDistribution
forall x.
Rep NegativeBinomialDistribution x -> NegativeBinomialDistribution
forall x.
NegativeBinomialDistribution -> Rep NegativeBinomialDistribution x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x.
NegativeBinomialDistribution -> Rep NegativeBinomialDistribution x
from :: forall x.
NegativeBinomialDistribution -> Rep NegativeBinomialDistribution x
$cto :: forall x.
Rep NegativeBinomialDistribution x -> NegativeBinomialDistribution
to :: forall x.
Rep NegativeBinomialDistribution x -> NegativeBinomialDistribution
Generic)

instance Show NegativeBinomialDistribution where
  showsPrec :: Int -> NegativeBinomialDistribution -> ShowS
showsPrec Int
i (NBD Double
r Double
p) = [Char] -> Double -> Double -> Int -> ShowS
forall a b. (Show a, Show b) => [Char] -> a -> b -> Int -> ShowS
defaultShow2 [Char]
"negativeBinomial" Double
r Double
p Int
i
instance Read NegativeBinomialDistribution where
  readPrec :: ReadPrec NegativeBinomialDistribution
readPrec = [Char]
-> (Double -> Double -> Maybe NegativeBinomialDistribution)
-> ReadPrec NegativeBinomialDistribution
forall a b r.
(Read a, Read b) =>
[Char] -> (a -> b -> Maybe r) -> ReadPrec r
defaultReadPrecM2 [Char]
"negativeBinomial" Double -> Double -> Maybe NegativeBinomialDistribution
negativeBinomialE

instance ToJSON NegativeBinomialDistribution
instance FromJSON NegativeBinomialDistribution where
  parseJSON :: Value -> Parser NegativeBinomialDistribution
parseJSON (Object Object
v) = do
    Double
r <- Object
v Object -> Key -> Parser Double
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"nbdSuccesses"
    Double
p <- Object
v Object -> Key -> Parser Double
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"nbdProbability"
    Parser NegativeBinomialDistribution
-> (NegativeBinomialDistribution
    -> Parser NegativeBinomialDistribution)
-> Maybe NegativeBinomialDistribution
-> Parser NegativeBinomialDistribution
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ([Char] -> Parser NegativeBinomialDistribution
forall a. [Char] -> Parser a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail ([Char] -> Parser NegativeBinomialDistribution)
-> [Char] -> Parser NegativeBinomialDistribution
forall a b. (a -> b) -> a -> b
$ Double -> Double -> [Char]
errMsg Double
r Double
p) NegativeBinomialDistribution -> Parser NegativeBinomialDistribution
forall a. a -> Parser a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe NegativeBinomialDistribution
 -> Parser NegativeBinomialDistribution)
-> Maybe NegativeBinomialDistribution
-> Parser NegativeBinomialDistribution
forall a b. (a -> b) -> a -> b
$ Double -> Double -> Maybe NegativeBinomialDistribution
negativeBinomialE Double
r Double
p
  parseJSON Value
_ = Parser NegativeBinomialDistribution
forall a. Parser a
forall (f :: * -> *) a. Alternative f => f a
empty

instance Binary NegativeBinomialDistribution where
  put :: NegativeBinomialDistribution -> Put
put (NBD Double
r Double
p) = Double -> Put
forall t. Binary t => t -> Put
put Double
r Put -> Put -> Put
forall a b. PutM a -> PutM b -> PutM b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Double -> Put
forall t. Binary t => t -> Put
put Double
p
  get :: Get NegativeBinomialDistribution
get = do
    Double
r <- Get Double
forall t. Binary t => Get t
get
    Double
p <- Get Double
forall t. Binary t => Get t
get
    Get NegativeBinomialDistribution
-> (NegativeBinomialDistribution
    -> Get NegativeBinomialDistribution)
-> Maybe NegativeBinomialDistribution
-> Get NegativeBinomialDistribution
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ([Char] -> Get NegativeBinomialDistribution
forall a. [Char] -> Get a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail ([Char] -> Get NegativeBinomialDistribution)
-> [Char] -> Get NegativeBinomialDistribution
forall a b. (a -> b) -> a -> b
$ Double -> Double -> [Char]
errMsg Double
r Double
p) NegativeBinomialDistribution -> Get NegativeBinomialDistribution
forall a. a -> Get a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe NegativeBinomialDistribution
 -> Get NegativeBinomialDistribution)
-> Maybe NegativeBinomialDistribution
-> Get NegativeBinomialDistribution
forall a b. (a -> b) -> a -> b
$ Double -> Double -> Maybe NegativeBinomialDistribution
negativeBinomialE Double
r Double
p

instance D.Distribution NegativeBinomialDistribution where
    cumulative :: NegativeBinomialDistribution -> Double -> Double
cumulative = NegativeBinomialDistribution -> Double -> Double
cumulative
    complCumulative :: NegativeBinomialDistribution -> Double -> Double
complCumulative = NegativeBinomialDistribution -> Double -> Double
complCumulative

instance D.DiscreteDistr NegativeBinomialDistribution where
    probability :: NegativeBinomialDistribution -> Int -> Double
probability    = NegativeBinomialDistribution -> Int -> Double
probability
    logProbability :: NegativeBinomialDistribution -> Int -> Double
logProbability = NegativeBinomialDistribution -> Int -> Double
logProbability

instance D.Mean NegativeBinomialDistribution where
    mean :: NegativeBinomialDistribution -> Double
mean = NegativeBinomialDistribution -> Double
mean

instance D.Variance NegativeBinomialDistribution where
    variance :: NegativeBinomialDistribution -> Double
variance = NegativeBinomialDistribution -> Double
variance

instance D.MaybeMean NegativeBinomialDistribution where
    maybeMean :: NegativeBinomialDistribution -> Maybe Double
maybeMean = Double -> Maybe Double
forall a. a -> Maybe a
Just (Double -> Maybe Double)
-> (NegativeBinomialDistribution -> Double)
-> NegativeBinomialDistribution
-> Maybe Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NegativeBinomialDistribution -> Double
forall d. Mean d => d -> Double
D.mean

instance D.MaybeVariance NegativeBinomialDistribution where
    maybeStdDev :: NegativeBinomialDistribution -> Maybe Double
maybeStdDev   = Double -> Maybe Double
forall a. a -> Maybe a
Just (Double -> Maybe Double)
-> (NegativeBinomialDistribution -> Double)
-> NegativeBinomialDistribution
-> Maybe Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NegativeBinomialDistribution -> Double
forall d. Variance d => d -> Double
D.stdDev
    maybeVariance :: NegativeBinomialDistribution -> Maybe Double
maybeVariance = Double -> Maybe Double
forall a. a -> Maybe a
Just (Double -> Maybe Double)
-> (NegativeBinomialDistribution -> Double)
-> NegativeBinomialDistribution
-> Maybe Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NegativeBinomialDistribution -> Double
forall d. Variance d => d -> Double
D.variance

instance D.Entropy NegativeBinomialDistribution where
   entropy :: NegativeBinomialDistribution -> Double
entropy = NegativeBinomialDistribution -> Double
directEntropy

instance D.MaybeEntropy NegativeBinomialDistribution where
   maybeEntropy :: NegativeBinomialDistribution -> Maybe Double
maybeEntropy = Double -> Maybe Double
forall a. a -> Maybe a
Just (Double -> Maybe Double)
-> (NegativeBinomialDistribution -> Double)
-> NegativeBinomialDistribution
-> Maybe Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NegativeBinomialDistribution -> Double
forall d. Entropy d => d -> Double
D.entropy

-- This could be slow for big n
probability :: NegativeBinomialDistribution -> Int -> Double
probability :: NegativeBinomialDistribution -> Int -> Double
probability d :: NegativeBinomialDistribution
d@(NBD Double
r Double
p) Int
k
  | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0          = Double
0
    -- Switch to log domain for large k + r to avoid overflows.
    --
    -- We also want to avoid underflow when computing (1-p)^k &
    -- p^r.
  | Double
k' Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
r Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
1000
  , Double
pK Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
m_tiny
  , Double
pR Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
m_tiny  = Double -> Int -> Double
gChoose (Double
k' Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
r Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1) Int
k Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
pK Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
pR
  | Bool
otherwise     = Double -> Double
forall a. Floating a => a -> a
exp (Double -> Double) -> Double -> Double
forall a b. (a -> b) -> a -> b
$ NegativeBinomialDistribution -> Int -> Double
logProbability NegativeBinomialDistribution
d Int
k
  where
    pK :: Double
pK  = Double -> Double
forall a. Floating a => a -> a
exp (Double -> Double) -> Double -> Double
forall a b. (a -> b) -> a -> b
$ Double -> Double
forall a. Floating a => a -> a
log1p (-Double
p) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
k'
    pR :: Double
pR  = Double
pDouble -> Double -> Double
forall a. Floating a => a -> a -> a
**Double
r
    k' :: Double
k'  = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k

logProbability :: NegativeBinomialDistribution -> Int -> Double
logProbability :: NegativeBinomialDistribution -> Int -> Double
logProbability (NBD Double
r Double
p) Int
k
  | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0                   = (-Double
1)Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
0
  | Bool
otherwise               = Double -> Double -> Double
logChooseFast (Double
k' Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
r Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1) Double
k'
                              Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double -> Double
forall a. Floating a => a -> a
log1p (-Double
p) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
k'
                              Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double -> Double
forall a. Floating a => a -> a
log Double
p Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
r
  where k' :: Double
k' = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k

cumulative :: NegativeBinomialDistribution -> Double -> Double
cumulative :: NegativeBinomialDistribution -> Double -> Double
cumulative (NBD Double
r Double
p) Double
x
  | Double -> Bool
forall a. RealFloat a => a -> Bool
isNaN Double
x      = [Char] -> Double
forall a. HasCallStack => [Char] -> a
error [Char]
"Statistics.Distribution.NegativeBinomial.cumulative: NaN input"
  | Double -> Bool
forall a. RealFloat a => a -> Bool
isInfinite Double
x = if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0 then Double
1 else Double
0
  | Integer
k Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0        = Double
0
  | Bool
otherwise    = Double -> Double -> Double -> Double
incompleteBeta Double
r (Integer -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer
kInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
+Integer
1)) Double
p
  where
    k :: Integer
k = Double -> Integer
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor Double
x :: Integer

complCumulative :: NegativeBinomialDistribution -> Double -> Double
complCumulative :: NegativeBinomialDistribution -> Double -> Double
complCumulative (NBD Double
r Double
p) Double
x
  | Double -> Bool
forall a. RealFloat a => a -> Bool
isNaN Double
x      = [Char] -> Double
forall a. HasCallStack => [Char] -> a
error [Char]
"Statistics.Distribution.NegativeBinomial.complCumulative: NaN input"
  | Double -> Bool
forall a. RealFloat a => a -> Bool
isInfinite Double
x = if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0 then Double
0 else Double
1
  | Integer
k Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0        = Double
1
  | Bool
otherwise    = Double -> Double -> Double -> Double
incompleteBeta (Integer -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer
kInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
+Integer
1)) Double
r (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
p)
  where
    k :: Integer
k = (Double -> Integer
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor Double
x)::Integer

mean :: NegativeBinomialDistribution -> Double
mean :: NegativeBinomialDistribution -> Double
mean (NBD Double
r Double
p) = Double
r Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
p)Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
p

variance :: NegativeBinomialDistribution -> Double
variance :: NegativeBinomialDistribution -> Double
variance (NBD Double
r Double
p) = Double
r Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
p)Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/(Double
p Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
p)

directEntropy :: NegativeBinomialDistribution -> Double
directEntropy :: NegativeBinomialDistribution -> Double
directEntropy NegativeBinomialDistribution
d =
  Double -> Double
forall a. Num a => a -> a
negate (Double -> Double) -> ([Double] -> Double) -> [Double] -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Double] -> Double) -> [Double] -> Double
forall a b. (a -> b) -> a -> b
$
  (Double -> Bool) -> [Double] -> [Double]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< -Double
m_epsilon) ([Double] -> [Double]) -> [Double] -> [Double]
forall a b. (a -> b) -> a -> b
$
  (Double -> Bool) -> [Double] -> [Double]
forall a. (a -> Bool) -> [a] -> [a]
dropWhile (Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= -Double
m_epsilon) ([Double] -> [Double]) -> [Double] -> [Double]
forall a b. (a -> b) -> a -> b
$
  [ let x :: Double
x = NegativeBinomialDistribution -> Int -> Double
probability NegativeBinomialDistribution
d Int
k in Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
log Double
x | Int
k <- [Int
0..]]

-- | Construct negative binomial distribution. Number of failures /r/
--   must be positive and probability must be in (0,1] range
negativeBinomial :: Double              -- ^ Number of successes.
                 -> Double              -- ^ Success probability.
                 -> NegativeBinomialDistribution
negativeBinomial :: Double -> Double -> NegativeBinomialDistribution
negativeBinomial Double
r Double
p = NegativeBinomialDistribution
-> (NegativeBinomialDistribution -> NegativeBinomialDistribution)
-> Maybe NegativeBinomialDistribution
-> NegativeBinomialDistribution
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ([Char] -> NegativeBinomialDistribution
forall a. HasCallStack => [Char] -> a
error ([Char] -> NegativeBinomialDistribution)
-> [Char] -> NegativeBinomialDistribution
forall a b. (a -> b) -> a -> b
$ Double -> Double -> [Char]
errMsg Double
r Double
p) NegativeBinomialDistribution -> NegativeBinomialDistribution
forall a. a -> a
id (Maybe NegativeBinomialDistribution
 -> NegativeBinomialDistribution)
-> Maybe NegativeBinomialDistribution
-> NegativeBinomialDistribution
forall a b. (a -> b) -> a -> b
$ Double -> Double -> Maybe NegativeBinomialDistribution
negativeBinomialE Double
r Double
p

-- | Construct negative binomial distribution. Number of failures /r/
--   must be positive and probability must be in (0,1] range
negativeBinomialE :: Double              -- ^ Number of successes.
                  -> Double              -- ^ Success probability.
                  -> Maybe NegativeBinomialDistribution
negativeBinomialE :: Double -> Double -> Maybe NegativeBinomialDistribution
negativeBinomialE Double
r Double
p
  | Double
r Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0 Bool -> Bool -> Bool
&& Double
0 Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
p Bool -> Bool -> Bool
&& Double
p Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
1            = NegativeBinomialDistribution -> Maybe NegativeBinomialDistribution
forall a. a -> Maybe a
Just (Double -> Double -> NegativeBinomialDistribution
NBD Double
r Double
p)
  | Bool
otherwise                           = Maybe NegativeBinomialDistribution
forall a. Maybe a
Nothing

errMsg :: Double -> Double -> String
errMsg :: Double -> Double -> [Char]
errMsg Double
r Double
p
  = [Char]
"Statistics.Distribution.NegativeBinomial.negativeBinomial: r=" [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Double -> [Char]
forall a. Show a => a -> [Char]
show Double
r
  [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
" p=" [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Double -> [Char]
forall a. Show a => a -> [Char]
show Double
p [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
", but need r>0 and p in (0,1]"