{-# LANGUAGE
    MultiParamTypeClasses,
    FlexibleInstances, FlexibleContexts,
    UndecidableInstances
  #-}

{-# OPTIONS_GHC -fno-warn-simplifiable-class-constraints #-}

module Data.Random.Distribution.Beta where

import Data.Random.RVar
import Data.Random.Distribution
import Data.Random.Distribution.Gamma
import Data.Random.Distribution.Uniform

import Numeric.SpecFunctions

{-# SPECIALIZE fractionalBeta :: Float  -> Float  -> RVarT m Float #-}
{-# SPECIALIZE fractionalBeta :: Double -> Double -> RVarT m Double #-}
fractionalBeta :: (Fractional a, Eq a, Distribution Gamma a, Distribution StdUniform a) => a -> a -> RVarT m a
fractionalBeta :: a -> a -> RVarT m a
fractionalBeta a
1 a
1 = RVarT m a
forall a (m :: * -> *). Distribution StdUniform a => RVarT m a
stdUniformT
fractionalBeta a
a a
b = do
    a
x <- a -> a -> RVarT m a
forall a (m :: * -> *). Distribution Gamma a => a -> a -> RVarT m a
gammaT a
a a
1
    a
y <- a -> a -> RVarT m a
forall a (m :: * -> *). Distribution Gamma a => a -> a -> RVarT m a
gammaT a
b a
1
    a -> RVarT m a
forall (m :: * -> *) a. Monad m => a -> m a
return (a
x a -> a -> a
forall a. Fractional a => a -> a -> a
/ (a
x a -> a -> a
forall a. Num a => a -> a -> a
+ a
y))

{-# SPECIALIZE beta :: Float  -> Float  -> RVar Float #-}
{-# SPECIALIZE beta :: Double -> Double -> RVar Double #-}
beta :: Distribution Beta a => a -> a -> RVar a
beta :: a -> a -> RVar a
beta a
a a
b = Beta a -> RVar a
forall (d :: * -> *) t. Distribution d t => d t -> RVar t
rvar (a -> a -> Beta a
forall a. a -> a -> Beta a
Beta a
a a
b)

{-# SPECIALIZE betaT :: Float  -> Float  -> RVarT m Float #-}
{-# SPECIALIZE betaT :: Double -> Double -> RVarT m Double #-}
betaT :: Distribution Beta a => a -> a -> RVarT m a
betaT :: a -> a -> RVarT m a
betaT a
a a
b = Beta a -> RVarT m a
forall (d :: * -> *) t (n :: * -> *).
Distribution d t =>
d t -> RVarT n t
rvarT (a -> a -> Beta a
forall a. a -> a -> Beta a
Beta a
a a
b)

data Beta a = Beta a a

-- FIXME: I am far from convinced that NaNs are a good idea.
logBetaPdf :: Double -> Double -> Double -> Double
logBetaPdf :: Double -> Double -> Double -> Double
logBetaPdf Double
a Double
b Double
x
   | Double
a Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 Bool -> Bool -> Bool
|| Double
b Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 = Double
nan
   | Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 = Double -> Double
forall a. Floating a => a -> a
log Double
0
   | Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
1 = Double -> Double
forall a. Floating a => a -> a
log Double
0
   | Bool
otherwise = (Double
aDouble -> 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 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ (Double
bDouble -> 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
1Double -> Double -> Double
forall a. Num a => a -> a -> a
-Double
x) Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double -> Double -> Double
logBeta Double
a Double
b
  where
    nan :: Double
nan = Double
0.0 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0.0

instance PDF Beta Double
  where
    pdf :: Beta Double -> Double -> Double
pdf (Beta Double
a Double
b) = Double -> Double
forall a. Floating a => a -> a
exp (Double -> Double) -> (Double -> Double) -> Double -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Double -> Double -> Double
logBetaPdf Double
a Double
b

instance PDF Beta Float
  where
    pdf :: Beta Float -> Float -> Double
pdf (Beta Float
a Float
b) = Double -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac (Double -> Double) -> (Float -> Double) -> Float -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Double
forall a. Floating a => a -> a
exp (Double -> Double) -> (Float -> Double) -> Float -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Double -> Double -> Double
logBetaPdf (Float -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac Float
a) (Float -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac Float
b) (Double -> Double) -> (Float -> Double) -> Float -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Float -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac

instance Distribution Beta Float
  where
    rvarT :: Beta Float -> RVarT n Float
rvarT (Beta Float
a Float
b) = Float -> Float -> RVarT n Float
forall a (m :: * -> *).
(Fractional a, Eq a, Distribution Gamma a,
 Distribution StdUniform a) =>
a -> a -> RVarT m a
fractionalBeta Float
a Float
b

instance Distribution Beta Double
  where
    rvarT :: Beta Double -> RVarT n Double
rvarT (Beta Double
a Double
b) = Double -> Double -> RVarT n Double
forall a (m :: * -> *).
(Fractional a, Eq a, Distribution Gamma a,
 Distribution StdUniform a) =>
a -> a -> RVarT m a
fractionalBeta Double
a Double
b