{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE TupleSections #-}

-- | This module defines the 'Distribution' monad and its operations.
--
-- A @'Distribution' a@ is a discrete probability distribution over values of
-- type @a@.
--
-- You can define distributions in several ways:
--
-- * Choosing from the common distributions exported by this module, such as
--   a 'categorical', 'uniform', 'geometric', 'bernoulli', 'binomial',
--   'negativeBinomial', 'hypergeometric', or 'poisson' distribution.
-- * Operating on existing distributions using the 'Functor', 'Applicative', and
--   'Monad' instances, or by conditioning on events using 'conditional' or
--   'finiteConditional'.
--
-- Once you have a distribution, you can sample from it using 'sample', list its
-- outcomes and their probabilities using 'probabilities' or 'possibilities',
-- and compute various statistics using 'probability', 'approxProbability',
-- 'expectation', 'variance', 'stddev', 'entropy', 'relativeEntropy', or
-- 'mutualInformation'.
--
-- It's important to make a distinction between *finite* and *infinite*
-- distributions.  An infinite distribution is one whose list of 'possibilities'
-- is infinite.  Note that this *includes* distributions for which there are
-- only finitely many distinct outcomes, but still an infinite number of paths
-- to reach these outcomes.  Infinite distributions typically arise from
-- recursive expressions.  Certain functions only work on finite distributions,
-- and will hang or OOM if given an infinite distribution.
--
-- For example, if you express the process of rolling a six-sided die, but
-- always rerolling if the result is one, then there are five distinct outcomes:
-- 2, 3, 4, 5, or 6.  Nevertheless, this is an infinite distribution, because
-- it's possible to roll any number of ones prior to the final result.
module Probability.Distribution
  ( -- * Types
    Distribution,
    Event,
    RandVar,
    EventView (..),

    -- * Basic operations
    possibilities,
    probabilities,
    simplify,
    sample,
    viewEvent,
    fromEventView,
    finitize,
    conditional,
    finiteConditional,
    bayesian,
    finiteBayesian,

    -- * Common distributions
    categorical,
    uniform,
    geometric,
    bernoulli,
    binomial,
    negativeBinomial,
    hypergeometric,
    poisson,

    -- * Analysis
    probability,
    probabilityBounds,
    approxProbability,
    expectation,
    variance,
    stddev,
    entropy,
    relativeEntropy,
    mutualInformation,
  )
where

import Control.Applicative (liftA2)
import Control.Monad (ap)
import Data.Bifunctor (Bifunctor (..))
import Data.Bool (bool)
import Data.Map (Map)
import qualified Data.Map as Map
import qualified Data.Sequence as Seq
import qualified Data.Set as Set
import Data.Tuple (swap)
import System.Random (randomRIO)

-- | A probability distribution of values.
data Distribution prob a
  = Certainly a
  | Choice prob (Distribution prob a) (Distribution prob a)
  deriving (a -> Distribution prob b -> Distribution prob a
(a -> b) -> Distribution prob a -> Distribution prob b
(forall a b.
 (a -> b) -> Distribution prob a -> Distribution prob b)
-> (forall a b. a -> Distribution prob b -> Distribution prob a)
-> Functor (Distribution prob)
forall a b. a -> Distribution prob b -> Distribution prob a
forall a b. (a -> b) -> Distribution prob a -> Distribution prob b
forall prob a b. a -> Distribution prob b -> Distribution prob a
forall prob a b.
(a -> b) -> Distribution prob a -> Distribution prob b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: a -> Distribution prob b -> Distribution prob a
$c<$ :: forall prob a b. a -> Distribution prob b -> Distribution prob a
fmap :: (a -> b) -> Distribution prob a -> Distribution prob b
$cfmap :: forall prob a b.
(a -> b) -> Distribution prob a -> Distribution prob b
Functor)

instance Bifunctor Distribution where
  bimap :: (a -> b) -> (c -> d) -> Distribution a c -> Distribution b d
bimap a -> b
_ c -> d
g (Certainly c
a) = d -> Distribution b d
forall prob a. a -> Distribution prob a
Certainly (c -> d
g c
a)
  bimap a -> b
f c -> d
g (Choice a
p Distribution a c
a Distribution a c
b) = b -> Distribution b d -> Distribution b d -> Distribution b d
forall prob a.
prob
-> Distribution prob a
-> Distribution prob a
-> Distribution prob a
Choice (a -> b
f a
p) ((a -> b) -> (c -> d) -> Distribution a c -> Distribution b d
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap a -> b
f c -> d
g Distribution a c
a) ((a -> b) -> (c -> d) -> Distribution a c -> Distribution b d
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap a -> b
f c -> d
g Distribution a c
b)

instance Applicative (Distribution prob) where
  pure :: a -> Distribution prob a
pure = a -> Distribution prob a
forall prob a. a -> Distribution prob a
Certainly
  <*> :: Distribution prob (a -> b)
-> Distribution prob a -> Distribution prob b
(<*>) = Distribution prob (a -> b)
-> Distribution prob a -> Distribution prob b
forall (m :: * -> *) a b. Monad m => m (a -> b) -> m a -> m b
ap

instance Monad (Distribution prob) where
  Certainly a
x >>= :: Distribution prob a
-> (a -> Distribution prob b) -> Distribution prob b
>>= a -> Distribution prob b
f = a -> Distribution prob b
f a
x
  Choice prob
p Distribution prob a
a Distribution prob a
b >>= a -> Distribution prob b
f = prob
-> Distribution prob b
-> Distribution prob b
-> Distribution prob b
forall prob a.
prob
-> Distribution prob a
-> Distribution prob a
-> Distribution prob a
Choice prob
p (Distribution prob a
a Distribution prob a
-> (a -> Distribution prob b) -> Distribution prob b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= a -> Distribution prob b
f) (Distribution prob a
b Distribution prob a
-> (a -> Distribution prob b) -> Distribution prob b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= a -> Distribution prob b
f)

instance Num a => Num (Distribution prob a) where
  + :: Distribution prob a -> Distribution prob a -> Distribution prob a
(+) = (a -> a -> a)
-> Distribution prob a
-> Distribution prob a
-> Distribution prob a
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 a -> a -> a
forall a. Num a => a -> a -> a
(+)
  (-) = (a -> a -> a)
-> Distribution prob a
-> Distribution prob a
-> Distribution prob a
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 (-)
  * :: Distribution prob a -> Distribution prob a -> Distribution prob a
(*) = (a -> a -> a)
-> Distribution prob a
-> Distribution prob a
-> Distribution prob a
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 a -> a -> a
forall a. Num a => a -> a -> a
(*)
  abs :: Distribution prob a -> Distribution prob a
abs = (a -> a) -> Distribution prob a -> Distribution prob a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> a
forall a. Num a => a -> a
abs
  signum :: Distribution prob a -> Distribution prob a
signum = (a -> a) -> Distribution prob a -> Distribution prob a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> a
forall a. Num a => a -> a
signum
  negate :: Distribution prob a -> Distribution prob a
negate = (a -> a) -> Distribution prob a -> Distribution prob a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> a
forall a. Num a => a -> a
negate
  fromInteger :: Integer -> Distribution prob a
fromInteger = a -> Distribution prob a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> Distribution prob a)
-> (Integer -> a) -> Integer -> Distribution prob a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> a
forall a. Num a => Integer -> a
fromInteger

instance Fractional a => Fractional (Distribution prob a) where
  / :: Distribution prob a -> Distribution prob a -> Distribution prob a
(/) = (a -> a -> a)
-> Distribution prob a
-> Distribution prob a
-> Distribution prob a
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 a -> a -> a
forall a. Fractional a => a -> a -> a
(/)
  recip :: Distribution prob a -> Distribution prob a
recip = (a -> a) -> Distribution prob a -> Distribution prob a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> a
forall a. Fractional a => a -> a
recip
  fromRational :: Rational -> Distribution prob a
fromRational = a -> Distribution prob a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> Distribution prob a)
-> (Rational -> a) -> Rational -> Distribution prob a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Rational -> a
forall a. Fractional a => Rational -> a
fromRational

-- | An event is a predicate on values from a sample space.
type Event s = s -> Bool

-- | A random variable is a function mapping each element of a sample space to
-- the corresponding value of the random variable.
type RandVar s a = s -> a

-- | Gives the list of all possible values of a given probability distribution.
-- The list will often contain multiple entries for the same outcome, in which
-- case the true probability for that outcome is the sum of probabilities of all
-- entries.
--
-- In the finite case, multiple entries can be combined by using 'simplify' on
-- the 'Distribution' first.
possibilities :: Num prob => Distribution prob a -> [(prob, a)]
possibilities :: Distribution prob a -> [(prob, a)]
possibilities Distribution prob a
dist = Seq (prob, Distribution prob a) -> [(prob, a)]
forall a b. Num a => Seq (a, Distribution a b) -> [(a, b)]
go ((prob, Distribution prob a) -> Seq (prob, Distribution prob a)
forall a. a -> Seq a
Seq.singleton (prob
1, Distribution prob a
dist))
  where
    go :: Seq (a, Distribution a b) -> [(a, b)]
go Seq (a, Distribution a b)
Seq.Empty = []
    go ((a
p, Certainly b
x) Seq.:<| Seq (a, Distribution a b)
queue') = (a
p, b
x) (a, b) -> [(a, b)] -> [(a, b)]
forall a. a -> [a] -> [a]
: Seq (a, Distribution a b) -> [(a, b)]
go Seq (a, Distribution a b)
queue'
    go ((a
p, Choice a
q Distribution a b
a Distribution a b
b) Seq.:<| Seq (a, Distribution a b)
queue') =
      Seq (a, Distribution a b) -> [(a, b)]
go (Seq (a, Distribution a b)
queue' Seq (a, Distribution a b)
-> (a, Distribution a b) -> Seq (a, Distribution a b)
forall a. Seq a -> a -> Seq a
Seq.:|> (a
p a -> a -> a
forall a. Num a => a -> a -> a
* a
q, Distribution a b
a) Seq (a, Distribution a b)
-> (a, Distribution a b) -> Seq (a, Distribution a b)
forall a. Seq a -> a -> Seq a
Seq.:|> (a
p a -> a -> a
forall a. Num a => a -> a -> a
* (a
1 a -> a -> a
forall a. Num a => a -> a -> a
- a
q), Distribution a b
b))

-- | Truncates an infinite distribution to make it finite.  The epsilon
-- parameter is the amount of tail probability that you're willing to ignore
-- and assign to an arbitrary outcome.
finitize ::
  (Fractional prob, Ord prob) =>
  prob ->
  Distribution prob a ->
  Distribution prob a
finitize :: prob -> Distribution prob a -> Distribution prob a
finitize prob
epsilon = [(prob, a)] -> Distribution prob a
forall prob a.
Fractional prob =>
[(prob, a)] -> Distribution prob a
categorical ([(prob, a)] -> Distribution prob a)
-> (Distribution prob a -> [(prob, a)])
-> Distribution prob a
-> Distribution prob a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. prob -> [(prob, a)] -> [(prob, a)]
forall b. prob -> [(prob, b)] -> [(prob, b)]
go prob
1 ([(prob, a)] -> [(prob, a)])
-> (Distribution prob a -> [(prob, a)])
-> Distribution prob a
-> [(prob, a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Distribution prob a -> [(prob, a)]
forall prob a. Num prob => Distribution prob a -> [(prob, a)]
possibilities
  where
    go :: prob -> [(prob, b)] -> [(prob, b)]
go prob
q ((prob
p, b
x) : [(prob, b)]
poss)
      | prob
q prob -> prob -> prob
forall a. Num a => a -> a -> a
- prob
p prob -> prob -> Bool
forall a. Ord a => a -> a -> Bool
< prob
epsilon = [(prob
q, b
x)]
      | Bool
otherwise = (prob
p, b
x) (prob, b) -> [(prob, b)] -> [(prob, b)]
forall a. a -> [a] -> [a]
: prob -> [(prob, b)] -> [(prob, b)]
go (prob
q prob -> prob -> prob
forall a. Num a => a -> a -> a
- prob
p) [(prob, b)]
poss
    go prob
_ [] = []

-- | Gives a map from outcomes to their probabilities in the given distribution.
--
-- This only works for finite distributions.  Infinite distributions (including
-- even distributions with finitely many outcomes, but infinitely many paths to
-- reach those outcomes) will hang.
probabilities :: (Num prob, Ord a) => Distribution prob a -> Map a prob
probabilities :: Distribution prob a -> Map a prob
probabilities = (prob -> prob -> prob) -> [(a, prob)] -> Map a prob
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith prob -> prob -> prob
forall a. Num a => a -> a -> a
(+) ([(a, prob)] -> Map a prob)
-> (Distribution prob a -> [(a, prob)])
-> Distribution prob a
-> Map a prob
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((prob, a) -> (a, prob)) -> [(prob, a)] -> [(a, prob)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (prob, a) -> (a, prob)
forall a b. (a, b) -> (b, a)
swap ([(prob, a)] -> [(a, prob)])
-> (Distribution prob a -> [(prob, a)])
-> Distribution prob a
-> [(a, prob)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Distribution prob a -> [(prob, a)]
forall prob a. Num prob => Distribution prob a -> [(prob, a)]
possibilities

-- | Simplifies a finite distribution.
--
-- This only works for finite distributions.  Infinite distributions (including
-- even distributions with finitely many outcomes, but infinitely many paths to
-- reach those outcomes) will hang.
simplify ::
  (Fractional prob, Ord a) => Distribution prob a -> Distribution prob a
simplify :: Distribution prob a -> Distribution prob a
simplify = [(prob, a)] -> Distribution prob a
forall prob a.
Fractional prob =>
[(prob, a)] -> Distribution prob a
categorical ([(prob, a)] -> Distribution prob a)
-> (Distribution prob a -> [(prob, a)])
-> Distribution prob a
-> Distribution prob a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((a, prob) -> (prob, a)) -> [(a, prob)] -> [(prob, a)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (a, prob) -> (prob, a)
forall a b. (a, b) -> (b, a)
swap ([(a, prob)] -> [(prob, a)])
-> (Distribution prob a -> [(a, prob)])
-> Distribution prob a
-> [(prob, a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map a prob -> [(a, prob)]
forall k a. Map k a -> [(k, a)]
Map.toList (Map a prob -> [(a, prob)])
-> (Distribution prob a -> Map a prob)
-> Distribution prob a
-> [(a, prob)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Distribution prob a -> Map a prob
forall prob a.
(Num prob, Ord a) =>
Distribution prob a -> Map a prob
probabilities

-- | Samples the probability distribution to produce a value.
sample :: Real prob => Distribution prob a -> IO a
sample :: Distribution prob a -> IO a
sample (Certainly a
x) = a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return a
x
sample (Choice prob
p Distribution prob a
a Distribution prob a
b) =
  IO a -> IO a -> Bool -> IO a
forall a. a -> a -> Bool -> a
bool (Distribution prob a -> IO a
forall prob a. Real prob => Distribution prob a -> IO a
sample Distribution prob a
b) (Distribution prob a -> IO a
forall prob a. Real prob => Distribution prob a -> IO a
sample Distribution prob a
a) (Bool -> IO a) -> (Double -> Bool) -> Double -> IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< prob -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac prob
p) (Double -> IO a) -> IO Double -> IO a
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Double, Double) -> IO Double
forall a (m :: * -> *). (Random a, MonadIO m) => (a, a) -> m a
randomRIO (Double
0 :: Double, Double
1)

-- | A view of a probability distribution from the point of view of a given
-- event.  The event either always happens, never happens, or happens sometimes
-- with some probability.  In the latter case, there are posterior distributions
-- for when the event does or does not happen.
data EventView prob s
  = Always (Distribution prob s)
  | Never (Distribution prob s)
  | Sometimes prob (Distribution prob s) (Distribution prob s)

-- | Gives a view on a probability distribution relative to some event.
--
-- The following are guaranteed.
-- 1. @'fromEventView' . 'viewEvent' ev = id@
-- 2. If @'viewEvent' ev dist = 'Always' dist'@, then @dist = dist'@ and
--    @'probability' ev dist = 1@.
-- 3. If @'viewEvent' ev dist = 'Never' dist'@, then @dist = dist'@ and
--    @'probability' ev dist = 0@.
-- 4. If @'viewEvent' ev dist = 'Sometimes' p a b@, then
--    @'probability' ev dist = p@ and:
--    * @dist = 'bernoulli' p >>= bool a b@
--    * @'probability' ev a = 1@
--    * @'probability' ev b = 0@
--
-- This only works for finite distributions.  Infinite distributions (including
-- even distributions with finitely many outcomes, but infinitely many paths to
-- reach those outcomes) will hang.
viewEvent ::
  Fractional prob =>
  Event s ->
  Distribution prob s ->
  EventView prob s
viewEvent :: Event s -> Distribution prob s -> EventView prob s
viewEvent Event s
event dist :: Distribution prob s
dist@(Certainly s
x)
  | Event s
event s
x = Distribution prob s -> EventView prob s
forall prob s. Distribution prob s -> EventView prob s
Always Distribution prob s
dist
  | Bool
otherwise = Distribution prob s -> EventView prob s
forall prob s. Distribution prob s -> EventView prob s
Never Distribution prob s
dist
viewEvent Event s
event dist :: Distribution prob s
dist@(Choice prob
p Distribution prob s
aa Distribution prob s
bb) =
  case (Event s -> Distribution prob s -> EventView prob s
forall prob s.
Fractional prob =>
Event s -> Distribution prob s -> EventView prob s
viewEvent Event s
event Distribution prob s
aa, Event s -> Distribution prob s -> EventView prob s
forall prob s.
Fractional prob =>
Event s -> Distribution prob s -> EventView prob s
viewEvent Event s
event Distribution prob s
bb) of
    (Never Distribution prob s
_, Never Distribution prob s
_) -> Distribution prob s -> EventView prob s
forall prob s. Distribution prob s -> EventView prob s
Never Distribution prob s
dist
    (Always Distribution prob s
_, Always Distribution prob s
_) -> Distribution prob s -> EventView prob s
forall prob s. Distribution prob s -> EventView prob s
Always Distribution prob s
dist
    (Always Distribution prob s
a, Never Distribution prob s
b) -> prob
-> Distribution prob s -> Distribution prob s -> EventView prob s
forall prob s.
prob
-> Distribution prob s -> Distribution prob s -> EventView prob s
Sometimes prob
p Distribution prob s
a Distribution prob s
b
    (Never Distribution prob s
a, Always Distribution prob s
b) -> prob
-> Distribution prob s -> Distribution prob s -> EventView prob s
forall prob s.
prob
-> Distribution prob s -> Distribution prob s -> EventView prob s
Sometimes (prob
1 prob -> prob -> prob
forall a. Num a => a -> a -> a
- prob
p) Distribution prob s
b Distribution prob s
a
    (Sometimes prob
q Distribution prob s
a1 Distribution prob s
a2, Never Distribution prob s
b) ->
      let (prob
p', prob
_, prob
p2) = prob -> prob -> (prob, prob, prob)
blend prob
q prob
0 in prob
-> Distribution prob s -> Distribution prob s -> EventView prob s
forall prob s.
prob
-> Distribution prob s -> Distribution prob s -> EventView prob s
Sometimes prob
p' Distribution prob s
a1 (prob
-> Distribution prob s
-> Distribution prob s
-> Distribution prob s
forall prob a.
prob
-> Distribution prob a
-> Distribution prob a
-> Distribution prob a
Choice prob
p2 Distribution prob s
a2 Distribution prob s
b)
    (Sometimes prob
q Distribution prob s
a1 Distribution prob s
a2, Always Distribution prob s
b) ->
      let (prob
p', prob
p1, prob
_) = prob -> prob -> (prob, prob, prob)
blend prob
q prob
1 in prob
-> Distribution prob s -> Distribution prob s -> EventView prob s
forall prob s.
prob
-> Distribution prob s -> Distribution prob s -> EventView prob s
Sometimes prob
p' (prob
-> Distribution prob s
-> Distribution prob s
-> Distribution prob s
forall prob a.
prob
-> Distribution prob a
-> Distribution prob a
-> Distribution prob a
Choice prob
p1 Distribution prob s
a1 Distribution prob s
b) Distribution prob s
a2
    (Never Distribution prob s
a, Sometimes prob
r Distribution prob s
b1 Distribution prob s
b2) ->
      let (prob
p', prob
_, prob
p2) = prob -> prob -> (prob, prob, prob)
blend prob
0 prob
r in prob
-> Distribution prob s -> Distribution prob s -> EventView prob s
forall prob s.
prob
-> Distribution prob s -> Distribution prob s -> EventView prob s
Sometimes prob
p' Distribution prob s
b1 (prob
-> Distribution prob s
-> Distribution prob s
-> Distribution prob s
forall prob a.
prob
-> Distribution prob a
-> Distribution prob a
-> Distribution prob a
Choice prob
p2 Distribution prob s
a Distribution prob s
b2)
    (Always Distribution prob s
a, Sometimes prob
r Distribution prob s
b1 Distribution prob s
b2) ->
      let (prob
p', prob
p1, prob
_) = prob -> prob -> (prob, prob, prob)
blend prob
1 prob
r in prob
-> Distribution prob s -> Distribution prob s -> EventView prob s
forall prob s.
prob
-> Distribution prob s -> Distribution prob s -> EventView prob s
Sometimes prob
p' (prob
-> Distribution prob s
-> Distribution prob s
-> Distribution prob s
forall prob a.
prob
-> Distribution prob a
-> Distribution prob a
-> Distribution prob a
Choice prob
p1 Distribution prob s
a Distribution prob s
b1) Distribution prob s
b2
    (Sometimes prob
q Distribution prob s
a1 Distribution prob s
a2, Sometimes prob
r Distribution prob s
b1 Distribution prob s
b2) ->
      let (prob
p', prob
p1, prob
p2) = prob -> prob -> (prob, prob, prob)
blend prob
q prob
r
       in prob
-> Distribution prob s -> Distribution prob s -> EventView prob s
forall prob s.
prob
-> Distribution prob s -> Distribution prob s -> EventView prob s
Sometimes prob
p' (prob
-> Distribution prob s
-> Distribution prob s
-> Distribution prob s
forall prob a.
prob
-> Distribution prob a
-> Distribution prob a
-> Distribution prob a
Choice prob
p1 Distribution prob s
a1 Distribution prob s
b1) (prob
-> Distribution prob s
-> Distribution prob s
-> Distribution prob s
forall prob a.
prob
-> Distribution prob a
-> Distribution prob a
-> Distribution prob a
Choice prob
p2 Distribution prob s
a2 Distribution prob s
b2)
  where
    blend :: prob -> prob -> (prob, prob, prob)
blend prob
q prob
r =
      let p' :: prob
p' = prob
p prob -> prob -> prob
forall a. Num a => a -> a -> a
* prob
q prob -> prob -> prob
forall a. Num a => a -> a -> a
+ (prob
1 prob -> prob -> prob
forall a. Num a => a -> a -> a
- prob
p) prob -> prob -> prob
forall a. Num a => a -> a -> a
* prob
r
       in (prob
p', prob
p prob -> prob -> prob
forall a. Num a => a -> a -> a
* prob
q prob -> prob -> prob
forall a. Fractional a => a -> a -> a
/ prob
p', prob
p prob -> prob -> prob
forall a. Num a => a -> a -> a
* (prob
1 prob -> prob -> prob
forall a. Num a => a -> a -> a
- prob
q) prob -> prob -> prob
forall a. Fractional a => a -> a -> a
/ (prob
1 prob -> prob -> prob
forall a. Num a => a -> a -> a
- prob
p'))

-- | Converts from 'EventView' back to a 'Distribution'.  The resulting
-- distribution is equivalent to the source distribution.
fromEventView :: EventView prob s -> Distribution prob s
fromEventView :: EventView prob s -> Distribution prob s
fromEventView (Always Distribution prob s
dist) = Distribution prob s
dist
fromEventView (Never Distribution prob s
dist) = Distribution prob s
dist
fromEventView (Sometimes prob
p Distribution prob s
a Distribution prob s
b) = prob
-> Distribution prob s
-> Distribution prob s
-> Distribution prob s
forall prob a.
prob
-> Distribution prob a
-> Distribution prob a
-> Distribution prob a
Choice prob
p Distribution prob s
a Distribution prob s
b

-- | Produces the conditional probability distribution, assuming some event.
-- This function works for all distributions, but always produces an infinite
-- distribution for non-trivial events.
conditional :: Event s -> Distribution prob s -> Distribution prob s
conditional :: Event s -> Distribution prob s -> Distribution prob s
conditional Event s
event Distribution prob s
dist = Distribution prob s
cdist
  where
    cdist :: Distribution prob s
cdist = do
      s
x <- Distribution prob s
dist
      if Event s
event s
x
        then s -> Distribution prob s
forall (m :: * -> *) a. Monad m => a -> m a
return s
x
        else Distribution prob s
cdist

-- | Produces the conditional probability distribution, assuming some event.
--
-- This only works for finite distributions.  Infinite distributions (including
-- even distributions with finitely many outcomes, but infinitely many paths to
-- reach those outcomes) will hang.
finiteConditional ::
  Fractional prob => Event s -> Distribution prob s -> Distribution prob s
finiteConditional :: Event s -> Distribution prob s -> Distribution prob s
finiteConditional Event s
event Distribution prob s
dist = [(prob, s)] -> Distribution prob s
forall prob a.
Fractional prob =>
[(prob, a)] -> Distribution prob a
categorical (((prob, s) -> (prob, s)) -> [(prob, s)] -> [(prob, s)]
forall a b. (a -> b) -> [a] -> [b]
map ((prob -> prob) -> (prob, s) -> (prob, s)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (prob -> prob -> prob
forall a. Fractional a => a -> a -> a
/ prob
p_event)) [(prob, s)]
filtered)
  where
    filtered :: [(prob, s)]
filtered = ((prob, s) -> Bool) -> [(prob, s)] -> [(prob, s)]
forall a. (a -> Bool) -> [a] -> [a]
filter (Event s
event Event s -> ((prob, s) -> s) -> (prob, s) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (prob, s) -> s
forall a b. (a, b) -> b
snd) (Distribution prob s -> [(prob, s)]
forall prob a. Num prob => Distribution prob a -> [(prob, a)]
possibilities Distribution prob s
dist)
    p_event :: prob
p_event = [prob] -> prob
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (((prob, s) -> prob) -> [(prob, s)] -> [prob]
forall a b. (a -> b) -> [a] -> [b]
map (prob, s) -> prob
forall a b. (a, b) -> a
fst [(prob, s)]
filtered)

-- | Updates a prior distribution of parameters for a model, based on an
-- observed event.  This implements Bayes' Law for distributions.
--
-- This function works for all distributions, but always produces an infinite
-- distribution for non-trivial events.
bayesian ::
  (param -> Distribution prob s) ->
  Event s ->
  Distribution prob param ->
  Distribution prob param
bayesian :: (param -> Distribution prob s)
-> Event s -> Distribution prob param -> Distribution prob param
bayesian param -> Distribution prob s
model Event s
event Distribution prob param
prior = Distribution prob param
posterior
  where
    posterior :: Distribution prob param
posterior = do
      param
param <- Distribution prob param
prior
      s
x <- param -> Distribution prob s
model param
param
      if Event s
event s
x then param -> Distribution prob param
forall (m :: * -> *) a. Monad m => a -> m a
return param
param else Distribution prob param
posterior

-- | Updates a prior distribution of parameters for a model, based on an
-- observed event.  This implements Bayes' Law for distributions.
--
-- This only works for finite distributions.  Infinite distributions (including
-- even distributions with finitely many outcomes, but infinitely many paths to
-- reach those outcomes) will hang.
finiteBayesian ::
  Fractional prob =>
  (param -> Distribution prob s) ->
  Event s ->
  Distribution prob param ->
  Distribution prob param
finiteBayesian :: (param -> Distribution prob s)
-> Event s -> Distribution prob param -> Distribution prob param
finiteBayesian param -> Distribution prob s
model Event s
event Distribution prob param
prior = case Event (param, s)
-> Distribution prob (param, s) -> EventView prob (param, s)
forall prob s.
Fractional prob =>
Event s -> Distribution prob s -> EventView prob s
viewEvent (Event s
event Event s -> ((param, s) -> s) -> Event (param, s)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (param, s) -> s
forall a b. (a, b) -> b
snd) Distribution prob (param, s)
withParam of
  Always Distribution prob (param, s)
dist -> (param, s) -> param
forall a b. (a, b) -> a
fst ((param, s) -> param)
-> Distribution prob (param, s) -> Distribution prob param
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Distribution prob (param, s)
dist
  Never Distribution prob (param, s)
_ -> [Char] -> Distribution prob param
forall a. HasCallStack => [Char] -> a
error [Char]
"Posterior is undefined for an impossible event"
  Sometimes prob
_ Distribution prob (param, s)
dist Distribution prob (param, s)
_ -> (param, s) -> param
forall a b. (a, b) -> a
fst ((param, s) -> param)
-> Distribution prob (param, s) -> Distribution prob param
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Distribution prob (param, s)
dist
  where
    withParam :: Distribution prob (param, s)
withParam = do
      param
param <- Distribution prob param
prior
      s
obs <- param -> Distribution prob s
model param
param
      (param, s) -> Distribution prob (param, s)
forall (m :: * -> *) a. Monad m => a -> m a
return (param
param, s
obs)

-- | A distribution with a fixed probability for each outcome.  The
-- probabilities should add to 1, but this is not checked.
categorical :: Fractional prob => [(prob, a)] -> Distribution prob a
categorical :: [(prob, a)] -> Distribution prob a
categorical = prob -> [(prob, a)] -> Distribution prob a
forall prob a.
Fractional prob =>
prob -> [(prob, a)] -> Distribution prob a
go prob
1
  where
    go :: prob -> [(prob, a)] -> Distribution prob a
go prob
_ [] = [Char] -> Distribution prob a
forall a. HasCallStack => [Char] -> a
error [Char]
"Empty distribution is not allowed"
    go prob
_ [(prob
_, a
x)] = a -> Distribution prob a
forall prob a. a -> Distribution prob a
Certainly a
x
    go prob
p ((prob
q, a
x) : [(prob, a)]
xs) = prob
-> Distribution prob a
-> Distribution prob a
-> Distribution prob a
forall prob a.
prob
-> Distribution prob a
-> Distribution prob a
-> Distribution prob a
Choice (prob
q prob -> prob -> prob
forall a. Fractional a => a -> a -> a
/ prob
p) (a -> Distribution prob a
forall prob a. a -> Distribution prob a
Certainly a
x) (prob -> [(prob, a)] -> Distribution prob a
go (prob
p prob -> prob -> prob
forall a. Num a => a -> a -> a
- prob
q) [(prob, a)]
xs)

-- | A uniform distribution over a list of values.
uniform :: Fractional prob => [a] -> Distribution prob a
uniform :: [a] -> Distribution prob a
uniform [a]
xs = [(prob, a)] -> Distribution prob a
forall prob a.
Fractional prob =>
[(prob, a)] -> Distribution prob a
categorical ([(prob, a)] -> Distribution prob a)
-> [(prob, a)] -> Distribution prob a
forall a b. (a -> b) -> a -> b
$ (prob -> prob
forall a. Fractional a => a -> a
recip prob
n,) (a -> (prob, a)) -> [a] -> [(prob, a)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [a]
xs where n :: prob
n = Int -> prob
forall a b. (Real a, Fractional b) => a -> b
realToFrac ([a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs)

-- | Geometric distribution over a list of possibilities.
geometric :: prob -> [a] -> Distribution prob a
geometric :: prob -> [a] -> Distribution prob a
geometric prob
_ [] = [Char] -> Distribution prob a
forall a. HasCallStack => [Char] -> a
error [Char]
"geometric: Empty distribution is not allowed"
geometric prob
_ [a
x] = a -> Distribution prob a
forall prob a. a -> Distribution prob a
Certainly a
x
geometric prob
p (a
x : [a]
xs) = prob
-> Distribution prob a
-> Distribution prob a
-> Distribution prob a
forall prob a.
prob
-> Distribution prob a
-> Distribution prob a
-> Distribution prob a
Choice prob
p (a -> Distribution prob a
forall prob a. a -> Distribution prob a
Certainly a
x) (prob -> [a] -> Distribution prob a
forall prob a. prob -> [a] -> Distribution prob a
geometric prob
p [a]
xs)

-- | A Bernoulli distribution.  This gives True with probability @p@, and False
-- otherwise.
bernoulli :: prob -> Distribution prob Bool
bernoulli :: prob -> Distribution prob Bool
bernoulli prob
p = prob
-> Distribution prob Bool
-> Distribution prob Bool
-> Distribution prob Bool
forall prob a.
prob
-> Distribution prob a
-> Distribution prob a
-> Distribution prob a
Choice prob
p (Bool -> Distribution prob Bool
forall prob a. a -> Distribution prob a
Certainly Bool
True) (Bool -> Distribution prob Bool
forall prob a. a -> Distribution prob a
Certainly Bool
False)

-- | Computes nCk.  This is a building block for several well-known discrete
-- distributions.
choose :: Integral t => t -> t -> t
t
n choose :: t -> t -> t
`choose` t
k
  | t
k t -> t -> Bool
forall a. Ord a => a -> a -> Bool
> t
n t -> t -> t
forall a. Integral a => a -> a -> a
`div` t
2 = t
n t -> t -> t
forall a. Integral a => a -> a -> a
`choose` (t
n t -> t -> t
forall a. Num a => a -> a -> a
- t
k)
  | Bool
otherwise = [t] -> t
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [t
n t -> t -> t
forall a. Num a => a -> a -> a
- t
k t -> t -> t
forall a. Num a => a -> a -> a
+ t
1 .. t
n] t -> t -> t
forall a. Integral a => a -> a -> a
`div` [t] -> t
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [t
1 .. t
k]

-- | A binomial distribution.  This gives the distribution of number of
-- successes in @n@ trials with probability @p@ of success.
binomial :: (Fractional prob, Integral n) => prob -> n -> Distribution prob n
binomial :: prob -> n -> Distribution prob n
binomial prob
p n
n =
  [(prob, n)] -> Distribution prob n
forall prob a.
Fractional prob =>
[(prob, a)] -> Distribution prob a
categorical
    [ (n -> prob
forall a b. (Integral a, Num b) => a -> b
fromIntegral (n
n n -> n -> n
forall a. Integral a => a -> a -> a
`choose` n
k) prob -> prob -> prob
forall a. Num a => a -> a -> a
* prob
p prob -> n -> prob
forall a b. (Num a, Integral b) => a -> b -> a
^ n
k prob -> prob -> prob
forall a. Num a => a -> a -> a
* (prob
1 prob -> prob -> prob
forall a. Num a => a -> a -> a
- prob
p) prob -> n -> prob
forall a b. (Num a, Integral b) => a -> b -> a
^ (n
n n -> n -> n
forall a. Num a => a -> a -> a
- n
k), n
k)
      | n
k <- [n
0 .. n
n]
    ]

-- | Negative binomial distribution.  This gives the distribution of number of
-- failures before @r@ successes with probability @p@ of success.
negativeBinomial ::
  (Fractional prob, Integral n) => prob -> n -> Distribution prob n
negativeBinomial :: prob -> n -> Distribution prob n
negativeBinomial prob
_ n
0 = n -> Distribution prob n
forall (f :: * -> *) a. Applicative f => a -> f a
pure n
0
negativeBinomial prob
p n
r =
  [(prob, n)] -> Distribution prob n
forall prob a.
Fractional prob =>
[(prob, a)] -> Distribution prob a
categorical
    [ (n -> prob
forall a b. (Integral a, Num b) => a -> b
fromIntegral ((n
k n -> n -> n
forall a. Num a => a -> a -> a
+ n
r n -> n -> n
forall a. Num a => a -> a -> a
- n
1) n -> n -> n
forall a. Integral a => a -> a -> a
`choose` (n
r n -> n -> n
forall a. Num a => a -> a -> a
- n
1)) prob -> prob -> prob
forall a. Num a => a -> a -> a
* prob
p prob -> n -> prob
forall a b. (Num a, Integral b) => a -> b -> a
^ n
r prob -> prob -> prob
forall a. Num a => a -> a -> a
* (prob
1 prob -> prob -> prob
forall a. Num a => a -> a -> a
- prob
p) prob -> n -> prob
forall a b. (Num a, Integral b) => a -> b -> a
^ n
k, n
k)
      | n
k <- [n
0 ..]
    ]

-- | Hypergeometric distribution.  This gives the distribution of number of
-- successful draws out of @n@ attempts without replacement, when @k@
-- possibilities are successful.
hypergeometric ::
  (Fractional prob, Integral n) => n -> n -> n -> Distribution prob n
hypergeometric :: n -> n -> n -> Distribution prob n
hypergeometric n
pop n
k n
n =
  [(prob, n)] -> Distribution prob n
forall prob a.
Fractional prob =>
[(prob, a)] -> Distribution prob a
categorical
    [ ( n -> prob
forall a b. (Integral a, Num b) => a -> b
fromIntegral ((n
k n -> n -> n
forall a. Integral a => a -> a -> a
`choose` n
m) n -> n -> n
forall a. Num a => a -> a -> a
* ((n
pop n -> n -> n
forall a. Num a => a -> a -> a
- n
k) n -> n -> n
forall a. Integral a => a -> a -> a
`choose` (n
n n -> n -> n
forall a. Num a => a -> a -> a
- n
m)))
          prob -> prob -> prob
forall a. Fractional a => a -> a -> a
/ n -> prob
forall a b. (Integral a, Num b) => a -> b
fromIntegral (n
pop n -> n -> n
forall a. Integral a => a -> a -> a
`choose` n
n),
        n
m
      )
      | n
m <- [n
lo .. n
hi]
    ]
  where
    lo :: n
lo = n -> n -> n
forall a. Ord a => a -> a -> a
max n
0 (n
n n -> n -> n
forall a. Num a => a -> a -> a
+ n
k n -> n -> n
forall a. Num a => a -> a -> a
- n
pop)
    hi :: n
hi = n -> n -> n
forall a. Ord a => a -> a -> a
min n
n n
k

-- | Poisson distribution.  Gives the number of independent events occurring in
-- a fixed time interval, if events are occurring at the given expected rate per
-- time interval.
poisson :: (Floating prob, Integral n) => prob -> Distribution prob n
poisson :: prob -> Distribution prob n
poisson prob
lambda =
  [(prob, n)] -> Distribution prob n
forall prob a.
Fractional prob =>
[(prob, a)] -> Distribution prob a
categorical
    [ (prob
lambda prob -> n -> prob
forall a b. (Num a, Integral b) => a -> b -> a
^ n
k prob -> prob -> prob
forall a. Num a => a -> a -> a
* prob -> prob
forall a. Floating a => a -> a
exp (-prob
lambda) prob -> prob -> prob
forall a. Fractional a => a -> a -> a
/ n -> prob
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([n] -> n
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [n
1 .. n
k]), n
k)
      | n
k <- [n
0 ..]
    ]

-- | Computes the probability of an event, represented by a predicate on values.
--
-- This only works for finite distributions.  Infinite distributions (including
-- even distributions with finitely many outcomes, but infinitely many paths to
-- reach those outcomes) will hang.
probability :: Num prob => Event s -> Distribution prob s -> prob
probability :: Event s -> Distribution prob s -> prob
probability Event s
event (Certainly s
x) = if Event s
event s
x then prob
1 else prob
0
probability Event s
event (Choice prob
p Distribution prob s
a Distribution prob s
b) =
  prob
p prob -> prob -> prob
forall a. Num a => a -> a -> a
* Event s -> Distribution prob s -> prob
forall prob s. Num prob => Event s -> Distribution prob s -> prob
probability Event s
event Distribution prob s
a prob -> prob -> prob
forall a. Num a => a -> a -> a
+ (prob
1 prob -> prob -> prob
forall a. Num a => a -> a -> a
- prob
p) prob -> prob -> prob
forall a. Num a => a -> a -> a
* Event s -> Distribution prob s -> prob
forall prob s. Num prob => Event s -> Distribution prob s -> prob
probability Event s
event Distribution prob s
b

-- | Like probability, but produces a lazy list of ever-improving bounds on the
-- probability.  This can be used on infinite distributions, for which the
-- exact probability cannot be calculated.
probabilityBounds ::
  Num prob => Event s -> Distribution prob s -> [(prob, prob)]
probabilityBounds :: Event s -> Distribution prob s -> [(prob, prob)]
probabilityBounds Event s
event Distribution prob s
dist = prob -> prob -> [(prob, s)] -> [(prob, prob)]
forall a. Num a => a -> a -> [(a, s)] -> [(a, a)]
go prob
0 prob
1 (Distribution prob s -> [(prob, s)]
forall prob a. Num prob => Distribution prob a -> [(prob, a)]
possibilities Distribution prob s
dist)
  where
    go :: a -> a -> [(a, s)] -> [(a, a)]
go a
p a
_ [] = [(a
p, a
p)]
    go a
p a
q ((a
q', s
x) : [(a, s)]
xs)
      | Event s
event s
x = (a
p, a
p a -> a -> a
forall a. Num a => a -> a -> a
+ a
q) (a, a) -> [(a, a)] -> [(a, a)]
forall a. a -> [a] -> [a]
: a -> a -> [(a, s)] -> [(a, a)]
go (a
p a -> a -> a
forall a. Num a => a -> a -> a
+ a
q') (a
q a -> a -> a
forall a. Num a => a -> a -> a
- a
q') [(a, s)]
xs
      | Bool
otherwise = (a
p, a
p a -> a -> a
forall a. Num a => a -> a -> a
+ a
q) (a, a) -> [(a, a)] -> [(a, a)]
forall a. a -> [a] -> [a]
: a -> a -> [(a, s)] -> [(a, a)]
go a
p (a
q a -> a -> a
forall a. Num a => a -> a -> a
- a
q') [(a, s)]
xs

-- | Like probability, but produces a value that differs from the true
-- probability by at most epsilon. This can be used on infinite distributions,
-- for which the exact probability cannot be calculated.
approxProbability ::
  (Ord prob, Fractional prob) =>
  prob ->
  Event s ->
  Distribution prob s ->
  prob
approxProbability :: prob -> Event s -> Distribution prob s -> prob
approxProbability prob
epsilon Event s
event Distribution prob s
dist =
  (prob -> prob -> prob
forall a. Fractional a => a -> a -> a
/ prob
2) (prob -> prob)
-> ([(prob, prob)] -> prob) -> [(prob, prob)] -> prob
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (prob -> prob -> prob) -> (prob, prob) -> prob
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry prob -> prob -> prob
forall a. Num a => a -> a -> a
(+) ((prob, prob) -> prob)
-> ([(prob, prob)] -> (prob, prob)) -> [(prob, prob)] -> prob
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(prob, prob)] -> (prob, prob)
forall a. [a] -> a
head ([(prob, prob)] -> (prob, prob))
-> ([(prob, prob)] -> [(prob, prob)])
-> [(prob, prob)]
-> (prob, prob)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((prob, prob) -> Bool) -> [(prob, prob)] -> [(prob, prob)]
forall a. (a -> Bool) -> [a] -> [a]
dropWhile ((prob -> prob -> Bool
forall a. Ord a => a -> a -> Bool
> prob
epsilon) (prob -> Bool) -> ((prob, prob) -> prob) -> (prob, prob) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. prob -> prob
forall a. Num a => a -> a
abs (prob -> prob) -> ((prob, prob) -> prob) -> (prob, prob) -> prob
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (prob -> prob -> prob) -> (prob, prob) -> prob
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (-)) ([(prob, prob)] -> prob) -> [(prob, prob)] -> prob
forall a b. (a -> b) -> a -> b
$
    Event s -> Distribution prob s -> [(prob, prob)]
forall prob s.
Num prob =>
Event s -> Distribution prob s -> [(prob, prob)]
probabilityBounds Event s
event Distribution prob s
dist

-- | Computes the expected value of a finite distribution.
--
-- This only works for finite distributions.  Infinite distributions (including
-- even distributions with finitely many outcomes, but infinitely many paths to
-- reach those outcomes) will hang.
expectation :: Num a => Distribution a a -> a
expectation :: Distribution a a -> a
expectation (Certainly a
x) = a
x
expectation (Choice a
p Distribution a a
a Distribution a a
b) =
  a
p a -> a -> a
forall a. Num a => a -> a -> a
* Distribution a a -> a
forall a. Num a => Distribution a a -> a
expectation Distribution a a
a a -> a -> a
forall a. Num a => a -> a -> a
+ (a
1 a -> a -> a
forall a. Num a => a -> a -> a
- a
p) a -> a -> a
forall a. Num a => a -> a -> a
* Distribution a a -> a
forall a. Num a => Distribution a a -> a
expectation Distribution a a
b

-- | Computes the variance of a finite distribution.
--
-- This only works for finite distributions.  Infinite distributions (including
-- even distributions with finitely many outcomes, but infinitely many paths to
-- reach those outcomes) will hang.
variance :: Num a => Distribution a a -> a
variance :: Distribution a a -> a
variance Distribution a a
dist = Distribution a a -> a
forall a. Num a => Distribution a a -> a
expectation ((a -> Int -> a
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int)) (a -> a) -> (a -> a) -> a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a -> a
forall a. Num a => a -> a -> a
subtract a
mean (a -> a) -> Distribution a a -> Distribution a a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Distribution a a
dist)
  where
    mean :: a
mean = Distribution a a -> a
forall a. Num a => Distribution a a -> a
expectation Distribution a a
dist

-- | Computes the standard deviation of a finite distribution.
--
-- This only works for finite distributions.  Infinite distributions (including
-- even distributions with finitely many outcomes, but infinitely many paths to
-- reach those outcomes) will hang.
stddev :: Floating a => Distribution a a -> a
stddev :: Distribution a a -> a
stddev Distribution a a
dist = a -> a
forall a. Floating a => a -> a
sqrt (Distribution a a -> a
forall a. Num a => Distribution a a -> a
variance Distribution a a
dist)

-- | Computes the entropy of a distribution in bits.
--
-- This only works for finite distributions.  Infinite distributions (including
-- even distributions with finitely many outcomes, but infinitely many paths to
-- reach those outcomes) will hang.
entropy :: (Floating prob, Ord a) => Distribution prob a -> prob
entropy :: Distribution prob a -> prob
entropy Distribution prob a
dist =
  [prob] -> prob
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum
    [ -prob
p prob -> prob -> prob
forall a. Num a => a -> a -> a
* prob -> prob -> prob
forall a. Floating a => a -> a -> a
logBase prob
2 prob
p
      | prob
p <- ((prob, a) -> prob) -> [(prob, a)] -> [prob]
forall a b. (a -> b) -> [a] -> [b]
map (prob, a) -> prob
forall a b. (a, b) -> a
fst (Distribution prob a -> [(prob, a)]
forall prob a. Num prob => Distribution prob a -> [(prob, a)]
possibilities (Distribution prob a -> Distribution prob a
forall prob a.
(Fractional prob, Ord a) =>
Distribution prob a -> Distribution prob a
simplify Distribution prob a
dist))
    ]

-- | Computes the relative entropy, also known as Kullback-Leibler divergence,
-- between two distributions in bits.
--
-- This only works for finite distributions.  Infinite distributions (including
-- even distributions with finitely many outcomes, but infinitely many paths to
-- reach those outcomes) will hang.
relativeEntropy ::
  (Eq prob, Floating prob, Ord a) =>
  Distribution prob a ->
  Distribution prob a ->
  prob
relativeEntropy :: Distribution prob a -> Distribution prob a -> prob
relativeEntropy Distribution prob a
post Distribution prob a
prior = [prob] -> prob
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (a -> prob
term (a -> prob) -> [a] -> [prob]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set a -> [a]
forall a. Set a -> [a]
Set.toList Set a
vals)
  where
    prob_post :: Map a prob
prob_post = Distribution prob a -> Map a prob
forall prob a.
(Num prob, Ord a) =>
Distribution prob a -> Map a prob
probabilities Distribution prob a
post
    prob_prior :: Map a prob
prob_prior = Distribution prob a -> Map a prob
forall prob a.
(Num prob, Ord a) =>
Distribution prob a -> Map a prob
probabilities Distribution prob a
prior
    vals :: Set a
vals = Map a prob -> Set a
forall k a. Map k a -> Set k
Map.keysSet Map a prob
prob_post Set a -> Set a -> Set a
forall a. Ord a => Set a -> Set a -> Set a
`Set.union` Map a prob -> Set a
forall k a. Map k a -> Set k
Map.keysSet Map a prob
prob_prior
    term :: a -> prob
term a
x =
      let p :: prob
p = prob -> a -> Map a prob -> prob
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault prob
0 a
x Map a prob
prob_post
          q :: prob
q = prob -> a -> Map a prob -> prob
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault prob
0 a
x Map a prob
prob_prior
       in if prob
p prob -> prob -> Bool
forall a. Eq a => a -> a -> Bool
== prob
0 then prob
0 else prob
p prob -> prob -> prob
forall a. Num a => a -> a -> a
* prob -> prob -> prob
forall a. Floating a => a -> a -> a
logBase prob
2 (prob
p prob -> prob -> prob
forall a. Fractional a => a -> a -> a
/ prob
q)

-- | Computes the mutual information between two random variables, in bits.  The
-- given distribution is taken as a definition of a probability space, and the
-- random variables are represented as functions from the sample space to values
-- taken by the random variable.
--
-- This only works for finite distributions.  Infinite distributions (including
-- even distributions with finitely many outcomes, but infinitely many paths to
-- reach those outcomes) will hang.
mutualInformation ::
  (Eq prob, Floating prob, Ord a, Ord b) =>
  RandVar s a ->
  RandVar s b ->
  Distribution prob s ->
  prob
mutualInformation :: RandVar s a -> RandVar s b -> Distribution prob s -> prob
mutualInformation RandVar s a
f RandVar s b
g Distribution prob s
dist =
  [prob] -> prob
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (a -> b -> prob
term (a -> b -> prob) -> [a] -> [b -> prob]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map a prob -> [a]
forall k a. Map k a -> [k]
Map.keys Map a prob
f_probs [b -> prob] -> [b] -> [prob]
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Map b prob -> [b]
forall k a. Map k a -> [k]
Map.keys Map b prob
g_probs)
  where
    joint_probs :: Map (a, b) prob
joint_probs = Distribution prob (a, b) -> Map (a, b) prob
forall prob a.
(Num prob, Ord a) =>
Distribution prob a -> Map a prob
probabilities ((\s
x -> (RandVar s a
f s
x, RandVar s b
g s
x)) (s -> (a, b)) -> Distribution prob s -> Distribution prob (a, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Distribution prob s
dist)
    f_probs :: Map a prob
f_probs = Distribution prob a -> Map a prob
forall prob a.
(Num prob, Ord a) =>
Distribution prob a -> Map a prob
probabilities (RandVar s a
f RandVar s a -> Distribution prob s -> Distribution prob a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Distribution prob s
dist)
    g_probs :: Map b prob
g_probs = Distribution prob b -> Map b prob
forall prob a.
(Num prob, Ord a) =>
Distribution prob a -> Map a prob
probabilities (RandVar s b
g RandVar s b -> Distribution prob s -> Distribution prob b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Distribution prob s
dist)
    term :: a -> b -> prob
term a
x b
y =
      let p_x :: prob
p_x = prob -> a -> Map a prob -> prob
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault prob
0 a
x Map a prob
f_probs
          p_y :: prob
p_y = prob -> b -> Map b prob -> prob
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault prob
0 b
y Map b prob
g_probs
          p_xy :: prob
p_xy = prob -> (a, b) -> Map (a, b) prob -> prob
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault prob
0 (a
x, b
y) Map (a, b) prob
joint_probs
       in if prob
p_xy prob -> prob -> Bool
forall a. Eq a => a -> a -> Bool
== prob
0 then prob
0 else prob
p_xy prob -> prob -> prob
forall a. Num a => a -> a -> a
* prob -> prob -> prob
forall a. Floating a => a -> a -> a
logBase prob
2 (prob
p_xy prob -> prob -> prob
forall a. Fractional a => a -> a -> a
/ (prob
p_x prob -> prob -> prob
forall a. Num a => a -> a -> a
* prob
p_y))