summaryrefslogtreecommitdiff
path: root/Statistics/Distribution/Geometric.hs
blob: 6eaa50243f7b2785efbf2620e5547c3e2343ddbf (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
-- |
-- Module    : Statistics.Distribution.Geometric
-- Copyright : (c) 2009 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- The Geometric distribution. There are two variants of
-- distribution. First is the probability distribution of the number
-- of Bernoulli trials needed to get one success, supported on the set
-- [1,2..] ('GeometricDistribution'). Sometimes it's referred to as
-- the /shifted/ geometric distribution to distinguish from another
-- one.
--
-- Second variant is probability distribution of the number of
-- failures before first success, defined over the set [0,1..]
-- ('GeometricDistribution0').
module Statistics.Distribution.Geometric
    (
      GeometricDistribution
    , GeometricDistribution0
    -- * Constructors
    , geometric
    , geometricE
    , geometric0
    , geometric0E
    -- ** Accessors
    , gdSuccess
    , gdSuccess0
    ) where

import Control.Applicative
import Control.Monad       (liftM)
import Data.Aeson          (FromJSON(..), ToJSON, Value(..), (.:))
import Data.Binary         (Binary(..))
import Data.Data           (Data, Typeable)
import GHC.Generics        (Generic)
import Numeric.MathFunctions.Constants (m_pos_inf, m_neg_inf)
import Numeric.SpecFunctions           (log1p,expm1)
import qualified System.Random.MWC.Distributions as MWC

import qualified Statistics.Distribution as D
import Statistics.Internal



----------------------------------------------------------------

-- | Distribution over [1..]
newtype GeometricDistribution = GD {
      gdSuccess :: Double
    } deriving (Eq, Typeable, Data, Generic)

instance Show GeometricDistribution where
  showsPrec i (GD x) = defaultShow1 "geometric" x i
instance Read GeometricDistribution where
  readPrec = defaultReadPrecM1 "geometric" geometricE

instance ToJSON GeometricDistribution
instance FromJSON GeometricDistribution where
  parseJSON (Object v) = do
    x <- v .: "gdSuccess"
    maybe (fail $ errMsg x) return  $ geometricE x
  parseJSON _ = empty

instance Binary GeometricDistribution where
  put (GD x) = put x
  get = do
    x <- get
    maybe (fail $ errMsg x) return  $ geometricE x


instance D.Distribution GeometricDistribution where
    cumulative      = cumulative
    complCumulative = complCumulative

instance D.DiscreteDistr GeometricDistribution where
    probability (GD s) n
      | n < 1     = 0
      | otherwise = s * (1-s) ** (fromIntegral n - 1)
    logProbability (GD s) n
       | n < 1     = m_neg_inf
       | otherwise = log s + log (1-s) * (fromIntegral n - 1)


instance D.Mean GeometricDistribution where
    mean (GD s) = 1 / s

instance D.Variance GeometricDistribution where
    variance (GD s) = (1 - s) / (s * s)

instance D.MaybeMean GeometricDistribution where
    maybeMean = Just . D.mean

instance D.MaybeVariance GeometricDistribution where
    maybeStdDev   = Just . D.stdDev
    maybeVariance = Just . D.variance

instance D.Entropy GeometricDistribution where
  entropy (GD s)
    | s == 0 = m_pos_inf
    | s == 1 = 0
    | otherwise = negate $ (s * log s + (1-s) * log (1-s)) / s

instance D.MaybeEntropy GeometricDistribution where
  maybeEntropy = Just . D.entropy

instance D.DiscreteGen GeometricDistribution where
  genDiscreteVar (GD s) g = MWC.geometric1 s g

instance D.ContGen GeometricDistribution where
  genContVar d g = fromIntegral `liftM` D.genDiscreteVar d g

cumulative :: GeometricDistribution -> Double -> Double
cumulative (GD s) x
  | x < 1        = 0
  | isInfinite x = 1
  | isNaN      x = error "Statistics.Distribution.Geometric.cumulative: NaN input"
  | otherwise    = negate $ expm1 $ fromIntegral (floor x :: Int) * log1p (-s)

complCumulative :: GeometricDistribution -> Double -> Double
complCumulative (GD s) x
  | x < 1        = 1
  | isInfinite x = 0
  | isNaN      x = error "Statistics.Distribution.Geometric.cumulative: NaN input"
  | otherwise    = exp $ fromIntegral (floor x :: Int) * log1p (-s)


-- | Create geometric distribution.
geometric :: Double                -- ^ Success rate
          -> GeometricDistribution
geometric x = maybe (error $ errMsg x) id $ geometricE x

-- | Create geometric distribution.
geometricE :: Double                -- ^ Success rate
           -> Maybe GeometricDistribution
geometricE x
  | x >= 0 && x <= 1 = Just (GD x)
  | otherwise        = Nothing

errMsg :: Double -> String
errMsg x = "Statistics.Distribution.Geometric.geometric: probability must be in [0,1] range. Got " ++ show x


----------------------------------------------------------------

-- | Distribution over [0..]
newtype GeometricDistribution0 = GD0 {
      gdSuccess0 :: Double
    } deriving (Eq, Typeable, Data, Generic)

instance Show GeometricDistribution0 where
  showsPrec i (GD0 x) = defaultShow1 "geometric0" x i
instance Read GeometricDistribution0 where
  readPrec = defaultReadPrecM1 "geometric0" geometric0E

instance ToJSON GeometricDistribution0
instance FromJSON GeometricDistribution0 where
  parseJSON (Object v) = do
    x <- v .: "gdSuccess0"
    maybe (fail $ errMsg x) return  $ geometric0E x
  parseJSON _ = empty

instance Binary GeometricDistribution0 where
  put (GD0 x) = put x
  get = do
    x <- get
    maybe (fail $ errMsg x) return  $ geometric0E x


instance D.Distribution GeometricDistribution0 where
    cumulative      (GD0 s) x = cumulative      (GD s) (x + 1)
    complCumulative (GD0 s) x = complCumulative (GD s) (x + 1)

instance D.DiscreteDistr GeometricDistribution0 where
    probability    (GD0 s) n = D.probability    (GD s) (n + 1)
    logProbability (GD0 s) n = D.logProbability (GD s) (n + 1)

instance D.Mean GeometricDistribution0 where
    mean (GD0 s) = 1 / s - 1

instance D.Variance GeometricDistribution0 where
    variance (GD0 s) = D.variance (GD s)

instance D.MaybeMean GeometricDistribution0 where
    maybeMean = Just . D.mean

instance D.MaybeVariance GeometricDistribution0 where
    maybeStdDev   = Just . D.stdDev
    maybeVariance = Just . D.variance

instance D.Entropy GeometricDistribution0 where
  entropy (GD0 s) = D.entropy (GD s)

instance D.MaybeEntropy GeometricDistribution0 where
  maybeEntropy = Just . D.entropy

instance D.DiscreteGen GeometricDistribution0 where
  genDiscreteVar (GD0 s) g = MWC.geometric0 s g

instance D.ContGen GeometricDistribution0 where
  genContVar d g = fromIntegral `liftM` D.genDiscreteVar d g


-- | Create geometric distribution.
geometric0 :: Double                -- ^ Success rate
           -> GeometricDistribution0
geometric0 x = maybe (error $ errMsg0 x) id $ geometric0E x

-- | Create geometric distribution.
geometric0E :: Double                -- ^ Success rate
            -> Maybe GeometricDistribution0
geometric0E x
  | x >= 0 && x <= 1 = Just (GD0 x)
  | otherwise        = Nothing

errMsg0 :: Double -> String
errMsg0 x = "Statistics.Distribution.Geometric.geometric0: probability must be in [0,1] range. Got " ++ show x