-- | Deterministic and probabilistic values

module Numeric.Probability.Distribution where

import Numeric.Probability.Show (showR)
import qualified Numeric.Probability.Shape as Shape

import qualified Control.Functor.HT as FuncHT
import Control.Applicative (Applicative(..))
import Control.Monad (liftM, liftM2, join, )

import qualified Data.Foldable as Fold
import qualified Data.List.HT as ListHT
import qualified Data.Map  as Map
import qualified Data.List as List
import qualified Data.Maybe as Maybe
import Data.Tuple.HT (mapFst, )
import Data.Ord.HT (comparing, )
import Data.Eq.HT (equating, )

import Prelude hiding (map, filter)


-- * Events
type Event a = a -> Bool

oneOf :: Eq a => [a] -> Event a
oneOf :: forall a. Eq a => [a] -> Event a
oneOf = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
elem

just :: Eq a => a -> Event a
just :: forall a. Eq a => a -> Event a
just = forall a. Eq a => a -> Event a
(==)



-- * Distributions

{- |
Probability disribution

The underlying data structure is a list.
Unfortunately we cannot use a more efficient data structure
because the key type must be of class 'Ord',
but the 'Monad' class does not allow constraints for result types.
The Monad instance is particularly useful
because many generic monad functions make sense here,
monad transformers can be used
and the monadic design allows to simulate probabilistic games in an elegant manner.

We have the same problem like making "Data.Set" an instance of 'Monad',
see <http://www.randomhacks.net/articles/2007/03/15/data-set-monad-haskell-macros>

If you need efficiency, you should remove redundant elements by 'norm'.
'norm' converts to 'Data.Map' and back internally
and you can hope that the compiler fuses the lists with the intermediate Map structure.

The defined monad is equivalent to
@WriterT (Product prob) [] a@.
See <http://www.randomhacks.net/articles/2007/02/21/refactoring-probability-distributions>.
-}
newtype T prob a = Cons {forall prob a. T prob a -> [(a, prob)]
decons :: [(a,prob)]}

certainly :: Num prob => a -> T prob a
certainly :: forall prob a. Num prob => a -> T prob a
certainly a
x = forall prob a. [(a, prob)] -> T prob a
Cons [(a
x,prob
1)]

instance Num prob => Monad (T prob) where
  return :: forall a. a -> T prob a
return   = forall prob a. Num prob => a -> T prob a
certainly
  T prob a
d >>= :: forall a b. T prob a -> (a -> T prob b) -> T prob b
>>= a -> T prob b
f  = forall prob a. [(a, prob)] -> T prob a
Cons [(b
y,prob
qforall a. Num a => a -> a -> a
*prob
p) | (a
x,prob
p) <- forall prob a. T prob a -> [(a, prob)]
decons T prob a
d, (b
y,prob
q) <- forall prob a. T prob a -> [(a, prob)]
decons (a -> T prob b
f a
x)]

instance Num prob => Applicative (T prob) where
  pure :: forall a. a -> T prob a
pure     = forall prob a. Num prob => a -> T prob a
certainly
  T prob (a -> b)
fm <*> :: forall a b. T prob (a -> b) -> T prob a -> T prob b
<*> T prob a
m = forall prob a. [(a, prob)] -> T prob a
Cons [(a -> b
f a
x,prob
qforall a. Num a => a -> a -> a
*prob
p) | (a -> b
f,prob
p) <- forall prob a. T prob a -> [(a, prob)]
decons T prob (a -> b)
fm, (a
x,prob
q) <- forall prob a. T prob a -> [(a, prob)]
decons T prob a
m]

{-
Dist cannot be an instance of MonadPlus,
because there is no mzero
(it would be an empty list of events, but their probabilities do not sum up to 1)
and thus it breaks the normalization for the >>= combinator.
See for instance the Boys example:

   do f <- family
      guard (existsBoy f)
      return f

mplus is not associative because we have to normalize the sum of probabilities to 1.

instance MonadPlus Dist where
  mzero      = Cons []
  mplus d d' =
     if isZero d || isZero d'
       then mzero
       else unfoldD $ choose 0.5 d d'

isZero :: Dist a -> Bool
isZero (Cons d) = null d
-}


instance Functor (T prob) where
  fmap :: forall a b. (a -> b) -> T prob a -> T prob b
fmap a -> b
f (Cons [(a, prob)]
d) = forall prob a. [(a, prob)] -> T prob a
Cons [(a -> b
f a
x,prob
p) | (a
x,prob
p) <- [(a, prob)]
d]



errorMargin :: RealFloat prob => prob
errorMargin :: forall prob. RealFloat prob => prob
errorMargin =
   let eps :: prob
eps = prob
10 forall a. Fractional a => a -> a -> a
/ forall a. Num a => Integer -> a
fromInteger (forall a. RealFloat a => a -> Integer
floatRadix prob
eps) forall a b. (Num a, Integral b) => a -> b -> a
^ forall a. RealFloat a => a -> Int
floatDigits prob
eps
   in  prob
eps

{- |
Check whether two distributions are equal when neglecting rounding errors.
We do not want to put this into an 'Eq' instance,
since it is not exact equivalence
and it seems to be too easy to mix it up with @liftM2 (==) x y@.
-}
approx :: (RealFloat prob, Ord a) =>
   T prob a -> T prob a ->
   Bool
approx :: forall prob a.
(RealFloat prob, Ord a) =>
T prob a -> T prob a -> Bool
approx (Cons [(a, prob)]
xs) (Cons [(a, prob)]
ys) =
   let ([a]
xse, [prob]
xsp) = forall a b. [(a, b)] -> ([a], [b])
unzip (forall prob a. (Num prob, Ord a) => [(a, prob)] -> [(a, prob)]
norm' [(a, prob)]
xs)
       ([a]
yse, [prob]
ysp) = forall a b. [(a, b)] -> ([a], [b])
unzip (forall prob a. (Num prob, Ord a) => [(a, prob)] -> [(a, prob)]
norm' [(a, prob)]
ys)
   in  [a]
xse forall a. Eq a => a -> Event a
== [a]
yse Bool -> Bool -> Bool
&&
       forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\prob
p -> forall a. Num a => a -> a
abs prob
p forall a. Ord a => a -> a -> Bool
< forall prob. RealFloat prob => prob
errorMargin) (forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (-) [prob]
xsp [prob]
ysp)


-- ** Auxiliary functions for constructing and working with distributions
lift :: (Num prob) => ([(a,prob)] -> [(b,prob)]) -> T prob a -> T prob b
lift :: forall prob a b.
Num prob =>
([(a, prob)] -> [(b, prob)]) -> T prob a -> T prob b
lift [(a, prob)] -> [(b, prob)]
f  = forall prob a. [(a, prob)] -> T prob a
Cons forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(a, prob)] -> [(b, prob)]
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall prob a. T prob a -> [(a, prob)]
decons

size :: T prob a -> Int
size :: forall prob a. T prob a -> Int
size = forall (t :: * -> *) a. Foldable t => t a -> Int
length forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall prob a. T prob a -> [(a, prob)]
decons

check :: (RealFloat prob, Show prob) => T prob a -> T prob a
check :: forall prob a. (RealFloat prob, Show prob) => T prob a -> T prob a
check (Cons [(a, prob)]
d) =
   if forall a. Num a => a -> a
abs (prob
1forall a. Num a => a -> a -> a
-forall prob a. Num prob => [(a, prob)] -> prob
sumP [(a, prob)]
d) forall a. Ord a => a -> a -> Bool
< forall prob. RealFloat prob => prob
errorMargin
     then forall prob a. [(a, prob)] -> T prob a
Cons [(a, prob)]
d
     else forall a. HasCallStack => [Char] -> a
error ([Char]
"Illegal distribution: total probability = "forall a. [a] -> [a] -> [a]
++forall a. Show a => a -> [Char]
show (forall prob a. Num prob => [(a, prob)] -> prob
sumP [(a, prob)]
d))

-- | can fail because of rounding errors, better use 'fromFreqs'
cons :: (RealFloat prob, Show prob) => [(a,prob)] -> T prob a
cons :: forall prob a.
(RealFloat prob, Show prob) =>
[(a, prob)] -> T prob a
cons = forall prob a. (RealFloat prob, Show prob) => T prob a -> T prob a
check forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall prob a. [(a, prob)] -> T prob a
Cons

sumP :: Num prob => [(a,prob)] -> prob
sumP :: forall prob a. Num prob => [(a, prob)] -> prob
sumP = forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
List.map forall a b. (a, b) -> b
snd

sortP :: Ord prob => [(a,prob)] -> [(a,prob)]
sortP :: forall prob a. Ord prob => [(a, prob)] -> [(a, prob)]
sortP = forall a. (a -> a -> Ordering) -> [a] -> [a]
List.sortBy (forall b a. Ord b => (a -> b) -> a -> a -> Ordering
comparing forall a b. (a, b) -> b
snd)

sortPDesc :: Ord prob => [(a,prob)] -> [(a,prob)]
sortPDesc :: forall prob a. Ord prob => [(a, prob)] -> [(a, prob)]
sortPDesc = forall a. (a -> a -> Ordering) -> [a] -> [a]
List.sortBy (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a b. (a -> b) -> a -> b
$ forall b a. Ord b => (a -> b) -> a -> a -> Ordering
comparing forall a b. (a, b) -> b
snd)

sortElem :: Ord a => [(a,prob)] -> [(a,prob)]
sortElem :: forall a prob. Ord a => [(a, prob)] -> [(a, prob)]
sortElem = forall a. (a -> a -> Ordering) -> [a] -> [a]
List.sortBy (forall b a. Ord b => (a -> b) -> a -> a -> Ordering
comparing forall a b. (a, b) -> a
fst)


-- ** Normalization = grouping
norm :: (Num prob, Ord a) => T prob a -> T prob a
norm :: forall prob a. (Num prob, Ord a) => T prob a -> T prob a
norm = forall prob a b.
Num prob =>
([(a, prob)] -> [(b, prob)]) -> T prob a -> T prob b
lift forall prob a. (Num prob, Ord a) => [(a, prob)] -> [(a, prob)]
norm'

norm' :: (Num prob, Ord a) => [(a,prob)] -> [(a,prob)]
norm' :: forall prob a. (Num prob, Ord a) => [(a, prob)] -> [(a, prob)]
norm' =
   forall k a. Map k a -> [(k, a)]
Map.toAscList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith forall a. Num a => a -> a -> a
(+)

norm'' :: (Num prob, Ord a) => [(a,prob)] -> [(a,prob)]
norm'' :: forall prob a. (Num prob, Ord a) => [(a, prob)] -> [(a, prob)]
norm'' =
   forall a b. (a -> b) -> [a] -> [b]
List.map (\[(a, prob)]
xs ->
      case [(a, prob)]
xs of
         ((a
x,prob
_):[(a, prob)]
_) -> (a
x, forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (forall a b. (a -> b) -> [a] -> [b]
List.map forall a b. (a, b) -> b
snd [(a, prob)]
xs))
         [(a, prob)]
_ -> forall a. HasCallStack => [Char] -> a
error [Char]
"Probability.Distribution.norm'': every sub-list in groupBy must be non-empty") forall b c a. (b -> c) -> (a -> b) -> a -> c
.
   forall a. (a -> a -> Bool) -> [a] -> [[a]]
ListHT.groupBy (forall b a. Eq b => (a -> b) -> a -> a -> Bool
equating forall a b. (a, b) -> a
fst) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a prob. Ord a => [(a, prob)] -> [(a, prob)]
sortElem


-- | pretty printing
pretty :: (Ord a, Show a, Num prob, Ord prob) =>
   (prob -> String) -> T prob a -> String
pretty :: forall a prob.
(Ord a, Show a, Num prob, Ord prob) =>
(prob -> [Char]) -> T prob a -> [Char]
pretty prob -> [Char]
_ (Cons []) = [Char]
"Impossible"
pretty prob -> [Char]
showProb (Cons [(a, prob)]
xs) =
   let w :: Int
w = forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum (forall a b. (a -> b) -> [a] -> [b]
List.map (forall (t :: * -> *) a. Foldable t => t a -> Int
lengthforall b c a. (b -> c) -> (a -> b) -> a -> c
.forall a. Show a => a -> [Char]
showforall b c a. (b -> c) -> (a -> b) -> a -> c
.forall a b. (a, b) -> a
fst) [(a, prob)]
xs)
   in  forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap
          (\(a
x,prob
p) -> forall a. Show a => Int -> a -> [Char]
showR Int
w a
xforall a. [a] -> [a] -> [a]
++Char
' 'forall a. a -> [a] -> [a]
: prob -> [Char]
showProb prob
pforall a. [a] -> [a] -> [a]
++[Char]
"\n")
          (forall prob a. Ord prob => [(a, prob)] -> [(a, prob)]
sortPDesc (forall prob a. (Num prob, Ord a) => [(a, prob)] -> [(a, prob)]
norm' [(a, prob)]
xs))

infix 0 //%

(//%) :: (Ord a, Show a) => T Rational a -> () -> IO ()
//% :: forall a. (Ord a, Show a) => T Rational a -> () -> IO ()
(//%) T Rational a
x () = [Char] -> IO ()
putStr (forall a prob.
(Ord a, Show a, Num prob, Ord prob) =>
(prob -> [Char]) -> T prob a -> [Char]
pretty forall a. Show a => a -> [Char]
show T Rational a
x)

instance (Num prob, Ord prob, Show prob, Ord a, Show a) =>
      Show (T prob a) where
   showsPrec :: Int -> T prob a -> ShowS
showsPrec Int
p (Cons [(a, prob)]
xs) =
      Bool -> ShowS -> ShowS
showParen (Int
pforall a. Ord a => a -> a -> Bool
>Int
10)
         ([Char] -> ShowS
showString [Char]
"fromFreqs " forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => Int -> a -> ShowS
showsPrec Int
11 (forall prob a. Ord prob => [(a, prob)] -> [(a, prob)]
sortPDesc (forall prob a. (Num prob, Ord a) => [(a, prob)] -> [(a, prob)]
norm' [(a, prob)]
xs)))


{- |
We would like to have an equality test of type

> (==) :: T prob a -> T prob a -> T prob Bool

that is consistent with the 'Num' instance.
We would certainly define

> x==y = norm (liftM2 (==) x y)   .

However the 'Eq' class enforces the type

> T prob a -> T prob a -> Bool    .

We could implement this as check for equal distributions.
This would be inconsistent with the 'Num' instance
because it compares entire distributions,
not only individual outcomes.
Thus we provide this function as 'equal'.

I would prefer to omit the 'Eq' instance completely,
but unfortunately the 'Num' instance requires 'Eq' as superclass.
-}
instance Eq (T prob a) where
   == :: T prob a -> T prob a -> Bool
(==) = forall a. HasCallStack => [Char] -> a
error [Char]
"Probability.Distribution.== cannot be implemented sensibly."

{-
instance (Num prob, Ord a) => Eq (T prob a) where
   (==) = equal
-}

equal :: (Num prob, Eq prob, Ord a) => T prob a -> T prob a -> Bool
equal :: forall prob a.
(Num prob, Eq prob, Ord a) =>
T prob a -> T prob a -> Bool
equal = forall b a. Eq b => (a -> b) -> a -> a -> Bool
equating (forall prob a. T prob a -> [(a, prob)]
decons forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall prob a. (Num prob, Ord a) => T prob a -> T prob a
norm)

{-
The Num operations consider their operands as independent distributions
(like all operations on distributions do).
All functions normalize their results if normalization is lost by the plain operation.
This is essential for performance.

Thus @sum $ replicate 10 d@ is significantly faster
than @fmap sum $ replicateM 10 d@
-}
instance (Num prob, Ord prob, Ord a, Num a) => Num (T prob a) where
   fromInteger :: Integer -> T prob a
fromInteger = forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Num a => Integer -> a
fromInteger
   T prob a
x + :: T prob a -> T prob a -> T prob a
+ T prob a
y = forall prob a. (Num prob, Ord a) => T prob a -> T prob a
norm (forall (m :: * -> *) a1 a2 r.
Monad m =>
(a1 -> a2 -> r) -> m a1 -> m a2 -> m r
liftM2 forall a. Num a => a -> a -> a
(+) T prob a
x T prob a
y)
   T prob a
x - :: T prob a -> T prob a -> T prob a
- T prob a
y = forall prob a. (Num prob, Ord a) => T prob a -> T prob a
norm (forall (m :: * -> *) a1 a2 r.
Monad m =>
(a1 -> a2 -> r) -> m a1 -> m a2 -> m r
liftM2 (-) T prob a
x T prob a
y)
   T prob a
x * :: T prob a -> T prob a -> T prob a
* T prob a
y = forall prob a. (Num prob, Ord a) => T prob a -> T prob a
norm (forall (m :: * -> *) a1 a2 r.
Monad m =>
(a1 -> a2 -> r) -> m a1 -> m a2 -> m r
liftM2 forall a. Num a => a -> a -> a
(*) T prob a
x T prob a
y)
   abs :: T prob a -> T prob a
abs T prob a
x = forall prob a. (Num prob, Ord a) => T prob a -> T prob a
norm (forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM forall a. Num a => a -> a
abs T prob a
x)
   signum :: T prob a -> T prob a
signum T prob a
x = forall prob a. (Num prob, Ord a) => T prob a -> T prob a
norm (forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM forall a. Num a => a -> a
signum T prob a
x)
   negate :: T prob a -> T prob a
negate T prob a
x = forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM forall a. Num a => a -> a
negate T prob a
x

instance (Num prob, Ord prob, Ord a, Fractional a) =>
      Fractional (T prob a) where
   fromRational :: Rational -> T prob a
fromRational = forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Fractional a => Rational -> a
fromRational
   recip :: T prob a -> T prob a
recip T prob a
x = forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM forall a. Fractional a => a -> a
recip T prob a
x
   T prob a
x / :: T prob a -> T prob a -> T prob a
/ T prob a
y = forall prob a. (Num prob, Ord a) => T prob a -> T prob a
norm (forall (m :: * -> *) a1 a2 r.
Monad m =>
(a1 -> a2 -> r) -> m a1 -> m a2 -> m r
liftM2 forall a. Fractional a => a -> a -> a
(/) T prob a
x T prob a
y)



-- * Spread: functions to convert a list of values into a distribution

-- | distribution generators
type Spread prob a = [a] -> T prob a

{- not a valid distribution
impossible :: T prob a
impossible = mzero
-}

choose :: Num prob => prob -> a -> a -> T prob a
choose :: forall prob a. Num prob => prob -> a -> a -> T prob a
choose prob
p a
x a
y = forall prob a. [(a, prob)] -> T prob a
Cons forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip [a
x,a
y] [prob
p,prob
1forall a. Num a => a -> a -> a
-prob
p]

enum :: Fractional prob => [Int] -> Spread prob a
enum :: forall prob a. Fractional prob => [Int] -> Spread prob a
enum  =  forall prob a. Fractional prob => [prob] -> Spread prob a
relative forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
List.map forall a b. (Integral a, Num b) => a -> b
fromIntegral

{- |
Give a list of frequencies, they do not need to sum up to 1.
-}
relative :: Fractional prob => [prob] -> Spread prob a
relative :: forall prob a. Fractional prob => [prob] -> Spread prob a
relative [prob]
ns = forall prob a. Fractional prob => [(a, prob)] -> T prob a
fromFreqs forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a b. [a] -> [b] -> [(a, b)]
zip [prob]
ns

shape :: Fractional prob =>
   (prob -> prob) -> Spread prob a
shape :: forall prob a. Fractional prob => (prob -> prob) -> Spread prob a
shape prob -> prob
_ [] = forall a. HasCallStack => [Char] -> a
error [Char]
"Probability.shape: empty list"
shape prob -> prob
f [a]
xs =
   let incr :: prob
incr = prob
1 forall a. Fractional a => a -> a -> a
/ forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs forall a. Num a => a -> a -> a
- Int
1)
       ps :: [prob]
ps = forall a b. (a -> b) -> [a] -> [b]
List.map prob -> prob
f (forall a. (a -> a) -> a -> [a]
iterate (forall a. Num a => a -> a -> a
+prob
incr) prob
0)
   in  forall prob a. Fractional prob => [(a, prob)] -> T prob a
fromFreqs (forall a b. [a] -> [b] -> [(a, b)]
zip [a]
xs [prob]
ps)

linear :: Fractional prob => Spread prob a
linear :: forall prob a. Fractional prob => Spread prob a
linear = forall prob a. Fractional prob => (prob -> prob) -> Spread prob a
shape forall a. Fractional a => a -> a
Shape.linear

uniform :: Fractional prob => Spread prob a
uniform :: forall prob a. Fractional prob => Spread prob a
uniform = forall prob a. Fractional prob => (prob -> prob) -> Spread prob a
shape forall a. Fractional a => a -> a
Shape.uniform

negExp :: Floating prob => Spread prob a
negExp :: forall prob a. Floating prob => Spread prob a
negExp = forall prob a. Fractional prob => (prob -> prob) -> Spread prob a
shape forall prob. Floating prob => T prob
Shape.negExp

normal :: Floating prob => Spread prob a
normal :: forall prob a. Floating prob => Spread prob a
normal = forall prob a. Fractional prob => (prob -> prob) -> Spread prob a
shape forall prob. Floating prob => T prob
Shape.normal



-- | extracting and mapping the domain of a distribution
extract :: T prob a -> [a]
extract :: forall prob a. T prob a -> [a]
extract = forall a b. (a -> b) -> [a] -> [b]
List.map forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall prob a. T prob a -> [(a, prob)]
decons

-- | 'fmap' with normalization
map :: (Num prob, Ord b) =>
   (a -> b) -> T prob a -> T prob b
map :: forall prob b a.
(Num prob, Ord b) =>
(a -> b) -> T prob a -> T prob b
map a -> b
f = forall prob a. (Num prob, Ord a) => T prob a -> T prob a
norm forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f


{- |
unfold a distribution of distributions into one distribution,
this is 'Control.Monad.join' with normalization.
-}
unfold :: (Num prob, Ord a) =>
   T prob (T prob a) -> T prob a
unfold :: forall prob a. (Num prob, Ord a) => T prob (T prob a) -> T prob a
unfold = forall prob a. (Num prob, Ord a) => T prob a -> T prob a
norm forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. Monad m => m (m a) -> m a
join


-- | conditional distribution
cond :: (Num prob) =>
   T prob Bool ->
   T prob a {-^ True -} ->
   T prob a {-^ False -} ->
   T prob a
cond :: forall prob a.
Num prob =>
T prob Bool -> T prob a -> T prob a -> T prob a
cond T prob Bool
b T prob a
d T prob a
d'  =  T prob Bool
b forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Bool
c -> if Bool
c then T prob a
d else T prob a
d'

truth :: (Num prob) => T prob Bool -> prob
truth :: forall prob. Num prob => T prob Bool -> prob
truth T prob Bool
dist =
   case forall prob a. T prob a -> [(a, prob)]
decons (forall prob a. (Num prob, Ord a) => T prob a -> T prob a
norm T prob Bool
dist) of
      (Bool
b,prob
p):[(Bool, prob)]
_ -> if Bool
b then prob
p else prob
1forall a. Num a => a -> a -> a
-prob
p
      [] -> forall a. HasCallStack => [Char] -> a
error [Char]
"Probability.truth: corrupt boolean random variable"


infixl 1 >>=?
infixr 1 ?=<<

-- | conditional probability, identical to 'Dist.filter'
(?=<<) :: (Fractional prob) =>
   (a -> Bool) -> T prob a -> T prob a
?=<< :: forall prob a.
Fractional prob =>
(a -> Bool) -> T prob a -> T prob a
(?=<<) = forall prob a.
Fractional prob =>
(a -> Bool) -> T prob a -> T prob a
filter

{- |
'Dist.filter' in infix form.
Can be considered an additional monadic combinator,
which can be used where you would want 'Control.Monad.guard' otherwise.
-}
(>>=?) :: (Fractional prob) =>
   T prob a -> (a -> Bool) -> T prob a
>>=? :: forall prob a.
Fractional prob =>
T prob a -> (a -> Bool) -> T prob a
(>>=?) = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall prob a.
Fractional prob =>
(a -> Bool) -> T prob a -> T prob a
filter


-- | filtering distributions
data Select a = Case a | Other
                deriving (Select a -> Select a -> Bool
forall a. Eq a => Select a -> Select a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Select a -> Select a -> Bool
$c/= :: forall a. Eq a => Select a -> Select a -> Bool
== :: Select a -> Select a -> Bool
$c== :: forall a. Eq a => Select a -> Select a -> Bool
Eq,Select a -> Select a -> Bool
Select a -> Select a -> Ordering
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall {a}. Ord a => Eq (Select a)
forall a. Ord a => Select a -> Select a -> Bool
forall a. Ord a => Select a -> Select a -> Ordering
forall a. Ord a => Select a -> Select a -> Select a
min :: Select a -> Select a -> Select a
$cmin :: forall a. Ord a => Select a -> Select a -> Select a
max :: Select a -> Select a -> Select a
$cmax :: forall a. Ord a => Select a -> Select a -> Select a
>= :: Select a -> Select a -> Bool
$c>= :: forall a. Ord a => Select a -> Select a -> Bool
> :: Select a -> Select a -> Bool
$c> :: forall a. Ord a => Select a -> Select a -> Bool
<= :: Select a -> Select a -> Bool
$c<= :: forall a. Ord a => Select a -> Select a -> Bool
< :: Select a -> Select a -> Bool
$c< :: forall a. Ord a => Select a -> Select a -> Bool
compare :: Select a -> Select a -> Ordering
$ccompare :: forall a. Ord a => Select a -> Select a -> Ordering
Ord,Int -> Select a -> ShowS
forall a. Show a => Int -> Select a -> ShowS
forall a. Show a => [Select a] -> ShowS
forall a. Show a => Select a -> [Char]
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
showList :: [Select a] -> ShowS
$cshowList :: forall a. Show a => [Select a] -> ShowS
show :: Select a -> [Char]
$cshow :: forall a. Show a => Select a -> [Char]
showsPrec :: Int -> Select a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> Select a -> ShowS
Show)

above, below :: (Num prob, Ord prob, Ord a) =>
   prob -> T prob a -> T prob (Select a)
above :: forall prob a.
(Num prob, Ord prob, Ord a) =>
prob -> T prob a -> T prob (Select a)
above prob
p = forall prob a.
(Num prob, Ord prob, Ord a) =>
(prob -> Bool) -> T prob a -> T prob (Select a)
select (forall a. Ord a => a -> a -> Bool
>=prob
p)
below :: forall prob a.
(Num prob, Ord prob, Ord a) =>
prob -> T prob a -> T prob (Select a)
below prob
p = forall prob a.
(Num prob, Ord prob, Ord a) =>
(prob -> Bool) -> T prob a -> T prob (Select a)
select (forall a. Ord a => a -> a -> Bool
<=prob
p)

select :: (Num prob, Ord prob, Ord a) =>
   (prob -> Bool) -> T prob a -> T prob (Select a)
select :: forall prob a.
(Num prob, Ord prob, Ord a) =>
(prob -> Bool) -> T prob a -> T prob (Select a)
select prob -> Bool
condp = forall prob a b.
Num prob =>
([(a, prob)] -> [(b, prob)]) -> T prob a -> T prob b
lift forall a b. (a -> b) -> a -> b
$ \[(a, prob)]
d ->
   let (Map a prob
d1,Map a prob
d2) = forall a k. (a -> Bool) -> Map k a -> (Map k a, Map k a)
Map.partition prob -> Bool
condp forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith forall a. Num a => a -> a -> a
(+) [(a, prob)]
d
   in  (forall a. Select a
Other, forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
Fold.sum Map a prob
d2) forall a. a -> [a] -> [a]
:
          forall a b. (a -> b) -> [a] -> [b]
List.map (forall a c b. (a -> c) -> (a, b) -> (c, b)
mapFst forall a. a -> Select a
Case) (forall k a. Map k a -> [(k, a)]
Map.toAscList Map a prob
d1)

fromFreqs :: (Fractional prob) => [(a,prob)] -> T prob a
fromFreqs :: forall prob a. Fractional prob => [(a, prob)] -> T prob a
fromFreqs [(a, prob)]
xs = forall prob a. [(a, prob)] -> T prob a
Cons (forall a b. (a -> b) -> [a] -> [b]
List.map (\(a
x,prob
p)->(a
x,prob
pforall a. Fractional a => a -> a -> a
/prob
q)) [(a, prob)]
xs)
           where q :: prob
q = forall prob a. Num prob => [(a, prob)] -> prob
sumP [(a, prob)]
xs

filter :: (Fractional prob) =>
   (a -> Bool) -> T prob a -> T prob a
filter :: forall prob a.
Fractional prob =>
(a -> Bool) -> T prob a -> T prob a
filter a -> Bool
p = forall prob a. Fractional prob => [(a, prob)] -> T prob a
fromFreqs forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. (a -> Bool) -> [a] -> [a]
List.filter (a -> Bool
p forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall prob a. T prob a -> [(a, prob)]
decons

mapMaybe :: (Fractional prob) =>
   (a -> Maybe b) -> T prob a -> T prob b
mapMaybe :: forall prob a b.
Fractional prob =>
(a -> Maybe b) -> T prob a -> T prob b
mapMaybe a -> Maybe b
f =
   forall prob a. Fractional prob => [(a, prob)] -> T prob a
fromFreqs forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> Maybe b) -> [a] -> [b]
Maybe.mapMaybe (forall (f :: * -> *) a c b.
Functor f =>
(a -> f c) -> (a, b) -> f (c, b)
FuncHT.mapFst a -> Maybe b
f) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall prob a. T prob a -> [(a, prob)]
decons


-- | selecting from distributions
selectP :: (Num prob, Ord prob) => T prob a -> prob -> a
selectP :: forall prob a. (Num prob, Ord prob) => T prob a -> prob -> a
selectP (Cons [(a, prob)]
d) prob
p = forall prob a. (Num prob, Ord prob) => prob -> [(a, prob)] -> a
scanP prob
p [(a, prob)]
d

scanP :: (Num prob, Ord prob) => prob -> [(a,prob)] -> a
scanP :: forall prob a. (Num prob, Ord prob) => prob -> [(a, prob)] -> a
scanP prob
p ((a
x,prob
q):[(a, prob)]
ps) =
   if prob
pforall a. Ord a => a -> a -> Bool
<=prob
q Bool -> Bool -> Bool
|| forall (t :: * -> *) a. Foldable t => t a -> Bool
null [(a, prob)]
ps
     then a
x
     else forall prob a. (Num prob, Ord prob) => prob -> [(a, prob)] -> a
scanP (prob
pforall a. Num a => a -> a -> a
-prob
q) [(a, prob)]
ps
scanP prob
_ [] = forall a. HasCallStack => [Char] -> a
error [Char]
"Probability.scanP: distribution must be non-empty"

infixr 1 ??

(??) :: Num prob => Event a -> T prob a -> prob
?? :: forall prob a. Num prob => Event a -> T prob a -> prob
(??) Event a
p = forall prob a. Num prob => [(a, prob)] -> prob
sumP forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. (a -> Bool) -> [a] -> [a]
List.filter (Event a
p forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall prob a. T prob a -> [(a, prob)]
decons


-- | expectation value
expected :: (Num a) => T a a -> a
expected :: forall a. Num a => T a a -> a
expected = forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
List.map (\(a
x,a
p) -> a
x forall a. Num a => a -> a -> a
* a
p) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall prob a. T prob a -> [(a, prob)]
decons

-- | statistical analyses
variance :: (Num a) => T a a -> a
variance :: forall a. Num a => T a a -> a
variance T a a
x =
   forall a. Num a => T a a -> a
expected (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((forall a b. (Num a, Integral b) => a -> b -> a
^(Int
2::Int)) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Num a => a -> a -> a
subtract (forall a. Num a => T a a -> a
expected T a a
x)) T a a
x)

stdDev :: (Floating a) => T a a -> a
stdDev :: forall a. Floating a => T a a -> a
stdDev = forall prob. Floating prob => T prob
sqrt forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Num a => T a a -> a
variance