-- SPDX-License-Identifier: MPL-2.0

{- | Checked probability models. Gamma uses shape/scale, geometric and negative
binomial count failures before successes, Gumbel/GEV model maxima, Student-t
is standardized, and triangular parameters are lower/mode/upper. Discrete
densities are zero away from integer support; CDF=P(X\<=x), survival=P(X>x).
Quantiles use inf{x:F(x)>=p}. Log zero is -infinity by contract.
Finite reference summation/integration budgets are explicit in the code.
DiscreteUniform endpoints must be exactly representable integers (|bound|\<=2^53),
and Skellam log-mass order is capped at 1000000; CDF tails outside its finite
numerical summation window saturate to zero/one.
-}
module HMath.Probability (
    Distribution (..),
    validateDistribution,
    density,
    logDensity,
    cdf,
    survival,
    quantile,
    hazard,
    logLikelihood,
    moments,
    entropy,
    sample,
    fitNormal,
    fitExponential,
    fitBernoulli,
    fitMLEGrid,
    fitMAPGrid,
    betaBernoulliUpdate,
    gammaPoissonUpdate,
    normalNormalUpdate,
    conditionalProbability,
    bayes,
    expectation,
    goodnessOfFit,
    gaussianKDE,
    multinomialPMF,
    dirichletDensity,
    multivariateNormalLogDensity,
    sampleMultivariateNormal,
    wishartLogDensity,
    sampleWishart,
    Markov,
    markov,
    transition,
    stationary,
    HMM,
    hmm,
    forward,
    backward,
    posteriorStates,
    viterbi,
    baumWelch,
) where

import Control.Monad (foldM)
import qualified Data.Foldable as Foldable (foldl')
import Data.List (maximumBy, sort)
import Data.Ord (comparing)
import HMath.Error
import HMath.Matrix (Matrix, fromListChecked, multiplyChecked, shape, toList, transposeMatrix)
import HMath.Matrix.Decomposition (cholesky)
import HMath.Matrix.LU (determinant, inverse, solve)
import qualified HMath.Random as R
import qualified HMath.Special as S

{- | Parameterized probability law. Constructors carry: Bernoulli p; Binomial trials p;
NegativeBinomial successes p; Geometric p; Hypergeometric population successes draws;
Poisson rate; Categorical probabilities; DiscreteUniform lo hi; Uniform lo hi; Normal
mean standardDeviation; LogNormal logMean logStandardDeviation; Exponential rate; Gamma
shape scale; Beta alpha beta; ChiSquare degrees; StudentT degrees; FisherF
numeratorDegrees denominatorDegrees; Cauchy/Laplace/Logistic location scale; Weibull
shape scale; Pareto scale shape; Rayleigh/Maxwell scale; Gumbel location scale;
GeneralizedExtremeValue location scale shape; BetaBinomial trials alpha beta; Skellam
rate1 rate2; VonMises meanAngle concentration; Triangular lo mode hi; TruncatedNormal
mean standardDeviation lo hi; Mixture weightedLaws; Empirical samples; GaussianKDE
bandwidth samples. Call validateDistribution to check finite
parameters,support,normalization,and family budgets before numerical use.
-}
data Distribution
    = Bernoulli Double
    | Binomial Int Double
    | NegativeBinomial Double Double
    | Geometric Double
    | Hypergeometric Int Int Int
    | Poisson Double
    | Categorical [Double]
    | DiscreteUniform Int Int
    | Uniform Double Double
    | Normal Double Double
    | LogNormal Double Double
    | Exponential Double
    | Gamma Double Double
    | Beta Double Double
    | ChiSquare Double
    | StudentT Double
    | FisherF Double Double
    | Cauchy Double Double
    | Laplace Double Double
    | Logistic Double Double
    | Weibull Double Double
    | Pareto Double Double
    | Rayleigh Double
    | Maxwell Double
    | Gumbel Double Double
    | GeneralizedExtremeValue Double Double Double
    | BetaBinomial Int Double Double
    | Skellam Double Double
    | VonMises Double Double
    | Triangular Double Double Double
    | TruncatedNormal Double Double Double Double
    | Mixture [(Double, Distribution)]
    | Empirical [Double]
    | GaussianKDE Double [Double]
    deriving (Distribution -> Distribution -> Bool
(Distribution -> Distribution -> Bool)
-> (Distribution -> Distribution -> Bool) -> Eq Distribution
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Distribution -> Distribution -> Bool
== :: Distribution -> Distribution -> Bool
$c/= :: Distribution -> Distribution -> Bool
/= :: Distribution -> Distribution -> Bool
Eq, Int -> Distribution -> ShowS
[Distribution] -> ShowS
Distribution -> String
(Int -> Distribution -> ShowS)
-> (Distribution -> String)
-> ([Distribution] -> ShowS)
-> Show Distribution
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Distribution -> ShowS
showsPrec :: Int -> Distribution -> ShowS
$cshow :: Distribution -> String
show :: Distribution -> String
$cshowList :: [Distribution] -> ShowS
showList :: [Distribution] -> ShowS
Show)

finite :: Double -> Bool
finite :: Double -> Bool
finite Double
x = Bool -> Bool
not (Double -> Bool
forall a. RealFloat a => a -> Bool
isNaN Double
x Bool -> Bool -> Bool
|| Double -> Bool
forall a. RealFloat a => a -> Bool
isInfinite Double
x)
positive :: Double -> Bool
positive :: Double -> Bool
positive Double
x = Double -> Bool
finite Double
x Bool -> Bool -> Bool
&& Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0
require :: Bool -> String -> Either MathError ()
require :: Bool -> String -> Either MathError ()
require Bool
p String
msg = if Bool
p then () -> Either MathError ()
forall a b. b -> Either a b
Right () else MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
msg)
checked :: Double -> Either MathError Double
checked :: Double -> Either MathError Double
checked Double
x = if Double -> Bool
forall a. RealFloat a => a -> Bool
isNaN Double
x Bool -> Bool -> Bool
|| Double -> Bool
forall a. RealFloat a => a -> Bool
isInfinite Double
x then MathError -> Either MathError Double
forall a b. a -> Either a b
Left MathError
NonFiniteResult else Double -> Either MathError Double
forall a b. b -> Either a b
Right Double
x
prob :: Double -> Bool
prob :: Double -> Bool
prob Double
x = Double -> Bool
finite Double
x Bool -> Bool -> Bool
&& Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
0 Bool -> Bool -> Bool
&& Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
1

-- | Check finite parameters, support ordering, probability normalization, and declared per-family reference budgets.
validateDistribution :: Distribution -> Either MathError ()
validateDistribution :: Distribution -> Either MathError ()
validateDistribution Distribution
d = case Distribution
d of
    Bernoulli Double
p -> Bool -> Either MathError ()
ok (Double -> Bool
prob Double
p)
    Binomial Int
n Double
p -> Bool -> Either MathError ()
ok (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
100000 Bool -> Bool -> Bool
&& Double -> Bool
prob Double
p)
    NegativeBinomial Double
r Double
p -> Bool -> Either MathError ()
ok (Double -> Bool
positive Double
r Bool -> Bool -> Bool
&& Double -> Bool
positive Double
p Bool -> Bool -> Bool
&& Double
p Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
1)
    Geometric Double
p -> Bool -> Either MathError ()
ok (Double -> Bool
positive Double
p Bool -> Bool -> Bool
&& Double
p Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
1)
    Hypergeometric Int
n Int
k Int
draws -> Bool -> Either MathError ()
ok (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
100000 Bool -> Bool -> Bool
&& Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
n Bool -> Bool -> Bool
&& Int
draws Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
draws Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
n)
    Poisson Double
r -> Bool -> Either MathError ()
ok (Double -> Bool
finite Double
r Bool -> Bool -> Bool
&& Double
r Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
0 Bool -> Bool -> Bool
&& Double
r Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
100000)
    Categorical [Double]
ps -> Bool -> Either MathError ()
ok (Bool -> Bool
not ([Double] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Double]
ps) Bool -> Bool -> Bool
&& (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Double -> Bool
prob [Double]
ps Bool -> Bool -> Bool
&& Double -> Double
forall a. Num a => a -> a
abs ([Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
ps Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1) Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
1e-12)
    DiscreteUniform Int
lo Int
hi ->
        Bool -> Either MathError ()
ok
            ( Int
lo Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
hi
                Bool -> Bool -> Bool
&& Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
hi Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
lo Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
1000000
                Bool -> Bool -> Bool
&& Integer -> Integer
forall a. Num a => a -> a
abs (Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
lo) Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
2 Integer -> Int -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
53 :: Int)
                Bool -> Bool -> Bool
&& Integer -> Integer
forall a. Num a => a -> a
abs (Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
hi) Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
2 Integer -> Int -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
53 :: Int)
            )
    Uniform Double
lo Double
hi -> Bool -> Either MathError ()
ok (Double -> Bool
finite Double
lo Bool -> Bool -> Bool
&& Double -> Bool
finite Double
hi Bool -> Bool -> Bool
&& Double -> Bool
positive (Double
hi Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
lo))
    Normal Double
mu Double
sd -> Double -> Double -> Either MathError ()
loc Double
mu Double
sd
    LogNormal Double
mu Double
sd -> Double -> Double -> Either MathError ()
loc Double
mu Double
sd
    Exponential Double
rate -> Bool -> Either MathError ()
ok (Double -> Bool
positive Double
rate)
    Gamma Double
a Double
b -> Double -> Double -> Either MathError ()
pair Double
a Double
b
    Beta Double
a Double
b -> Double -> Double -> Either MathError ()
pair Double
a Double
b
    ChiSquare Double
a -> Bool -> Either MathError ()
ok (Double -> Bool
positive Double
a)
    StudentT Double
a -> Bool -> Either MathError ()
ok (Double -> Bool
positive Double
a)
    FisherF Double
a Double
b -> Double -> Double -> Either MathError ()
pair Double
a Double
b
    Cauchy Double
mu Double
scale -> Double -> Double -> Either MathError ()
loc Double
mu Double
scale
    Laplace Double
mu Double
scale -> Double -> Double -> Either MathError ()
loc Double
mu Double
scale
    Logistic Double
mu Double
scale -> Double -> Double -> Either MathError ()
loc Double
mu Double
scale
    Weibull Double
a Double
b -> Double -> Double -> Either MathError ()
pair Double
a Double
b
    Pareto Double
a Double
b -> Double -> Double -> Either MathError ()
pair Double
a Double
b
    Rayleigh Double
s -> Bool -> Either MathError ()
ok (Double -> Bool
positive Double
s)
    Maxwell Double
s -> Bool -> Either MathError ()
ok (Double -> Bool
positive Double
s)
    Gumbel Double
a Double
b -> Double -> Double -> Either MathError ()
loc Double
a Double
b
    GeneralizedExtremeValue Double
mu Double
sc Double
xi -> Bool -> Either MathError ()
ok (Double -> Bool
finite (Double
mu Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
xi) Bool -> Bool -> Bool
&& Double -> Bool
positive Double
sc)
    BetaBinomial Int
n Double
a Double
b -> Bool -> Either MathError ()
ok (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
100000 Bool -> Bool -> Bool
&& Double -> Bool
positive Double
a Bool -> Bool -> Bool
&& Double -> Bool
positive Double
b)
    Skellam Double
a Double
b -> Bool -> Either MathError ()
ok (Double -> Bool
positive Double
a Bool -> Bool -> Bool
&& Double -> Bool
positive Double
b Bool -> Bool -> Bool
&& Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
b Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
500)
    VonMises Double
mu Double
k -> Bool -> Either MathError ()
ok (Double -> Bool
finite Double
mu Bool -> Bool -> Bool
&& Double -> Bool
finite Double
k Bool -> Bool -> Bool
&& Double
k Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
0 Bool -> Bool -> Bool
&& Double
k Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
500)
    Triangular Double
a Double
b Double
c -> Bool -> Either MathError ()
ok ((Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Double -> Bool
finite [Double
a, Double
b, Double
c] Bool -> Bool -> Bool
&& Double
a Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
b Bool -> Bool -> Bool
&& Double
b Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
c Bool -> Bool -> Bool
&& Double
a Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
c)
    TruncatedNormal Double
mu Double
sd Double
a Double
b -> Bool -> Either MathError ()
ok (Double -> Bool
finite (Double
mu Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
b) Bool -> Bool -> Bool
&& Double -> Bool
positive Double
sd Bool -> Bool -> Bool
&& Double
a Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
b)
    Mixture [(Double, Distribution)]
pairs -> do
        Bool -> Either MathError ()
ok
            ( Bool -> Bool
not ([(Double, Distribution)] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [(Double, Distribution)]
pairs)
                Bool -> Bool -> Bool
&& ((Double, Distribution) -> Bool)
-> [(Double, Distribution)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\(Double
p, Distribution
_) -> Double -> Bool
finite Double
p Bool -> Bool -> Bool
&& Double
p Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
0) [(Double, Distribution)]
pairs
                Bool -> Bool -> Bool
&& Double -> Double
forall a. Num a => a -> a
abs ([Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (((Double, Distribution) -> Double)
-> [(Double, Distribution)] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double, Distribution) -> Double
forall a b. (a, b) -> a
fst [(Double, Distribution)]
pairs) Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1) Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
1e-12
            )
        ((Double, Distribution) -> Either MathError ())
-> [(Double, Distribution)] -> Either MathError ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Distribution -> Either MathError ()
validateDistribution (Distribution -> Either MathError ())
-> ((Double, Distribution) -> Distribution)
-> (Double, Distribution)
-> Either MathError ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Double, Distribution) -> Distribution
forall a b. (a, b) -> b
snd) [(Double, Distribution)]
pairs
        Bool -> Either MathError ()
ok (((Double, Distribution) -> Bool)
-> [(Double, Distribution)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ((Bool -> Bool -> Bool
forall a. Eq a => a -> a -> Bool
== Distribution -> Bool
isDiscrete ((Double, Distribution) -> Distribution
forall a b. (a, b) -> b
snd ([(Double, Distribution)]
pairs [(Double, Distribution)] -> Int -> (Double, Distribution)
forall a. HasCallStack => [a] -> Int -> a
!! Int
0))) (Bool -> Bool)
-> ((Double, Distribution) -> Bool)
-> (Double, Distribution)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Distribution -> Bool
isDiscrete (Distribution -> Bool)
-> ((Double, Distribution) -> Distribution)
-> (Double, Distribution)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Double, Distribution) -> Distribution
forall a b. (a, b) -> b
snd) [(Double, Distribution)]
pairs)
    Empirical [Double]
xs -> Bool -> Either MathError ()
ok (Bool -> Bool
not ([Double] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Double]
xs) Bool -> Bool -> Bool
&& (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Double -> Bool
finite [Double]
xs)
    GaussianKDE Double
h [Double]
xs -> Bool -> Either MathError ()
ok (Double -> Bool
positive Double
h Bool -> Bool -> Bool
&& Bool -> Bool
not ([Double] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Double]
xs) Bool -> Bool -> Bool
&& (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Double -> Bool
finite [Double]
xs)
  where
    ok :: Bool -> Either MathError ()
ok Bool
p = Bool -> String -> Either MathError ()
require Bool
p (String
"invalid distribution parameters: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Distribution -> String
forall a. Show a => a -> String
show Distribution
d)
    loc :: Double -> Double -> Either MathError ()
loc Double
a Double
b = Bool -> Either MathError ()
ok (Double -> Bool
finite Double
a Bool -> Bool -> Bool
&& Double -> Bool
positive Double
b)
    pair :: Double -> Double -> Either MathError ()
pair Double
a Double
b = Bool -> Either MathError ()
ok (Double -> Bool
positive Double
a Bool -> Bool -> Bool
&& Double -> Bool
positive Double
b)

isDiscrete :: Distribution -> Bool
isDiscrete :: Distribution -> Bool
isDiscrete Distribution
d = case Distribution
d of
    Bernoulli Double
_ -> Bool
True
    Binomial Int
_ Double
_ -> Bool
True
    NegativeBinomial Double
_ Double
_ -> Bool
True
    Geometric Double
_ -> Bool
True
    Hypergeometric Int
_ Int
_ Int
_ -> Bool
True
    Poisson Double
_ -> Bool
True
    Categorical [Double]
_ -> Bool
True
    DiscreteUniform Int
_ Int
_ -> Bool
True
    BetaBinomial Int
_ Double
_ Double
_ -> Bool
True
    Skellam Double
_ Double
_ -> Bool
True
    Empirical [Double]
_ -> Bool
True
    Mixture [(Double, Distribution)]
ps -> ((Double, Distribution) -> Bool)
-> [(Double, Distribution)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Distribution -> Bool
isDiscrete (Distribution -> Bool)
-> ((Double, Distribution) -> Distribution)
-> (Double, Distribution)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Double, Distribution) -> Distribution
forall a b. (a, b) -> b
snd) [(Double, Distribution)]
ps
    Distribution
_ -> Bool
False
logChoose :: Int -> Int -> Either MathError Double
logChoose :: Int -> Int -> Either MathError Double
logChoose Int
n Int
k
    | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
n = Double -> Either MathError Double
forall a b. b -> Either a b
Right (-Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0)
    | Bool
otherwise = do
        a <- Double -> Either MathError Double
S.logGamma (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1)
        b <- S.logGamma (fromIntegral k + 1)
        c <- S.logGamma (fromIntegral (n - k) + 1)
        pure (a - b - c)
xlog :: Double -> Double -> Double
xlog :: Double -> Double -> Double
xlog Double
a Double
b = if Double
a Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 then Double
0 else Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
log Double
b

-- | Log density or mass; zero probability is negative infinity, while an integrable endpoint singularity is positive infinity. Observations must be finite.
logDensity :: Distribution -> Double -> Either MathError Double
logDensity :: Distribution -> Double -> Either MathError Double
logDensity Distribution
d Double
x = do
    Distribution -> Either MathError ()
validateDistribution Distribution
d
    Bool -> String -> Either MathError ()
require (Double -> Bool
finite Double
x) String
"density requires finite observation"
    value <-
        if Distribution -> Bool
isDiscrete Distribution
d Bool -> Bool -> Bool
&& Bool -> Bool
not (Distribution -> Bool
empirical Distribution
d) Bool -> Bool -> Bool
&& Double
x Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
/= Integer -> Double
forall a. Num a => Integer -> a
fromInteger (Double -> Integer
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
round Double
x) then Double -> Either MathError Double
forall a b. b -> Either a b
Right (-Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0) else Distribution -> Double -> Either MathError Double
raw Distribution
d Double
x
    if isNaN value then Left NonFiniteResult else Right value
  where
    empirical :: Distribution -> Bool
empirical (Empirical [Double]
_) = Bool
True
    empirical (Mixture [(Double, Distribution)]
ps) = ((Double, Distribution) -> Bool)
-> [(Double, Distribution)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Distribution -> Bool
empirical (Distribution -> Bool)
-> ((Double, Distribution) -> Distribution)
-> (Double, Distribution)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Double, Distribution) -> Distribution
forall a b. (a, b) -> b
snd) [(Double, Distribution)]
ps
    empirical Distribution
_ = Bool
False
    raw :: Distribution -> Double -> Either MathError Double
raw Distribution
dist Double
y = case Distribution
dist of
        Bernoulli Double
p -> Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (if Double
y Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 then Double -> Double
forall a. Floating a => a -> a
log (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
p) else if Double
y Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
1 then Double -> Double
forall a. Floating a => a -> a
log Double
p else -Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0)
        Binomial Int
n Double
p ->
            if Double
y Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0 Bool -> Bool -> Bool
|| Double
y Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n
                then Either MathError Double
forall {a}. Either a Double
zero
                else do c <- Int -> Int -> Either MathError Double
logChoose Int
n (Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
round Double
y); pure (c + xlog y p + xlog (fromIntegral n - y) (1 - p))
        NegativeBinomial Double
r Double
p ->
            if Double
y Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0
                then Either MathError Double
forall {a}. Either a Double
zero
                else do
                    a <- Double -> Either MathError Double
S.logGamma (Double
y Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
r)
                    b <- S.logGamma r
                    c <- S.logGamma (y + 1)
                    pure (a - b - c + r * log p + xlog y (1 - p))
        Geometric Double
p -> Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (if Double
y Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0 then -Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0 else Double -> Double
forall a. Floating a => a -> a
log Double
p Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double -> Double -> Double
xlog Double
y (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
p))
        Hypergeometric Int
n Int
k Int
draws ->
            if Double
y Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int
draws Int -> Int -> Int
forall a. Num a => a -> a -> a
- (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
k))) Bool -> Bool -> Bool
|| Double
y Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
draws Int
k)
                then Either MathError Double
forall {a}. Either a Double
zero
                else do
                    a <- Int -> Int -> Either MathError Double
logChoose Int
k (Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
round Double
y)
                    b <- logChoose (n - k) (draws - round y)
                    c <- logChoose n draws
                    pure (a + b - c)
        Poisson Double
rate ->
            if Double
y Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0
                then Either MathError Double
forall {a}. Either a Double
zero
                else
                    if Double
rate Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0
                        then Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (if Double
y Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 then Double
0 else -Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0)
                        else do g <- Double -> Either MathError Double
S.logGamma (Double
y Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1); pure (y * log rate - rate - g)
        Categorical [Double]
ps -> Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (if Double
y Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0 Bool -> Bool -> Bool
|| Double
y Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
ps) then -Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0 else Double -> Double
forall a. Floating a => a -> a
log ([Double]
ps [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
round Double
y))
        DiscreteUniform Int
lo Int
hi ->
            Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
                ( if Double
y Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
lo Bool -> Bool -> Bool
|| Double
y Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
hi
                    then -Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0
                    else -Double -> Double
forall a. Floating a => a -> a
log (Integer -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
hi Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
lo Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1))
                )
        Uniform Double
lo Double
hi -> Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (if Double
y Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
lo Bool -> Bool -> Bool
|| Double
y Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
hi then -Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0 else -Double -> Double
forall a. Floating a => a -> a
log (Double
hi Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
lo))
        Normal Double
mu Double
sd -> Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (-Double
0.5 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
log (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
forall a. Floating a => a
pi) Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double -> Double
forall a. Floating a => a -> a
log Double
sd Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
0.5 Double -> Double -> Double
forall a. Num a => a -> a -> a
* ((Double
y Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
mu) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
sd) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int))
        LogNormal Double
mu Double
sd -> if Double
y Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 then Either MathError Double
forall {a}. Either a Double
zero else (Double -> Double -> Double
forall a. Num a => a -> a -> a
subtract (Double -> Double
forall a. Floating a => a -> a
log Double
y)) (Double -> Double)
-> Either MathError Double -> Either MathError Double
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Distribution -> Double -> Either MathError Double
raw (Double -> Double -> Distribution
Normal Double
mu Double
sd) (Double -> Double
forall a. Floating a => a -> a
log Double
y)
        Exponential Double
rate -> Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (if Double
y Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0 then -Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0 else Double -> Double
forall a. Floating a => a -> a
log Double
rate Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
rate Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
y)
        Gamma Double
a Double
sc ->
            if Double
y Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0
                then Either MathError Double
forall {a}. Either a Double
zero
                else
                    if Double
y Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0
                        then Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (if Double
a Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
1 then -Double -> Double
forall a. Floating a => a -> a
log Double
sc else if Double
a Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
1 then -Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0 else Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0)
                        else do g <- Double -> Either MathError Double
S.logGamma Double
a; pure ((a - 1) * log y - y / sc - g - a * log sc)
        Beta Double
a Double
b ->
            if Double
y Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0 Bool -> Bool -> Bool
|| Double
y Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
1
                then Either MathError Double
forall {a}. Either a Double
zero
                else do lb <- Double -> Double -> Either MathError Double
S.logBeta Double
a Double
b; pure (xlog (a - 1) y + xlog (b - 1) (1 - y) - lb)
        ChiSquare Double
k -> Distribution -> Double -> Either MathError Double
raw (Double -> Double -> Distribution
Gamma (Double
k Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2) Double
2) Double
y
        StudentT Double
nu -> do
            a <- Double -> Either MathError Double
S.logGamma ((Double
nu Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2)
            b <- S.logGamma (nu / 2)
            pure (a - b - 0.5 * log (nu * pi) - (nu + 1) / 2 * log (1 + y * y / nu))
        FisherF Double
a Double
b ->
            if Double
y Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0
                then Either MathError Double
forall {a}. Either a Double
zero
                else do
                    lb <- Double -> Double -> Either MathError Double
S.logBeta (Double
a Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2) (Double
b Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2)
                    pure (a / 2 * log (a / b) + (a / 2 - 1) * log y - (a + b) / 2 * log (1 + a * y / b) - lb)
        Cauchy Double
mu Double
sc -> Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (-Double -> Double
forall a. Floating a => a -> a
log (Double
forall a. Floating a => a
pi Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sc) Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double -> Double
forall a. Floating a => a -> a
log (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ ((Double
y Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
mu) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
sc) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int)))
        Laplace Double
mu Double
sc -> Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (-Double -> Double
forall a. Floating a => a -> a
log (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sc) Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double -> Double
forall a. Num a => a -> a
abs (Double
y Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
mu) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
sc)
        Logistic Double
mu Double
sc -> let z :: Double
z = Double -> Double
forall a. Num a => a -> a
abs ((Double
y Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
mu) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
sc) in Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (-Double
z Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
log (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double -> Double
forall a. Floating a => a -> a
exp (-Double
z)) Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double -> Double
forall a. Floating a => a -> a
log Double
sc)
        Weibull Double
k Double
sc -> if Double
y Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0 then Either MathError Double
forall {a}. Either a Double
zero else Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double -> Double
forall a. Floating a => a -> a
log (Double
k Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
sc) Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double -> Double -> Double
xlog (Double
k Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1) (Double
y Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
sc) Double -> Double -> Double
forall a. Num a => a -> a -> a
- (Double
y Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
sc) Double -> Double -> Double
forall a. Floating a => a -> a -> a
** Double
k)
        Pareto Double
xm Double
a -> Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (if Double
y Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
xm then -Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0 else Double -> Double
forall a. Floating a => a -> a
log Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
log Double
xm Double -> Double -> Double
forall a. Num a => a -> a -> a
- (Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
log Double
y)
        Rayleigh Double
sc -> Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (if Double
y Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 then -Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0 else Double -> Double
forall a. Floating a => a -> a
log Double
y Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
log Double
sc Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
y Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
y Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sc Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sc))
        Maxwell Double
sc ->
            Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
                (if Double
y Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 then -Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0 else Double
0.5 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
log (Double
2 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
forall a. Floating a => a
pi) Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
log Double
y Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
3 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
log Double
sc Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
y Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
y Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sc Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sc))
        Gumbel Double
mu Double
sc -> let z :: Double
z = (Double
y Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
mu) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
sc in Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (if Double -> Bool
forall a. RealFloat a => a -> Bool
isInfinite Double
z then -Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0 else -Double -> Double
forall a. Floating a => a -> a
log Double
sc Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
z Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double -> Double
forall a. Floating a => a -> a
exp (-Double
z))
        GeneralizedExtremeValue Double
mu Double
sc Double
xi ->
            if Double -> Double
forall a. Num a => a -> a
abs Double
xi Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
1e-12
                then Distribution -> Double -> Either MathError Double
raw (Double -> Double -> Distribution
Gumbel Double
mu Double
sc) Double
y
                else
                    let t :: Double
t = Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
xi Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
y Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
mu) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
sc
                     in Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (if Double
t Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 then -Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0 else -Double -> Double
forall a. Floating a => a -> a
log Double
sc Double -> Double -> Double
forall a. Num a => a -> a -> a
- (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
xi) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
log Double
t Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
t Double -> Double -> Double
forall a. Floating a => a -> a -> a
** (-Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
xi))
        BetaBinomial Int
n Double
a Double
b ->
            if Double
y Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0 Bool -> Bool -> Bool
|| Double
y Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n
                then Either MathError Double
forall {a}. Either a Double
zero
                else do
                    c <- Int -> Int -> Either MathError Double
logChoose Int
n (Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
round Double
y)
                    u <- S.logBeta (y + a) (fromIntegral n - y + b)
                    v <- S.logBeta a b
                    pure (c + u - v)
        Skellam Double
a Double
b -> do
            let order :: Integer
order = Integer -> Integer
forall a. Num a => a -> a
abs (Double -> Integer
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
round Double
y) :: Integer
            Bool -> String -> Either MathError ()
require (Integer
order Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
1000000) String
"Skellam absolute integer order exceeds 1000000 reference budget"
            let n :: Int
n = Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
order :: Int; t :: Double
t = Double -> Double
forall a. Floating a => a -> a
sqrt (Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
b)
            lg <- Double -> Either MathError Double
S.logGamma (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1)
            let go t
k Double
term Double
accum
                    | t
k t -> t -> Bool
forall a. Ord a => a -> a -> Bool
> t
10000 = MathError -> Either MathError Double
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"Skellam Bessel series" Int
10000)
                    | Bool
otherwise =
                        let u :: Double
u = Double
term Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
t Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
t Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (t -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral t
k Double -> Double -> Double
forall a. Num a => a -> a -> a
* (t -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral t
k Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n)); v :: Double
v = Double
accum Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
u
                         in if Double
u Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
1e-14 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
v then Double -> Either MathError Double
forall a b. b -> Either a b
Right Double
v else t -> Double -> Double -> Either MathError Double
go (t
k t -> t -> t
forall a. Num a => a -> a -> a
+ t
1) Double
u Double
v
            v <- go (1 :: Int) 1 1
            pure (-(a + b) + y / 2 * log (a / b) + fromIntegral n * log t - lg + log v)
        VonMises Double
mu Double
k -> do v <- Int -> Double -> Either MathError SFResult
S.scaledBesselI Int
0 Double
k; pure (k * (cos (y - mu) - 1) - log (2 * pi * S.sfValue v))
        Triangular Double
a Double
b Double
c ->
            Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
                ( if Double
y Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
a Bool -> Bool -> Bool
|| Double
y Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
c
                    then -Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0
                    else
                        if Double
y Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
b
                            then
                                if Double
a Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
b then Double -> Double
forall a. Floating a => a -> a
log (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
c Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
y) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ ((Double
c Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
a) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
c Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
b))) else Double -> Double
forall a. Floating a => a -> a
log (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
y Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
a) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ ((Double
c Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
a) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
a)))
                            else
                                if Double
b Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
c then Double -> Double
forall a. Floating a => a -> a
log (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
y Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
a) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ ((Double
c Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
a) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
a))) else Double -> Double
forall a. Floating a => a -> a
log (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
c Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
y) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ ((Double
c Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
a) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
c Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
b)))
                )
        TruncatedNormal Double
mu Double
sd Double
a Double
b ->
            if Double
y Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
a Bool -> Bool -> Bool
|| Double
y Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
b
                then Either MathError Double
forall {a}. Either a Double
zero
                else do
                    u <- Double -> Either MathError Double
S.normalCDF ((Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
mu) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
sd)
                    v <- S.normalCDF ((b - mu) / sd)
                    require (v > u) "truncated normal has unresolved mass"
                    p <- raw (Normal mu sd) y
                    pure (p - log (v - u))
        Mixture [(Double, Distribution)]
ps -> do ls <- ((Double, Distribution) -> Either MathError Double)
-> [(Double, Distribution)] -> Either MathError [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
p, Distribution
m) -> (Double -> Double
forall a. Floating a => a -> a
log Double
p Double -> Double -> Double
forall a. Num a => a -> a -> a
+) (Double -> Double)
-> Either MathError Double -> Either MathError Double
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Distribution -> Double -> Either MathError Double
logDensity Distribution
m Double
y) [(Double, Distribution)]
ps; pure (logSumExp ls)
        Empirical [Double]
xs -> Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double -> Double
forall a. Floating a => a -> a
log (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ((Double -> Bool) -> [Double] -> [Double]
forall a. (a -> Bool) -> [a] -> [a]
filter (Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
y) [Double]
xs)) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs)))
        GaussianKDE Double
h [Double]
xs -> do ls <- (Double -> Either MathError Double)
-> [Double] -> Either MathError [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
mu -> Distribution -> Double -> Either MathError Double
raw (Double -> Double -> Distribution
Normal Double
mu Double
h) Double
y) [Double]
xs; pure (logSumExp ls - log (fromIntegral (length xs)))
    zero :: Either a Double
zero = Double -> Either a Double
forall a b. b -> Either a b
Right (-Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0)

log1pexp :: Double -> Double
log1pexp :: Double -> Double
log1pexp Double
x = if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0 then Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double -> Double
forall a. Floating a => a -> a
log (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double -> Double
forall a. Floating a => a -> a
exp (-Double
x)) else Double -> Double
forall a. Floating a => a -> a
log (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double -> Double
forall a. Floating a => a -> a
exp Double
x)
logSumExp :: [Double] -> Double
logSumExp :: [Double] -> Double
logSumExp [Double]
xs = let m :: Double
m = [Double] -> Double
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [Double]
xs in if Double -> Bool
forall a. RealFloat a => a -> Bool
isInfinite Double
m then Double
m else Double
m Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double -> Double
forall a. Floating a => a -> a
log ([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]
map (\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]
xs))

-- | Probability mass for discrete families or density for continuous families. At integrable density singularities the result can be positive infinity; outside support it is zero.
density :: Distribution -> Double -> Either MathError Double
density :: Distribution -> Double -> Either MathError Double
density Distribution
d Double
x = Double -> Double
forall a. Floating a => a -> a
exp (Double -> Double)
-> Either MathError Double -> Either MathError Double
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Distribution -> Double -> Either MathError Double
logDensity Distribution
d Double
x

-- | P(X\<=x). Gamma/beta and classical transformed families use analytic special-function identities. Remaining finite-support continuous cases use 2048-panel Simpson integration; Skellam drops a tail beyond 14 standard deviations plus 20 counts.
cdf :: Distribution -> Double -> Either MathError Double
cdf :: Distribution -> Double -> Either MathError Double
cdf Distribution
d Double
x = do
    Distribution -> Either MathError ()
validateDistribution Distribution
d
    Bool -> String -> Either MathError ()
require (Double -> Bool
finite Double
x) String
"CDF requires finite argument"
    p <- case Distribution
d of
        Poisson Double
rate -> if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0 then Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
0 else Double -> Double -> Either MathError Double
S.gammaQ (Integer -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Double -> Integer
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor Double
x :: Integer) Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1) Double
rate
        Binomial Int
n Double
p ->
            if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0
                then Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
0
                else
                    if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n
                        then Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
1
                        else
                            if Double
p Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0
                                then Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
1
                                else
                                    if Double
p Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
1
                                        then Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
0
                                        else
                                            Double -> Double -> Double -> Either MathError Double
S.betaI
                                                (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n Double -> Double -> Double
forall a. Num a => a -> a -> a
- Integer -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Double -> Integer
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor Double
x :: Integer))
                                                (Integer -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Double -> Integer
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor Double
x :: Integer) Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1)
                                                (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
p)
        NegativeBinomial Double
r Double
p ->
            if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0
                then Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
0
                else if Double
p Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
1 then Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
1 else Double -> Double -> Double -> Either MathError Double
S.betaI Double
r (Integer -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Double -> Integer
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor Double
x :: Integer) Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1) Double
p
        Geometric Double
p -> Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0 then Double
0 else Double
1 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. Floating a => a -> a -> a
** (Integer -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Double -> Integer
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor Double
x :: Integer) Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1))
        Normal Double
mu Double
sd -> Double -> Either MathError Double
S.normalCDF ((Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
mu) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
sd)
        LogNormal Double
mu Double
sd -> if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 then Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
0 else Double -> Either MathError Double
S.normalCDF ((Double -> Double
forall a. Floating a => a -> a
log Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
mu) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
sd)
        Uniform Double
a Double
b -> Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
a) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
a))
        Exponential Double
rate -> Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 then Double
0 else Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double -> Double
forall a. Floating a => a -> a
exp (-Double
rate Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x))
        Gamma Double
a Double
sc -> if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 then Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
0 else Double -> Double -> Either MathError Double
S.gammaP Double
a (Double
x Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
sc)
        Beta Double
a Double
b -> if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 then Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
0 else if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
1 then Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
1 else Double -> Double -> Double -> Either MathError Double
S.betaI Double
a Double
b Double
x
        ChiSquare Double
k -> Distribution -> Double -> Either MathError Double
cdf (Double -> Double -> Distribution
Gamma (Double
k Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2) Double
2) Double
x
        StudentT Double
nu -> do b <- Double -> Double -> Double -> Either MathError Double
S.betaI (Double
nu Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2) Double
0.5 (Double
nu Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
nu Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x)); pure (if x >= 0 then 1 - b / 2 else b / 2)
        FisherF Double
a Double
b -> if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 then Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
0 else Double -> Double -> Double -> Either MathError Double
S.betaI (Double
a Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2) (Double
b Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2) (Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
b))
        Cauchy Double
mu Double
sc -> Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double
0.5 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double -> Double
forall a. Floating a => a -> a
atan ((Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
mu) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
sc) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
forall a. Floating a => a
pi)
        Laplace Double
mu Double
sc -> Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
mu then Double
0.5 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
exp ((Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
mu) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
sc) else Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
0.5 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
exp (-(Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
mu) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
sc))
        Logistic Double
mu Double
sc -> Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double -> Double
forall a. Floating a => a -> a
exp (-Double -> Double
log1pexp (-(Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
mu) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
sc)))
        Weibull Double
k Double
sc -> Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 then Double
0 else Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double -> Double
forall a. Floating a => a -> a
exp (-(Double
x Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
sc) Double -> Double -> Double
forall a. Floating a => a -> a -> a
** Double
k))
        Pareto Double
xm Double
a -> Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
xm then Double
0 else Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- (Double
xm Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
x) Double -> Double -> Double
forall a. Floating a => a -> a -> a
** Double
a)
        Rayleigh Double
sc -> Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 then Double
0 else Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double -> Double
forall a. Floating a => a -> a
exp (-Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sc Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sc)))
        Maxwell Double
sc ->
            if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0
                then Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
0
                else do e <- Double -> Either MathError Double
S.erf (Double
x Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double -> Double
forall a. Floating a => a -> a
sqrt Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sc)); pure (e - sqrt (2 / pi) * x / sc * exp (-x * x / (2 * sc * sc)))
        Gumbel Double
mu Double
sc -> Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double -> Double
forall a. Floating a => a -> a
exp (-Double -> Double
forall a. Floating a => a -> a
exp (-(Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
mu) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
sc)))
        GeneralizedExtremeValue Double
mu Double
sc Double
xi ->
            if Double -> Double
forall a. Num a => a -> a
abs Double
xi Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
1e-12
                then Distribution -> Double -> Either MathError Double
cdf (Double -> Double -> Distribution
Gumbel Double
mu Double
sc) Double
x
                else
                    let t :: Double
t = Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
xi Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
mu) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
sc
                     in Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (if Double
t Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 then if Double
xi Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0 then Double
0 else Double
1 else Double -> Double
forall a. Floating a => a -> a
exp (-Double
t Double -> Double -> Double
forall a. Floating a => a -> a -> a
** (-Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
xi)))
        Triangular Double
a Double
b Double
c ->
            Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
                ( if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
a
                    then Double
0
                    else
                        if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
c
                            then Double
1
                            else
                                if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
b
                                    then (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
a) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ ((Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
a) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
c Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
a))
                                    else Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- (Double
c Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
x) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ ((Double
c Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
a) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
c Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
b))
                )
        TruncatedNormal Double
mu Double
sd Double
a Double
b ->
            if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
a
                then Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
0
                else
                    if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
b
                        then Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
1
                        else do
                            l <- Distribution -> Double -> Either MathError Double
cdf (Double -> Double -> Distribution
Normal Double
mu Double
sd) Double
a
                            u <- cdf (Normal mu sd) b
                            v <- cdf (Normal mu sd) x
                            require (u > l) "unresolved truncation probability"
                            pure ((v - l) / (u - l))
        Empirical [Double]
xs -> Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ((Double -> Bool) -> [Double] -> [Double]
forall a. (a -> Bool) -> [a] -> [a]
filter (Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
x) [Double]
xs)) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs))
        Mixture [(Double, Distribution)]
ps -> [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Double] -> Double)
-> Either MathError [Double] -> Either MathError Double
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Double, Distribution) -> Either MathError Double)
-> [(Double, Distribution)] -> Either MathError [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
w, Distribution
m) -> (Double
w Double -> Double -> Double
forall a. Num a => a -> a -> a
*) (Double -> Double)
-> Either MathError Double -> Either MathError Double
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Distribution -> Double -> Either MathError Double
cdf Distribution
m Double
x) [(Double, Distribution)]
ps
        GaussianKDE Double
h [Double]
xs -> do ps <- (Double -> Either MathError Double)
-> [Double] -> Either MathError [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
mu -> Distribution -> Double -> Either MathError Double
cdf (Double -> Double -> Distribution
Normal Double
mu Double
h) Double
x) [Double]
xs; pure (sum ps / fromIntegral (length xs))
        Skellam Double
a Double
b -> do
            let low :: Double
low = Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
14 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
sqrt (Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
b) Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
20; high :: Double
high = Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
14 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
sqrt (Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
b) Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
20
            if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
low
                then Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
0
                else
                    if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
high
                        then Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
1
                        else do
                            let start :: Integer
start = Double -> Integer
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor Double
low :: Integer; end :: Integer
end = Double -> Integer
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor Double
x
                            [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Double] -> Double)
-> Either MathError [Double] -> Either MathError Double
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Integer -> Either MathError Double)
-> [Integer] -> Either MathError [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 (Distribution -> Double -> Either MathError Double
density Distribution
d (Double -> Either MathError Double)
-> (Integer -> Double) -> Integer -> Either MathError Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Double
forall a. Num a => Integer -> a
fromInteger) [Integer
start .. Integer
end]
        VonMises Double
mu Double
_ -> if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
mu Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
forall a. Floating a => a
pi then Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
0 else if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
mu Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
forall a. Floating a => a
pi then Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
1 else Distribution -> Double -> Double -> Either MathError Double
integrateDensity Distribution
d (Double
mu Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
forall a. Floating a => a
pi) Double
x
        Distribution
_
            | Distribution -> Bool
isDiscrete Distribution
d -> do
                let (Double
lo, Double
hi) = Distribution -> (Double, Double)
support Distribution
d
                if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
lo
                    then Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
0
                    else
                        if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
hi
                            then Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
1
                            else do
                                let start :: Integer
start = Double -> Integer
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
ceiling Double
lo :: Integer; end :: Integer
end = Double -> Integer
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor Double
x; count :: Integer
count = Integer
end Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
start Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1
                                Bool -> String -> Either MathError ()
require (Integer
count Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
1000000) String
"discrete CDF summation budget exceeded"
                                [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Double] -> Double)
-> Either MathError [Double] -> Either MathError Double
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Integer -> Either MathError Double)
-> [Integer] -> Either MathError [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 (Distribution -> Double -> Either MathError Double
density Distribution
d (Double -> Either MathError Double)
-> (Integer -> Double) -> Integer -> Either MathError Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Double
forall a. Num a => Integer -> a
fromInteger) [Integer
start .. Integer
end]
            | Bool
otherwise -> let (Double
lo, Double
_) = Distribution -> (Double, Double)
support Distribution
d in if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
lo then Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
0 else Distribution -> Double -> Double -> Either MathError Double
integrateDensity Distribution
d Double
lo Double
x
    _ <- checked p
    pure (max 0 (min 1 p))

-- | P(X>x), using direct upper tails for normal, gamma, beta, t, F, binomial and Poisson to avoid cancellation. Other families use the CDF complement.
survival :: Distribution -> Double -> Either MathError Double
survival :: Distribution -> Double -> Either MathError Double
survival Distribution
d Double
x = do
    Distribution -> Either MathError ()
validateDistribution Distribution
d
    Bool -> String -> Either MathError ()
require (Double -> Bool
finite Double
x) String
"survival finite argument required"
    case Distribution
d of
        Normal Double
mu Double
sd -> (Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2) (Double -> Double)
-> Either MathError Double -> Either MathError Double
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Double -> Either MathError Double
S.erfc ((Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
mu) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
sd Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
sqrt Double
2))
        LogNormal Double
mu Double
sd -> if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 then Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
1 else Distribution -> Double -> Either MathError Double
survival (Double -> Double -> Distribution
Normal Double
mu Double
sd) (Double -> Double
forall a. Floating a => a -> a
log Double
x)
        Beta Double
a Double
b -> if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 then Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
1 else if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
1 then Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
0 else Double -> Double -> Double -> Either MathError Double
S.betaI Double
b Double
a (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
x)
        StudentT Double
nu -> do p <- Double -> Double -> Double -> Either MathError Double
S.betaI (Double
nu Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2) Double
0.5 (Double
nu Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
nu Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x)); pure (if x >= 0 then p / 2 else 1 - p / 2)
        FisherF Double
a Double
b -> if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 then Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
1 else Double -> Double -> Double -> Either MathError Double
S.betaI (Double
b Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2) (Double
a Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2) (Double
b Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x))
        Gamma Double
a Double
sc -> if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 then Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
1 else Double -> Double -> Either MathError Double
S.gammaQ Double
a (Double
x Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
sc)
        ChiSquare Double
k -> Distribution -> Double -> Either MathError Double
survival (Double -> Double -> Distribution
Gamma (Double
k Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2) Double
2) Double
x
        Exponential Double
r -> Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 then Double
1 else Double -> Double
forall a. Floating a => a -> a
exp (-Double
r Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x))
        Poisson Double
r -> if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0 then Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
1 else Double -> Double -> Either MathError Double
S.gammaP (Integer -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Double -> Integer
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor Double
x :: Integer) Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1) Double
r
        Binomial Int
n Double
p ->
            if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0
                then Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
1
                else
                    if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n
                        then Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
0
                        else
                            if Double
p Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0
                                then Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
0
                                else
                                    if Double
p Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
1
                                        then Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
1
                                        else
                                            Double -> Double -> Double -> Either MathError Double
S.betaI
                                                (Integer -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Double -> Integer
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor Double
x :: Integer) Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1)
                                                (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n Double -> Double -> Double
forall a. Num a => a -> a -> a
- Integer -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Double -> Integer
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor Double
x :: Integer))
                                                Double
p
        Distribution
_ -> (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
-) (Double -> Double)
-> Either MathError Double -> Either MathError Double
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Distribution -> Double -> Either MathError Double
cdf Distribution
d Double
x

support :: Distribution -> (Double, Double)
support :: Distribution -> (Double, Double)
support Distribution
d = case Distribution
d of
    Bernoulli Double
_ -> (Double
0, Double
1)
    Binomial Int
n Double
_ -> (Double
0, Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n)
    NegativeBinomial Double
_ Double
_ -> (Double
0, Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0)
    Geometric Double
_ -> (Double
0, Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0)
    Hypergeometric Int
n Int
k Int
draws -> (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int
draws Int -> Int -> Int
forall a. Num a => a -> a -> a
- (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
k))), Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
draws Int
k))
    Poisson Double
_ -> (Double
0, Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0)
    Categorical [Double]
ps -> (Double
0, Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
ps Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))
    DiscreteUniform Int
a Int
b -> (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
a, Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
b)
    Uniform Double
a Double
b -> (Double
a, Double
b)
    Beta Double
_ Double
_ -> (Double
0, Double
1)
    Pareto Double
xm Double
_ -> (Double
xm, Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0)
    BetaBinomial Int
n Double
_ Double
_ -> (Double
0, Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n)
    Triangular Double
a Double
_ Double
b -> (Double
a, Double
b)
    TruncatedNormal Double
_ Double
_ Double
a Double
b -> (Double
a, Double
b)
    VonMises Double
mu Double
_ -> (Double
mu Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
forall a. Floating a => a
pi, Double
mu Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
forall a. Floating a => a
pi)
    Empirical [Double]
xs -> ([Double] -> Double
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [Double]
xs, [Double] -> Double
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [Double]
xs)
    Skellam Double
_ Double
_ -> (-Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0, Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0)
    Normal Double
_ Double
_ -> (-Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0, Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0)
    StudentT Double
_ -> (-Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0, Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0)
    Cauchy Double
_ Double
_ -> (-Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0, Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0)
    Laplace Double
_ Double
_ -> (-Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0, Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0)
    Logistic Double
_ Double
_ -> (-Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0, Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0)
    Gumbel Double
_ Double
_ -> (-Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0, Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0)
    GeneralizedExtremeValue Double
mu Double
sc Double
xi ->
        if Double
xi Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0 then (Double
mu Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
sc Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
xi, Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0) else if Double
xi Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0 then (-Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0, Double
mu Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
sc Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
xi) else (-Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0, Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0)
    GaussianKDE Double
_ [Double]
_ -> (-Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0, Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0)
    Mixture [(Double, Distribution)]
ps -> ([Double] -> Double
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum (((Double, Distribution) -> Double)
-> [(Double, Distribution)] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map ((Double, Double) -> Double
forall a b. (a, b) -> a
fst ((Double, Double) -> Double)
-> ((Double, Distribution) -> (Double, Double))
-> (Double, Distribution)
-> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Distribution -> (Double, Double)
support (Distribution -> (Double, Double))
-> ((Double, Distribution) -> Distribution)
-> (Double, Distribution)
-> (Double, Double)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Double, Distribution) -> Distribution
forall a b. (a, b) -> b
snd) [(Double, Distribution)]
ps), [Double] -> Double
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum (((Double, Distribution) -> Double)
-> [(Double, Distribution)] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map ((Double, Double) -> Double
forall a b. (a, b) -> b
snd ((Double, Double) -> Double)
-> ((Double, Distribution) -> (Double, Double))
-> (Double, Distribution)
-> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Distribution -> (Double, Double)
support (Distribution -> (Double, Double))
-> ((Double, Distribution) -> Distribution)
-> (Double, Distribution)
-> (Double, Double)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Double, Distribution) -> Distribution
forall a b. (a, b) -> b
snd) [(Double, Distribution)]
ps))
    Distribution
_ -> (Double
0, Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0)
integrateDensity :: Distribution -> Double -> Double -> Either MathError Double
integrateDensity :: Distribution -> Double -> Double -> Either MathError Double
integrateDensity Distribution
d Double
a Double
b = do
    let n :: Int
n = Int
2048 :: Int; h :: Double
h = (Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
a) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n
    Bool -> String -> Either MathError ()
require (Double -> Bool
finite Double
h Bool -> Bool -> Bool
&& Double
h Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
0) String
"integration requires finite ordered bounds"
    ys <- (Int -> Either MathError Double)
-> [Int] -> Either MathError [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 (Distribution -> Double -> Either MathError Double
density Distribution
d (Double -> Either MathError Double)
-> (Int -> Double) -> Int -> Either MathError Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+) (Double -> Double) -> (Int -> Double) -> Int -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Double
h Double -> Double -> Double
forall a. Num a => a -> a -> a
*) (Double -> Double) -> (Int -> Double) -> Int -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral) [Int
0 .. Int
n]
    checked
        ( h
            / 3
            * sum [if i == 0 || i == n then y else if odd i then 4 * y else 2 * y | (i, y) <- zip [0 ..] ys]
        )

-- | Generalized inverse CDF, p in [0,1]; endpoints return mathematical support bounds, including infinities. Interior values use analytic transforms or bounded bracketing/bisection.
quantile :: Distribution -> Double -> Either MathError Double
quantile :: Distribution -> Double -> Either MathError Double
quantile Distribution
d Double
p = do
    Distribution -> Either MathError ()
validateDistribution Distribution
d
    Bool -> String -> Either MathError ()
require (Double -> Bool
prob Double
p) String
"quantile p outside [0,1]"
    let (Double
a, Double
b) = Distribution -> (Double, Double)
support Distribution
d
    if Double
p Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0
        then Double -> Either MathError Double
forall a b. b -> Either a b
Right Double
a
        else
            if Double
p Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
1
                then Double -> Either MathError Double
forall a b. b -> Either a b
Right Double
b
                else case Distribution
d of
                    Normal Double
mu Double
sd -> (\Double
z -> Double
mu Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
sd Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
z) (Double -> Double)
-> Either MathError Double -> Either MathError Double
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Double -> Either MathError Double
S.normalQuantile Double
p
                    LogNormal Double
mu Double
sd -> Double -> Double
forall a. Floating a => a -> a
exp (Double -> Double)
-> Either MathError Double -> Either MathError Double
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Distribution -> Double -> Either MathError Double
quantile (Double -> Double -> Distribution
Normal Double
mu Double
sd) Double
p
                    Uniform Double
lo Double
hi -> Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double
lo Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
p Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
hi Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
lo))
                    Exponential Double
r -> Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (-Double -> Double
forall a. Floating a => a -> a
log (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
r)
                    Empirical [Double]
xs -> Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Double] -> [Double]
forall a. Ord a => [a] -> [a]
sort [Double]
xs [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
ceiling (Double
p Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs)) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))
                    Distribution
_ -> do
                        lo <- if Double -> Bool
finite Double
a then Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
a else Double
-> (Double -> Either MathError Bool) -> Either MathError Double
forall {t}.
Num t =>
t -> (t -> Either MathError Bool) -> Either MathError t
bracket (-Double
1) (\Double
v -> (Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
p) (Double -> Bool)
-> Either MathError Double -> Either MathError Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Distribution -> Double -> Either MathError Double
cdf Distribution
d Double
v)
                        hi <- if finite b && b < 1000000 then pure b else bracket 1 (\Double
v -> (Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
p) (Double -> Bool)
-> Either MathError Double -> Either MathError Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Distribution -> Double -> Either MathError Double
cdf Distribution
d Double
v)
                        if isDiscrete d then discreteSearch (ceiling lo) (ceiling hi) else continuousSearch 80 lo hi
  where
    bracket :: t -> (t -> Either MathError Bool) -> Either MathError t
bracket t
v t -> Either MathError Bool
predicate = Int -> t -> Either MathError t
forall {t}. (Eq t, Num t) => t -> t -> Either MathError t
go (Int
64 :: Int) t
v
      where
        go :: t -> t -> Either MathError t
go t
0 t
_ = MathError -> Either MathError t
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"quantile bracket" Int
64)
        go t
k t
x = do ok <- t -> Either MathError Bool
predicate t
x; if ok then pure x else go (k - 1) (2 * x)
    discreteSearch :: Integer -> Integer -> Either MathError Double
    discreteSearch :: Integer -> Integer -> Either MathError Double
discreteSearch Integer
lo Integer
hi
        | Integer
lo Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Integer
hi = Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Integer -> Double
forall a. Num a => Integer -> a
fromInteger Integer
lo)
        | Bool
otherwise = do
            let mid :: Integer
mid = Integer
lo Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ (Integer
hi Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
lo) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
2
            v <- Distribution -> Double -> Either MathError Double
cdf Distribution
d (Integer -> Double
forall a. Num a => Integer -> a
fromInteger Integer
mid)
            if v >= p then discreteSearch lo mid else discreteSearch (mid + 1) hi
    continuousSearch :: Int -> Double -> Double -> Either MathError Double
    continuousSearch :: Int -> Double -> Double -> Either MathError Double
continuousSearch Int
0 Double
lo Double
hi = Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Double
lo Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
hi) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2)
    continuousSearch Int
n Double
lo Double
hi = do
        let m :: Double
m = (Double
lo Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
hi) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2
        v <- Distribution -> Double -> Either MathError Double
cdf Distribution
d Double
m
        if v >= p then continuousSearch (n - 1) lo m else continuousSearch (n - 1) m hi

-- | Density divided by survival. For zero survival returns positive infinity by convention.
hazard :: Distribution -> Double -> Either MathError Double
hazard :: Distribution -> Double -> Either MathError Double
hazard Distribution
d Double
x = do f <- Distribution -> Double -> Either MathError Double
density Distribution
d Double
x; s <- survival d x; pure (if s == 0 then 1 / 0 else f / s)

-- | Sum of independent observation log probabilities. Impossible data produce negative infinity.
logLikelihood :: Distribution -> [Double] -> Either MathError Double
logLikelihood :: Distribution -> [Double] -> Either MathError Double
logLikelihood Distribution
d [Double]
xs = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Double] -> Double)
-> Either MathError [Double] -> Either MathError Double
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Double -> Either MathError Double)
-> [Double] -> Either MathError [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 (Distribution -> Double -> Either MathError Double
logDensity Distribution
d) [Double]
xs

{- | Mean and population variance. Infinite/nonexistent moments return a domain
error rather than a numeric value that could be mistaken for a finite moment.
-}
moments :: Distribution -> Either MathError (Double, Double)
moments :: Distribution -> Either MathError (Double, Double)
moments Distribution
d = do
    Distribution -> Either MathError ()
validateDistribution Distribution
d
    pair <- case Distribution
d of
        Bernoulli Double
p -> (Double, Double) -> Either MathError (Double, Double)
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double
p, Double
p Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
p))
        Binomial Int
n Double
p -> (Double, Double) -> Either MathError (Double, Double)
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
p, Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
p Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
p))
        NegativeBinomial Double
r Double
p -> (Double, Double) -> Either MathError (Double, Double)
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (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
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))
        Geometric Double
p -> (Double, Double) -> Either MathError (Double, Double)
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((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
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))
        Hypergeometric Int
n Int
k Int
draws ->
            let p :: Double
p = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n; m :: Double
m = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
draws Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
p
             in (Double, Double) -> Either MathError (Double, Double)
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double
m, if Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 then Double
0 else Double
m 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. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
draws) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))
        Poisson Double
r -> (Double, Double) -> Either MathError (Double, Double)
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double
r, Double
r)
        Categorical [Double]
ps ->
            let m :: Double
m = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) [Double
0 ..] [Double]
ps)
             in (Double, Double) -> Either MathError (Double, Double)
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double
m, [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Double
i Double
p -> Double
p Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
i Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
m) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int)) [Double
0 ..] [Double]
ps))
        DiscreteUniform Int
a Int
b ->
            (Double, Double) -> Either MathError (Double, Double)
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
                ( (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
b) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2
                , ((Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
b Double -> Double -> Double
forall a. Num a => a -> a -> a
- Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int) Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
12
                )
        Uniform Double
a Double
b -> (Double, Double) -> Either MathError (Double, Double)
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
b) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2, (Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
a) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
12)
        Normal Double
mu Double
sd -> (Double, Double) -> Either MathError (Double, Double)
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double
mu, Double
sd Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sd)
        LogNormal Double
mu Double
sd -> let m :: Double
m = Double -> Double
forall a. Floating a => a -> a
exp (Double
mu Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
sd Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sd Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2) in (Double, Double) -> Either MathError (Double, Double)
checkedPair (Double
m, Double
m Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
m Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double -> Double
forall a. Floating a => a -> a
exp (Double
sd Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sd) Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1))
        Exponential Double
r -> (Double, Double) -> Either MathError (Double, Double)
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
r, Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
r Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
r))
        Gamma Double
a Double
sc -> (Double, Double) -> Either MathError (Double, Double)
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sc, Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sc Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sc)
        Beta Double
a Double
b -> (Double, Double) -> Either MathError (Double, Double)
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double
a Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
b), Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
b Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ ((Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
b) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1)))
        ChiSquare Double
k -> (Double, Double) -> Either MathError (Double, Double)
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double
k, Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
k)
        StudentT Double
nu -> do Bool -> String -> Either MathError ()
require (Double
nu Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
2) String
"Student-t variance needs nu>2"; (Double, Double) -> Either MathError (Double, Double)
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double
0, Double
nu Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
nu Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
2))
        FisherF Double
a Double
b -> do
            Bool -> String -> Either MathError ()
require (Double
b Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
4) String
"F variance needs d2>4"
            (Double, Double) -> Either MathError (Double, Double)
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double
b Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
2), Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
2) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
2) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
4)))
        Cauchy Double
_ Double
_ -> MathError -> Either MathError (Double, Double)
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"Cauchy moments do not exist")
        Laplace Double
mu Double
sc -> (Double, Double) -> Either MathError (Double, Double)
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double
mu, Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sc Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sc)
        Logistic Double
mu Double
sc -> (Double, Double) -> Either MathError (Double, Double)
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double
mu, Double
forall a. Floating a => a
pi Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
forall a. Floating a => a
pi Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sc Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sc Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
3)
        Weibull Double
k Double
sc -> do a <- Double -> Either MathError Double
S.gamma (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
k); b <- S.gamma (1 + 2 / k); pure (sc * a, sc * sc * (b - a * a))
        Pareto Double
xm Double
a -> do
            Bool -> String -> Either MathError ()
require (Double
a Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
2) String
"Pareto variance needs shape>2"
            (Double, Double) -> Either MathError (Double, Double)
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
xm Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1), Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
xm Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
xm Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ ((Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
2)))
        Rayleigh Double
sc -> (Double, Double) -> Either MathError (Double, Double)
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double
sc Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
sqrt (Double
forall a. Floating a => a
pi Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2), (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
forall a. Floating a => a
pi Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sc Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sc)
        Maxwell Double
sc -> (Double, Double) -> Either MathError (Double, Double)
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sc Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
sqrt (Double
2 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
forall a. Floating a => a
pi), (Double
3 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
8 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
forall a. Floating a => a
pi) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sc Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sc)
        Gumbel Double
mu Double
sc -> (Double, Double) -> Either MathError (Double, Double)
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double
mu Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
0.5772156649015329 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sc, Double
forall a. Floating a => a
pi Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
forall a. Floating a => a
pi Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sc Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sc Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
6)
        GeneralizedExtremeValue Double
mu Double
sc Double
xi ->
            if Double -> Double
forall a. Num a => a -> a
abs Double
xi Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
1e-10
                then Distribution -> Either MathError (Double, Double)
moments (Double -> Double -> Distribution
Gumbel Double
mu Double
sc)
                else do
                    Bool -> String -> Either MathError ()
require (Double
xi Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0.5) String
"GEV variance needs xi<1/2"
                    g1 <- Double -> Either MathError Double
S.gamma (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
xi)
                    g2 <- S.gamma (1 - 2 * xi)
                    pure (mu + sc * (g1 - 1) / xi, sc * sc / (xi * xi) * (g2 - g1 * g1))
        BetaBinomial Int
n Double
a Double
b ->
            (Double, Double) -> Either MathError (Double, Double)
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
                ( Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
a Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
b)
                , Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ ((Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
b) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1))
                )
        Skellam Double
a Double
b -> (Double, Double) -> Either MathError (Double, Double)
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
b, Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
b)
        Triangular Double
a Double
b Double
c -> (Double, Double) -> Either MathError (Double, Double)
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
c) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
3, (Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
c Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
c Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
c Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
c) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
18)
        Empirical [Double]
xs ->
            let m :: Double
m = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
xs Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs)
             in (Double, Double) -> Either MathError (Double, Double)
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double
m, [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]
map (\Double
x -> (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
m) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int)) [Double]
xs) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs))
        GaussianKDE Double
h [Double]
xs -> do (m, v) <- Distribution -> Either MathError (Double, Double)
moments ([Double] -> Distribution
Empirical [Double]
xs); pure (m, v + h * h)
        Mixture [(Double, Distribution)]
ps -> do
            ms <- ((Double, Distribution)
 -> Either MathError (Double, Double, Double))
-> [(Double, Distribution)]
-> Either MathError [(Double, Double, 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
w, Distribution
m) -> do (u, v) <- Distribution -> Either MathError (Double, Double)
moments Distribution
m; pure (w, u, v)) [(Double, Distribution)]
ps
            let mean = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double
w Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
u | (Double
w, Double
u, Double
_) <- [(Double, Double, Double)]
ms]
             in pure (mean, sum [w * (v + (u - mean) ^ (2 :: Int)) | (w, u, v) <- ms])
        Distribution
_ -> do m <- Distribution -> (Double -> Double) -> Either MathError Double
expectation Distribution
d Double -> Double
forall a. a -> a
id; second <- expectation d (\Double
x -> Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x); pure (m, max 0 (second - m * m))
    checkedPair pair
  where
    checkedPair :: (Double, Double) -> Either MathError (Double, Double)
checkedPair (Double
a, Double
b) = do _ <- Double -> Either MathError Double
checked Double
a; _ <- checked b; pure (a, b)

-- | Shannon or differential entropy in nats. Closed forms are used when available; other families use the documented expectation quadrature.
entropy :: Distribution -> Either MathError Double
entropy :: Distribution -> Either MathError Double
entropy Distribution
d = do
    Distribution -> Either MathError ()
validateDistribution Distribution
d
    case Distribution
d of
        Normal Double
_ Double
sd -> Double -> Either MathError Double
checked (Double
0.5 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
log (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
forall a. Floating a => a
pi Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
exp Double
1) Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double -> Double
forall a. Floating a => a -> a
log Double
sd)
        Uniform Double
a Double
b -> Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double -> Double
forall a. Floating a => a -> a
log (Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
a))
        Exponential Double
r -> Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double -> Double
forall a. Floating a => a -> a
log Double
r)
        Bernoulli Double
p -> Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (-Double -> Double -> Double
xlog Double
p Double
p Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double -> Double -> Double
xlog (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
p) (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
p))
        Categorical [Double]
ps -> Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (-[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]
map (\Double
p -> Double -> Double -> Double
xlog Double
p Double
p) [Double]
ps))
        Distribution
_ ->
            Distribution -> (Double -> Double) -> Either MathError Double
expectation
                Distribution
d
                (\Double
x -> case Distribution -> Double -> Either MathError Double
logDensity Distribution
d Double
x of Right Double
lp -> if Double -> Bool
forall a. RealFloat a => a -> Bool
isInfinite Double
lp then Double
0 else -Double
lp; Left MathError
_ -> Double
0 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0)

-- | Reference expectation under a distribution. Empirical/mixture cases are exact weighted sums; count families truncate at 1e-10 tail quantiles, continuous families use 4096 midpoint probability-space nodes. No certified quadrature error is claimed.
expectation :: Distribution -> (Double -> Double) -> Either MathError Double
expectation :: Distribution -> (Double -> Double) -> Either MathError Double
expectation Distribution
d Double -> Double
f = do
    Distribution -> Either MathError ()
validateDistribution Distribution
d
    case Distribution
d of
        Empirical [Double]
xs -> Double -> Either MathError Double
checked ([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]
map Double -> Double
f [Double]
xs) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs))
        Mixture [(Double, Distribution)]
pairs -> do values <- ((Double, Distribution) -> Either MathError Double)
-> [(Double, Distribution)] -> Either MathError [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
w, Distribution
m) -> (Double
w Double -> Double -> Double
forall a. Num a => a -> a -> a
*) (Double -> Double)
-> Either MathError Double -> Either MathError Double
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Distribution -> (Double -> Double) -> Either MathError Double
expectation Distribution
m Double -> Double
f) [(Double, Distribution)]
pairs; checked (sum values)
        Distribution
_ -> Either MathError Double
numeric
  where
    numeric :: Either MathError Double
numeric =
        if Distribution -> Bool
isDiscrete Distribution
d
            then do
                lo <- Distribution -> Double -> Either MathError Double
quantile Distribution
d Double
1e-10
                hi <- quantile d (1 - 1e-10)
                let a = Double -> Integer
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
ceiling Double
lo :: Integer; b = Double -> Integer
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor Double
hi
                require (b - a <= 1000000) "expectation summation budget"
                terms <- mapM (\Integer
k -> do p <- Distribution -> Double -> Either MathError Double
density Distribution
d (Integer -> Double
forall a. Num a => Integer -> a
fromInteger Integer
k); pure (p * f (fromInteger k))) [a .. b]
                checked (sum terms)
            else do
                -- Midpoint integration over probability space handles unbounded supports.
                xs <- (Int -> Either MathError Double)
-> [Int] -> Either MathError [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 (\Int
i -> Distribution -> Double -> Either MathError Double
quantile Distribution
d ((Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
0.5) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
4096)) [Int
0 .. Int
4095 :: Int]
                checked (sum (map f xs) / 4096)

-- | Draw a model sample and updated immutable RNG; specialized gamma/normal/count samplers are reused, with inverse-CDF sampling for remaining families.
sample :: Distribution -> R.RNG -> Either MathError (Double, R.RNG)
sample :: Distribution -> RNG -> Either MathError (Double, RNG)
sample Distribution
d RNG
r = do
    Distribution -> Either MathError ()
validateDistribution Distribution
d
    case Distribution
d of
        Normal Double
mu Double
sd -> Double -> Double -> RNG -> Either MathError (Double, RNG)
R.normal Double
mu Double
sd RNG
r
        Gamma Double
a Double
sc -> Double -> Double -> RNG -> Either MathError (Double, RNG)
R.gamma Double
a Double
sc RNG
r
        Beta Double
a Double
b -> Double -> Double -> RNG -> Either MathError (Double, RNG)
R.beta Double
a Double
b RNG
r
        Exponential Double
rate -> Double -> RNG -> Either MathError (Double, RNG)
R.exponential Double
rate RNG
r
        Poisson Double
rate -> do (x, s) <- Double -> RNG -> Either MathError (Int, RNG)
R.poisson Double
rate RNG
r; pure (fromIntegral x, s)
        Binomial Int
n Double
p -> do (x, s) <- Int -> Double -> RNG -> Either MathError (Int, RNG)
R.binomial Int
n Double
p RNG
r; pure (fromIntegral x, s)
        Categorical [Double]
ps -> do (x, s) <- [Double] -> RNG -> Either MathError (Int, RNG)
R.weighted [Double]
ps RNG
r; pure (fromIntegral x, s)
        Mixture [(Double, Distribution)]
ps -> do (i, s) <- [Double] -> RNG -> Either MathError (Int, RNG)
R.weighted (((Double, Distribution) -> Double)
-> [(Double, Distribution)] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double, Distribution) -> Double
forall a b. (a, b) -> a
fst [(Double, Distribution)]
ps) RNG
r; sample (snd (ps !! i)) s
        GaussianKDE Double
h [Double]
xs -> do (i, s) <- Int -> Int -> RNG -> Either MathError (Int, RNG)
R.uniformInt Int
0 ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) RNG
r; R.normal (xs !! i) h s
        Empirical [Double]
xs -> do (i, s) <- Int -> Int -> RNG -> Either MathError (Int, RNG)
R.uniformInt Int
0 ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) RNG
r; pure (xs !! i, s)
        Distribution
_ -> let (Double
u, RNG
s) = RNG -> (Double, RNG)
R.uniformOpen RNG
r in do x <- Distribution -> Double -> Either MathError Double
quantile Distribution
d Double
u; pure (x, s)

-- | Normal maximum-likelihood fit; variance divides by n. Constant data are rejected because this continuous family requires positive standard deviation.
fitNormal :: [Double] -> Either MathError Distribution
fitNormal :: [Double] -> Either MathError Distribution
fitNormal [Double]
xs = do
    Bool -> String -> Either MathError ()
require (Bool -> Bool
not ([Double] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Double]
xs) Bool -> Bool -> Bool
&& (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Double -> Bool
finite [Double]
xs) String
"normal fit needs finite nonempty data"
    let m :: Double
m = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
xs Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs)
        v :: Double
v = [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]
map (\Double
x -> (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
m) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int)) [Double]
xs) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs)
        d :: Distribution
d = Double -> Double -> Distribution
Normal Double
m (Double -> Double
forall a. Floating a => a -> a
sqrt Double
v)
     in Distribution -> Either MathError ()
validateDistribution Distribution
d Either MathError ()
-> Either MathError Distribution -> Either MathError Distribution
forall a b.
Either MathError a -> Either MathError b -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Distribution -> Either MathError Distribution
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Distribution
d

-- | Maximum-likelihood rate for nonnegative observations with strictly positive total.
fitExponential :: [Double] -> Either MathError Distribution
fitExponential :: [Double] -> Either MathError Distribution
fitExponential [Double]
xs = do
    Bool -> String -> Either MathError ()
require
        (Bool -> Bool
not ([Double] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Double]
xs) Bool -> Bool -> Bool
&& (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\Double
x -> Double -> Bool
finite Double
x Bool -> Bool -> Bool
&& Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
0) [Double]
xs Bool -> Bool -> Bool
&& [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
xs Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0)
        String
"exponential fit nonnegative data with positive sum required"
    Distribution -> Either MathError Distribution
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double -> Distribution
Exponential (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
xs))

-- | Maximum-likelihood probability for nonempty binary observations.
fitBernoulli :: [Double] -> Either MathError Distribution
fitBernoulli :: [Double] -> Either MathError Distribution
fitBernoulli [Double]
xs = do
    Bool -> String -> Either MathError ()
require (Bool -> Bool
not ([Double] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Double]
xs) Bool -> Bool -> Bool
&& (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\Double
x -> Double
x Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 Bool -> Bool -> Bool
|| Double
x Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
1) [Double]
xs) String
"Bernoulli fit requires binary data"
    Distribution -> Either MathError Distribution
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double -> Distribution
Bernoulli ([Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
xs Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs)))

-- | Select the maximum-likelihood model from an explicit finite candidate grid; this is grid search, not continuous parameter optimization.
fitMLEGrid :: [Distribution] -> [Double] -> Either MathError Distribution
fitMLEGrid :: [Distribution] -> [Double] -> Either MathError Distribution
fitMLEGrid [Distribution]
ds [Double]
xs = [Distribution]
-> [Double]
-> (Distribution -> Double)
-> Either MathError Distribution
fitMAPGrid [Distribution]
ds [Double]
xs (Double -> Distribution -> Double
forall a b. a -> b -> a
const Double
0)

-- | Select the maximum log-likelihood plus caller-supplied log prior over a finite model grid. Log priors must be finite.
fitMAPGrid ::
    [Distribution] -> [Double] -> (Distribution -> Double) -> Either MathError Distribution
fitMAPGrid :: [Distribution]
-> [Double]
-> (Distribution -> Double)
-> Either MathError Distribution
fitMAPGrid [Distribution]
ds [Double]
xs Distribution -> Double
prior = do
    Bool -> String -> Either MathError ()
require (Bool -> Bool
not ([Distribution] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Distribution]
ds) Bool -> Bool -> Bool
&& Bool -> Bool
not ([Double] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Double]
xs)) String
"grid fit requires models and observations"
    scores <-
        (Distribution -> Either MathError (Distribution, Double))
-> [Distribution] -> Either MathError [(Distribution, 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
            ( \Distribution
d -> do
                l <- Distribution -> [Double] -> Either MathError Double
logLikelihood Distribution
d [Double]
xs
                let p = Distribution -> Double
prior Distribution
d in require (finite p) "nonfinite log prior" >> pure (d, l + p)
            )
            [Distribution]
ds
    pure (fst (maximumBy (comparing snd) scores))

-- | Conjugate Beta prior update from success/failure observations, returning posterior concentrations (alpha,beta).
betaBernoulliUpdate :: Double -> Double -> [Bool] -> Either MathError (Double, Double)
betaBernoulliUpdate :: Double -> Double -> [Bool] -> Either MathError (Double, Double)
betaBernoulliUpdate Double
a Double
b [Bool]
xs = do
    Bool -> String -> Either MathError ()
require (Double -> Bool
positive Double
a Bool -> Bool -> Bool
&& Double -> Bool
positive Double
b) String
"positive beta prior required"
    (Double, Double) -> Either MathError (Double, Double)
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Bool] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ((Bool -> Bool) -> [Bool] -> [Bool]
forall a. (a -> Bool) -> [a] -> [a]
filter Bool -> Bool
forall a. a -> a
id [Bool]
xs)), Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Bool] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ((Bool -> Bool) -> [Bool] -> [Bool]
forall a. (a -> Bool) -> [a] -> [a]
filter Bool -> Bool
not [Bool]
xs)))

-- | Conjugate Poisson-rate update; arguments and result use Gamma shape/rate, deliberately distinct from the Gamma distribution constructor shape/scale.
gammaPoissonUpdate :: Double -> Double -> [Int] -> Either MathError (Double, Double)
gammaPoissonUpdate :: Double -> Double -> [Int] -> Either MathError (Double, Double)
gammaPoissonUpdate Double
shapePrior Double
ratePrior [Int]
xs = do
    Bool -> String -> Either MathError ()
require
        (Double -> Bool
positive Double
shapePrior Bool -> Bool -> Bool
&& Double -> Bool
positive Double
ratePrior Bool -> Bool -> Bool
&& (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0) [Int]
xs)
        String
"positive Gamma shape/rate prior and counts required"
    (Double, Double) -> Either MathError (Double, Double)
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double
shapePrior Double -> Double -> Double
forall a. Num a => a -> a -> a
+ [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Int -> Double) -> [Int] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral [Int]
xs), Double
ratePrior Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
xs))

-- | Known-variance conjugate normal mean update. Arguments are prior mean, prior variance, observation variance and observations; returns posterior mean/variance.
normalNormalUpdate :: Double -> Double -> Double -> [Double] -> Either MathError (Double, Double)
normalNormalUpdate :: Double
-> Double
-> Double
-> [Double]
-> Either MathError (Double, Double)
normalNormalUpdate Double
mu Double
priorVariance Double
observationVariance [Double]
xs = do
    Bool -> String -> Either MathError ()
require
        (Double -> Bool
finite Double
mu Bool -> Bool -> Bool
&& Double -> Bool
positive Double
priorVariance Bool -> Bool -> Bool
&& Double -> Bool
positive Double
observationVariance Bool -> Bool -> Bool
&& (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Double -> Bool
finite [Double]
xs)
        String
"normal update requires finite mean and positive variances"
    let precision :: Double
precision = Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
priorVariance Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
observationVariance
     in (Double, Double) -> Either MathError (Double, Double)
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Double
mu Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
priorVariance Double -> Double -> Double
forall a. Num a => a -> a -> a
+ [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
xs Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
observationVariance) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
precision, Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
precision)

-- | Compute P(A|B) from joint P(A intersection B) and strictly positive P(B).
conditionalProbability :: Double -> Double -> Either MathError Double
conditionalProbability :: Double -> Double -> Either MathError Double
conditionalProbability Double
joint Double
given = do
    Bool -> String -> Either MathError ()
require
        (Double -> Bool
prob Double
joint Bool -> Bool -> Bool
&& Double -> Bool
positive Double
given Bool -> Bool -> Bool
&& Double
given Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
1 Bool -> Bool -> Bool
&& Double
joint Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
given)
        String
"conditional probability invalid joint/marginal"
    Double -> Either MathError Double
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double
joint Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
given)

-- | Normalize prior probabilities times observation likelihoods. A zero-evidence observation is rejected.
bayes :: [Double] -> [Double] -> Either MathError [Double]
bayes :: [Double] -> [Double] -> Either MathError [Double]
bayes [Double]
priors [Double]
likelihoods = do
    Bool -> String -> Either MathError ()
require
        ( [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
priors Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
likelihoods
            Bool -> Bool -> Bool
&& Bool -> Bool
not ([Double] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Double]
priors)
            Bool -> Bool -> Bool
&& (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Double -> Bool
prob [Double]
priors
            Bool -> Bool -> Bool
&& (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Double -> Bool
prob [Double]
likelihoods
            Bool -> Bool -> Bool
&& Double -> Double
forall a. Num a => a -> a
abs ([Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
priors Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1) Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
1e-12
        )
        String
"Bayes invalid probabilities"
    let ws :: [Double]
ws = (Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) [Double]
priors [Double]
likelihoods
     in Bool -> String -> Either MathError ()
require ([Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
ws Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0) String
"zero evidence" Either MathError ()
-> Either MathError [Double] -> Either MathError [Double]
forall a b.
Either MathError a -> Either MathError b -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> [Double] -> Either MathError [Double]
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
ws) [Double]
ws)

-- | Two-sided one-sample Kolmogorov distance against a specified CDF. Returns the distance only; no discrete-null p-value is implied.
goodnessOfFit :: Distribution -> [Double] -> Either MathError Double
goodnessOfFit :: Distribution -> [Double] -> Either MathError Double
goodnessOfFit Distribution
d [Double]
xs = do
    Bool -> String -> Either MathError ()
require (Bool -> Bool
not ([Double] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Double]
xs) Bool -> Bool -> Bool
&& (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Double -> Bool
finite [Double]
xs) String
"KS goodness needs finite observations"
    ps <- (Double -> Either MathError Double)
-> [Double] -> Either MathError [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 (Distribution -> Double -> Either MathError Double
cdf Distribution
d) ([Double] -> [Double]
forall a. Ord a => [a] -> [a]
sort [Double]
xs)
    let n = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs)
     in pure
            ( maximum
                [maximum [p - fromIntegral i / n, fromIntegral (i + 1) / n - p] | (i, p) <- zip [0 :: Int ..] ps]
            )

-- | Construct an equal-weight Gaussian kernel density with strictly positive bandwidth and finite nonempty observations.
gaussianKDE :: Double -> [Double] -> Either MathError Distribution
gaussianKDE :: Double -> [Double] -> Either MathError Distribution
gaussianKDE Double
h [Double]
xs = let d :: Distribution
d = Double -> [Double] -> Distribution
GaussianKDE Double
h [Double]
xs in Distribution -> Either MathError ()
validateDistribution Distribution
d Either MathError ()
-> Either MathError Distribution -> Either MathError Distribution
forall a b.
Either MathError a -> Either MathError b -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Distribution -> Either MathError Distribution
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Distribution
d

-- | Multinomial mass for probabilities summing to one and nonnegative count vector; log-gamma arithmetic avoids factorial overflow.
multinomialPMF :: [Double] -> [Int] -> Either MathError Double
multinomialPMF :: [Double] -> [Int] -> Either MathError Double
multinomialPMF [Double]
ps [Int]
counts = do
    Bool -> String -> Either MathError ()
require
        ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
ps Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
counts Bool -> Bool -> Bool
&& (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Double -> Bool
prob [Double]
ps Bool -> Bool -> Bool
&& Double -> Double
forall a. Num a => a -> a
abs ([Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
ps Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1) Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
1e-12 Bool -> Bool -> Bool
&& (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0) [Int]
counts)
        String
"multinomial invalid probabilities/counts"
    a <- Double -> Either MathError Double
S.logGamma (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Int]
counts) Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1)
    bs <- mapM (\Int
c -> Double -> Either MathError Double
S.logGamma (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
c Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1)) counts
    pure (exp (a - sum bs + sum (zipWith (\Int
c Double
p -> Double -> Double -> Double
xlog (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
c) Double
p) counts ps)))

-- | Dirichlet density on the strictly positive simplex interior. Concentration and coordinate vectors must have matching nonzero length.
dirichletDensity :: [Double] -> [Double] -> Either MathError Double
dirichletDensity :: [Double] -> [Double] -> Either MathError Double
dirichletDensity [Double]
as [Double]
xs = do
    Bool -> String -> Either MathError ()
require
        ( Bool -> Bool
not ([Double] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Double]
as)
            Bool -> Bool -> Bool
&& [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
as Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs
            Bool -> Bool -> Bool
&& (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Double -> Bool
positive [Double]
as
            Bool -> Bool -> Bool
&& (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Double -> Bool
positive [Double]
xs
            Bool -> Bool -> Bool
&& Double -> Double
forall a. Num a => a -> a
abs ([Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
xs Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1) Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
1e-10
        )
        String
"Dirichlet requires positive simplex interior and concentrations"
    a <- Double -> Either MathError Double
S.logGamma ([Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
as)
    bs <- mapM S.logGamma as
    checked (exp (a - sum bs + sum (zipWith (\Double
p Double
x -> (Double
p Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
log Double
x) as xs)))

-- | Log density for a finite mean/point and symmetric positive-definite covariance, validated by Cholesky.
multivariateNormalLogDensity :: [Double] -> Matrix Double -> [Double] -> Either MathError Double
multivariateNormalLogDensity :: [Double] -> Matrix Double -> [Double] -> Either MathError Double
multivariateNormalLogDensity [Double]
mu Matrix Double
covariance [Double]
x = do
    Bool -> String -> Either MathError ()
require
        ( Bool -> Bool
not ([Double] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Double]
mu)
            Bool -> Bool -> Bool
&& [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
mu Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
x
            Bool -> Bool -> Bool
&& Matrix Double -> (Int, Int)
forall a. Matrix a -> (Int, Int)
shape Matrix Double
covariance (Int, Int) -> (Int, Int) -> Bool
forall a. Eq a => a -> a -> Bool
== ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
mu, [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
mu)
            Bool -> Bool -> Bool
&& (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Double -> Bool
finite ([Double]
mu [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ [Double]
x)
        )
        String
"multivariate normal shape or input"
    _ <- Matrix Double -> Either MathError (Matrix Double)
forall a. RealFloat a => Matrix a -> Either MathError (Matrix a)
cholesky Matrix Double
covariance
    det <- determinant covariance
    z <- solve covariance (zipWith (-) x mu)
    checked
        (-0.5 * (fromIntegral (length mu) * log (2 * pi) + log det + sum (zipWith (*) z (zipWith (-) x mu))))

-- | Cholesky-transform independent standard normals; returns a vector and updated state.
sampleMultivariateNormal :: [Double] -> Matrix Double -> R.RNG -> Either MathError ([Double], R.RNG)
sampleMultivariateNormal :: [Double]
-> Matrix Double -> RNG -> Either MathError ([Double], RNG)
sampleMultivariateNormal [Double]
mu Matrix Double
covariance RNG
r = do
    Bool -> String -> Either MathError ()
require
        (Bool -> Bool
not ([Double] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Double]
mu) Bool -> Bool -> Bool
&& Matrix Double -> (Int, Int)
forall a. Matrix a -> (Int, Int)
shape Matrix Double
covariance (Int, Int) -> (Int, Int) -> Bool
forall a. Eq a => a -> a -> Bool
== ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
mu, [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
mu) Bool -> Bool -> Bool
&& (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Double -> Bool
finite [Double]
mu)
        String
"multivariate normal sampling shape"
    l <- Matrix Double -> Either MathError (Matrix Double)
forall a. RealFloat a => Matrix a -> Either MathError (Matrix a)
cholesky Matrix Double
covariance
    (zs, s) <- R.sampleN (length mu) (R.normal 0 1) r
    let n = [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
mu
        ys = [[Double]
mu [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
i Double -> Double -> Double
forall a. Num a => a -> a -> a
+ [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Matrix Double -> [Double]
forall a. Matrix a -> [a]
toList Matrix Double
l [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j) Double -> Double -> Double
forall a. Num a => a -> a -> a
* [Double]
zs [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
j | Int
j <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]] | Int
i <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
    pure (ys, s)

-- | Wishart log density for real degrees of freedom nu>p-1 and positive-definite p-by-p scale and observation matrices.
wishartLogDensity :: Double -> Matrix Double -> Matrix Double -> Either MathError Double
wishartLogDensity :: Double -> Matrix Double -> Matrix Double -> Either MathError Double
wishartLogDensity Double
nu Matrix Double
scale Matrix Double
x = do
    let p :: Int
p = (Int, Int) -> Int
forall a b. (a, b) -> a
fst (Matrix Double -> (Int, Int)
forall a. Matrix a -> (Int, Int)
shape Matrix Double
scale)
    Bool -> String -> Either MathError ()
require
        (Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& Matrix Double -> (Int, Int)
forall a. Matrix a -> (Int, Int)
shape Matrix Double
scale (Int, Int) -> (Int, Int) -> Bool
forall a. Eq a => a -> a -> Bool
== (Int
p, Int
p) Bool -> Bool -> Bool
&& Matrix Double -> (Int, Int)
forall a. Matrix a -> (Int, Int)
shape Matrix Double
x (Int, Int) -> (Int, Int) -> Bool
forall a. Eq a => a -> a -> Bool
== (Int
p, Int
p) Bool -> Bool -> Bool
&& Double -> Bool
finite Double
nu Bool -> Bool -> Bool
&& Double
nu Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))
        String
"Wishart degrees of freedom must exceed p-1"
    _ <- Matrix Double -> Either MathError (Matrix Double)
forall a. RealFloat a => Matrix a -> Either MathError (Matrix a)
cholesky Matrix Double
scale
    _ <- cholesky x
    ds <- determinant scale
    dx <- determinant x
    inv <- inverse scale
    prod <- multiplyChecked inv x
    gs <- mapM (\Int
j -> Double -> Either MathError Double
S.logGamma ((Double
nu Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
j) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2)) [1 .. p]
    let tr = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Matrix Double -> [Double]
forall a. Matrix a -> [a]
toList Matrix Double
prod [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j) | Int
j <- [Int
0 .. Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
        lg = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
* (Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
4 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
log Double
forall a. Floating a => a
pi Double -> Double -> Double
forall a. Num a => a -> a -> a
+ [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
gs
    checked
        ( (nu - fromIntegral p - 1) / 2 * log dx
            - tr / 2
            - nu * fromIntegral p / 2 * log 2
            - nu / 2 * log ds
            - lg
        )

-- | Reference Wishart draw as the sum of nu Gaussian outer products; requires integer nu>=p and nu\<=100000.
sampleWishart :: Int -> Matrix Double -> R.RNG -> Either MathError (Matrix Double, R.RNG)
sampleWishart :: Int
-> Matrix Double -> RNG -> Either MathError (Matrix Double, RNG)
sampleWishart Int
nu Matrix Double
scale RNG
r = do
    let p :: Int
p = (Int, Int) -> Int
forall a b. (a, b) -> a
fst (Matrix Double -> (Int, Int)
forall a. Matrix a -> (Int, Int)
shape Matrix Double
scale)
    Bool -> String -> Either MathError ()
require
        (Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& Matrix Double -> (Int, Int)
forall a. Matrix a -> (Int, Int)
shape Matrix Double
scale (Int, Int) -> (Int, Int) -> Bool
forall a. Eq a => a -> a -> Bool
== (Int
p, Int
p) Bool -> Bool -> Bool
&& Int
nu Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
p Bool -> Bool -> Bool
&& Int
nu Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
100000)
        String
"Wishart reference sampler uses integer nu>=p"
    (ys, s) <- Int
-> (RNG -> Either MathError ([Double], RNG))
-> RNG
-> Either MathError ([[Double]], RNG)
forall a.
Int
-> (RNG -> Either MathError (a, RNG))
-> RNG
-> Either MathError ([a], RNG)
R.sampleN Int
nu ([Double]
-> Matrix Double -> RNG -> Either MathError ([Double], RNG)
sampleMultivariateNormal (Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate Int
p Double
0) Matrix Double
scale) RNG
r
    mat <- fromListChecked nu p (concat ys)
    x <- multiplyChecked (transposeMatrix mat) mat
    pure (x, s)

-- | Opaque finite row-stochastic transition matrix.
newtype Markov = Markov [[Double]] deriving (Markov -> Markov -> Bool
(Markov -> Markov -> Bool)
-> (Markov -> Markov -> Bool) -> Eq Markov
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Markov -> Markov -> Bool
== :: Markov -> Markov -> Bool
$c/= :: Markov -> Markov -> Bool
/= :: Markov -> Markov -> Bool
Eq, Int -> Markov -> ShowS
[Markov] -> ShowS
Markov -> String
(Int -> Markov -> ShowS)
-> (Markov -> String) -> ([Markov] -> ShowS) -> Show Markov
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Markov -> ShowS
showsPrec :: Int -> Markov -> ShowS
$cshow :: Markov -> String
show :: Markov -> String
$cshowList :: [Markov] -> ShowS
showList :: [Markov] -> ShowS
Show)

-- | Validate and construct a row-stochastic matrix with 1..1024 states; row sums tolerate absolute error 1e-12.
markov :: [[Double]] -> Either MathError Markov
markov :: [[Double]] -> Either MathError Markov
markov [[Double]]
rows = do
    let n :: Int
n = [[Double]] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Double]]
rows
    Bool -> String -> Either MathError ()
require
        (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1024 Bool -> Bool -> Bool
&& ([Double] -> Bool) -> [[Double]] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\[Double]
r -> [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
r Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n Bool -> Bool -> Bool
&& (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Double -> Bool
prob [Double]
r Bool -> Bool -> Bool
&& Double -> Double
forall a. Num a => a -> a
abs ([Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
r Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1) Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
1e-12) [[Double]]
rows)
        String
"Markov transition matrix rows must be probabilities"
    Markov -> Either MathError Markov
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([[Double]] -> Markov
Markov [[Double]]
rows)

-- | Advance a normalized state-probability row vector by one Markov step.
transition :: Markov -> [Double] -> Either MathError [Double]
transition :: Markov -> [Double] -> Either MathError [Double]
transition (Markov [[Double]]
rows) [Double]
ps = do
    Bool -> String -> Either MathError ()
require
        ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
ps Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [[Double]] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Double]]
rows Bool -> Bool -> Bool
&& (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Double -> Bool
prob [Double]
ps Bool -> Bool -> Bool
&& Double -> Double
forall a. Num a => a -> a
abs ([Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
ps Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1) Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
1e-10)
        String
"invalid state probabilities"
    [Double] -> Either MathError [Double]
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [[Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [[Double]
ps [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
i Double -> Double -> Double
forall a. Num a => a -> a -> a
* [[Double]]
rows [[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! Int
i [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
j | Int
i <- [Int
0 .. [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
ps Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]] | Int
j <- [Int
0 .. [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
ps Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]

-- | Stationary probabilities by damped power iteration, with caller iteration budget and positive residual tolerance. Reducible chains can have multiple stationary laws.
stationary :: Int -> Double -> Markov -> Either MathError [Double]
stationary :: Int -> Double -> Markov -> Either MathError [Double]
stationary Int
budget Double
tol m :: Markov
m@(Markov [[Double]]
rows) = do
    Bool -> String -> Either MathError ()
require (Int
budget Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& Double -> Bool
positive Double
tol) String
"stationary requires positive budget and tolerance"
    -- Cesaro damping preserves stationary states and converges for periodic chains.
    let go :: t -> [Double] -> Either MathError [Double]
go t
0 [Double]
_ = MathError -> Either MathError [Double]
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"stationary distribution" Int
budget)
        go t
n [Double]
p = do
            q <- Markov -> [Double] -> Either MathError [Double]
transition Markov
m [Double]
p
            let z = (Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Double
a Double
b -> (Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
b) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2) [Double]
p [Double]
q
             in if maximum (zipWith (\Double
a Double
b -> Double -> Double
forall a. Num a => a -> a
abs (Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
b)) p q) < tol then pure q else go (n - 1) z
    Int -> [Double] -> Either MathError [Double]
forall {t}.
(Eq t, Num t) =>
t -> [Double] -> Either MathError [Double]
go Int
budget (Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate ([[Double]] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Double]]
rows) (Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([[Double]] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Double]]
rows)))

-- | Opaque finite-state categorical-emission hidden Markov model.
data HMM = HMM [Double] [[Double]] [[Double]] deriving (HMM -> HMM -> Bool
(HMM -> HMM -> Bool) -> (HMM -> HMM -> Bool) -> Eq HMM
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: HMM -> HMM -> Bool
== :: HMM -> HMM -> Bool
$c/= :: HMM -> HMM -> Bool
/= :: HMM -> HMM -> Bool
Eq, Int -> HMM -> ShowS
[HMM] -> ShowS
HMM -> String
(Int -> HMM -> ShowS)
-> (HMM -> String) -> ([HMM] -> ShowS) -> Show HMM
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> HMM -> ShowS
showsPrec :: Int -> HMM -> ShowS
$cshow :: HMM -> String
show :: HMM -> String
$cshowList :: [HMM] -> ShowS
showList :: [HMM] -> ShowS
Show)

-- | Validate initial probabilities, square transition matrix and equally sized categorical-emission rows.
hmm :: [Double] -> [[Double]] -> [[Double]] -> Either MathError HMM
hmm :: [Double] -> [[Double]] -> [[Double]] -> Either MathError HMM
hmm [Double]
initial [[Double]]
trans [[Double]]
emissions = do
    _ <- [[Double]] -> Either MathError Markov
markov [[Double]]
trans
    let n = [[Double]] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Double]]
trans; k = case [[Double]]
emissions of [] -> Int
0; [Double]
r : [[Double]]
_ -> [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
r
    require
        ( length initial == n
            && all prob initial
            && abs (sum initial - 1) < 1e-12
            && length emissions == n
            && k > 0
            && all (\[Double]
r -> [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
r Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
k Bool -> Bool -> Bool
&& (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Double -> Bool
prob [Double]
r Bool -> Bool -> Bool
&& Double -> Double
forall a. Num a => a -> a
abs ([Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
r Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1) Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
1e-12) emissions
        )
        "HMM invalid initial/emission probabilities"
    pure (HMM initial trans emissions)

-- | Check a nonempty HMM observation sequence of zero-based emission-symbol indices.
validateObservations :: HMM -> [Int] -> Either MathError ()
validateObservations :: HMM -> [Int] -> Either MathError ()
validateObservations (HMM [Double]
_ [[Double]]
_ [[Double]]
es) [Int]
obs =
    Bool -> String -> Either MathError ()
require
        (Bool -> Bool
not ([Int] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Int]
obs) Bool -> Bool -> Bool
&& (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\Int
x -> Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([[Double]]
es [[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! Int
0)) [Int]
obs)
        String
"HMM requires valid nonempty observation sequence"

{- | Scaled forward returns log likelihood, normalized alpha rows and per-time
scaling constants. This avoids the product underflow of raw probability DP.
-}
forward :: HMM -> [Int] -> Either MathError (Double, [[Double]], [Double])
forward :: HMM -> [Int] -> Either MathError (Double, [[Double]], [Double])
forward model :: HMM
model@(HMM [Double]
initial [[Double]]
trans [[Double]]
es) [Int]
obs = do
    HMM -> [Int] -> Either MathError ()
validateObservations HMM
model [Int]
obs
    let norm :: [Double] -> Either MathError ([Double], Double)
norm [Double]
xs =
            let c :: Double
c = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
xs in Bool -> String -> Either MathError ()
require (Double -> Bool
positive Double
c) String
"HMM zero observation likelihood" Either MathError ()
-> Either MathError ([Double], Double)
-> Either MathError ([Double], Double)
forall a b.
Either MathError a -> Either MathError b -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> ([Double], Double) -> Either MathError ([Double], Double)
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
c) [Double]
xs, Double
c)
    (a, c) <- [Double] -> Either MathError ([Double], Double)
norm ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) [Double]
initial [[[Double]]
es [[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! Int
i [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! ([Int]
obs [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
0) | Int
i <- [Int
0 .. [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
initial Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]])
    let step ([[Double]]
as, [Double]
cs) Int
o = do
            let prev :: [Double]
prev = [[Double]] -> [Double]
forall a. HasCallStack => [a] -> a
last [[Double]]
as
                next :: [Double]
next =
                    [ [[Double]]
es [[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! Int
j [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
o Double -> Double -> Double
forall a. Num a => a -> a -> a
* [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [[Double]
prev [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
i Double -> Double -> Double
forall a. Num a => a -> a -> a
* [[Double]]
trans [[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! Int
i [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
j | Int
i <- [Int
0 .. [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
initial Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
                    | Int
j <- [Int
0 .. [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
initial Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                    ]
            (v, s) <- [Double] -> Either MathError ([Double], Double)
norm [Double]
next
            pure (as ++ [v], cs ++ [s])
    (as, cs) <- foldM step ([a], [c]) (drop 1 obs)
    pure (sum (map log cs), as, cs)

-- | Scaled backward messages using the same per-time normalizers as forward; terminal message is the all-ones vector.
backward :: HMM -> [Int] -> Either MathError [[Double]]
backward :: HMM -> [Int] -> Either MathError [[Double]]
backward model :: HMM
model@(HMM [Double]
initial [[Double]]
trans [[Double]]
es) [Int]
obs = do
    (_, _, cs) <- HMM -> [Int] -> Either MathError (Double, [[Double]], [Double])
forward HMM
model [Int]
obs
    let n = [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
initial
        t = [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
obs
        go Int
k [[Double]]
rows
            | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = [[Double]]
rows
            | Bool
otherwise =
                let next :: [Double]
next = (case [[Double]]
rows of [Double]
a : [[Double]]
_ -> [Double]
a; [] -> Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate Int
n Double
1)
                    b :: [Double]
b =
                        [ [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [[[Double]]
trans [[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! Int
i [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
j Double -> Double -> Double
forall a. Num a => a -> a -> a
* [[Double]]
es [[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! Int
j [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! ([Int]
obs [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)) Double -> Double -> Double
forall a. Num a => a -> a -> a
* [Double]
next [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
j | Int
j <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
                            Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ ([Double]
cs [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
                        | Int
i <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                        ]
                 in Int -> [[Double]] -> [[Double]]
go (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) ([Double]
b [Double] -> [[Double]] -> [[Double]]
forall a. a -> [a] -> [a]
: [[Double]]
rows)
    pure (go (t - 2) [replicate n 1])

-- | Posterior state probabilities at each observation time from normalized alpha-beta products.
posteriorStates :: HMM -> [Int] -> Either MathError [[Double]]
posteriorStates :: HMM -> [Int] -> Either MathError [[Double]]
posteriorStates HMM
model [Int]
obs = do
    (_, as, _) <- HMM -> [Int] -> Either MathError (Double, [[Double]], [Double])
forward HMM
model [Int]
obs
    bs <- backward model obs
    pure [let ps = (Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) [Double]
a [Double]
b in map (/ sum ps) ps | (a, b) <- zip as bs]

-- | Most likely state path and its joint log probability, via max-product dynamic programming in log space.
viterbi :: HMM -> [Int] -> Either MathError ([Int], Double)
viterbi :: HMM -> [Int] -> Either MathError ([Int], Double)
viterbi model :: HMM
model@(HMM [Double]
initial [[Double]]
trans [[Double]]
es) [Int]
obs = do
    HMM -> [Int] -> Either MathError ()
validateObservations HMM
model [Int]
obs
    let start :: [(Double, [Int])]
start = [(Double -> Double
forall a. Floating a => a -> a
log Double
p Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double -> Double
forall a. Floating a => a -> a
log ([[Double]]
es [[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! Int
i [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! ([Int]
obs [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
0)), [Int
i]) | (Int
i, Double
p) <- [Int] -> [Double] -> [(Int, Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [Double]
initial]
        step :: [(Double, [Int])] -> Int -> [(Double, [Int])]
step [(Double, [Int])]
states Int
o =
            [ ((Double, [Int]) -> (Double, [Int]) -> Ordering)
-> [(Double, [Int])] -> (Double, [Int])
forall (t :: * -> *) a.
Foldable t =>
(a -> a -> Ordering) -> t a -> a
maximumBy
                (((Double, [Int]) -> Double)
-> (Double, [Int]) -> (Double, [Int]) -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing (Double, [Int]) -> Double
forall a b. (a, b) -> a
fst)
                [ (Double
l Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double -> Double
forall a. Floating a => a -> a
log ([[Double]]
trans [[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! Int
i [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
j) Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double -> Double
forall a. Floating a => a -> a
log ([[Double]]
es [[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! Int
j [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
o), [Int]
previousPath [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int
j])
                | (Int
i, (Double
l, [Int]
previousPath)) <- [Int] -> [(Double, [Int])] -> [(Int, (Double, [Int]))]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [(Double, [Int])]
states
                ]
            | Int
j <- [Int
0 .. [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
initial Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
            ]
        (Double
score, [Int]
path) = ((Double, [Int]) -> (Double, [Int]) -> Ordering)
-> [(Double, [Int])] -> (Double, [Int])
forall (t :: * -> *) a.
Foldable t =>
(a -> a -> Ordering) -> t a -> a
maximumBy (((Double, [Int]) -> Double)
-> (Double, [Int]) -> (Double, [Int]) -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing (Double, [Int]) -> Double
forall a b. (a, b) -> a
fst) (([(Double, [Int])] -> Int -> [(Double, [Int])])
-> [(Double, [Int])] -> [Int] -> [(Double, [Int])]
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl' [(Double, [Int])] -> Int -> [(Double, [Int])]
step [(Double, [Int])]
start (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
drop Int
1 [Int]
obs))
    Bool -> String -> Either MathError ()
require (Double -> Bool
finite Double
score) String
"HMM no possible state path"
    ([Int], Double) -> Either MathError ([Int], Double)
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Int]
path, Double
score)

{- | Baum--Welch EM for one observation sequence, pseudocount>=0. Transition
rows with zero expected occupancy retain their old row. Fixed iteration budget.
-}
baumWelch :: Int -> Double -> HMM -> [Int] -> Either MathError HMM
baumWelch :: Int -> Double -> HMM -> [Int] -> Either MathError HMM
baumWelch Int
count Double
pseudo HMM
initialModel [Int]
obs = do
    Bool -> String -> Either MathError ()
require
        (Int
count Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
count Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
10000 Bool -> Bool -> Bool
&& Double -> Bool
finite Double
pseudo Bool -> Bool -> Bool
&& Double
pseudo Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
0)
        String
"Baum-Welch invalid iteration count/pseudocount"
    (HMM -> Int -> Either MathError HMM)
-> HMM -> [Int] -> Either MathError HMM
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (\HMM
model Int
_ -> HMM -> Either MathError HMM
step HMM
model) HMM
initialModel [Int
1 .. Int
count]
  where
    step :: HMM -> Either MathError HMM
step model :: HMM
model@(HMM [Double]
_ [[Double]]
trans [[Double]]
es) = do
        (_, as, cs) <- HMM -> [Int] -> Either MathError (Double, [[Double]], [Double])
forward HMM
model [Int]
obs
        bs <- backward model obs
        gs <- posteriorStates model obs
        let n = [[Double]] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Double]]
trans
            k = [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([[Double]]
es [[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! Int
0)
            t = [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
obs
            norm [b]
old [b]
xs = if [b] -> b
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [b]
xs b -> b -> Bool
forall a. Eq a => a -> a -> Bool
== b
0 then [b]
old else (b -> b) -> [b] -> [b]
forall a b. (a -> b) -> [a] -> [b]
map (b -> b -> b
forall a. Fractional a => a -> a -> a
/ [b] -> b
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [b]
xs) [b]
xs
            xis Int
i Int
j =
                [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum
                    [ ([[Double]]
as [[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! Int
u [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
i)
                        Double -> Double -> Double
forall a. Num a => a -> a -> a
* [[Double]]
trans [[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! Int
i [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
j
                        Double -> Double -> Double
forall a. Num a => a -> a -> a
* [[Double]]
es [[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! Int
j [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! ([Int]
obs [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! (Int
u Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
                        Double -> Double -> Double
forall a. Num a => a -> a -> a
* ([[Double]]
bs [[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! (Int
u Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
j)
                        Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ ([Double]
cs [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! (Int
u Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
                    | Int
u <- [Int
0 .. Int
t Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2]
                    ]
            tr = [[Double] -> [Double] -> [Double]
forall {b}. (Eq b, Fractional b) => [b] -> [b] -> [b]
norm ([[Double]]
trans [[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! Int
i) [Double
pseudo Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Int -> Int -> Double
xis Int
i Int
j | Int
j <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]] | Int
i <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
            em =
                [ [Double] -> [Double] -> [Double]
forall {b}. (Eq b, Fractional b) => [b] -> [b] -> [b]
norm
                    ([[Double]]
es [[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! Int
i)
                    [Double
pseudo Double -> Double -> Double
forall a. Num a => a -> a -> a
+ [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [[[Double]]
gs [[Double]] -> Int -> [Double]
forall a. HasCallStack => [a] -> Int -> a
!! Int
u [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
i | Int
u <- [Int
0 .. Int
t Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1], [Int]
obs [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
u Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
symbol] | Int
symbol <- [Int
0 .. Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
                | Int
i <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                ]
        hmm (gs !! 0) tr em