{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE RankNTypes #-}
{-# OPTIONS_GHC -Wno-deprecations #-}

-- |
-- Module      : Control.Monad.Bayes.Population
-- Description : Representation of distributions using multiple samples
-- Copyright   : (c) Adam Scibior, 2015-2020
-- License     : MIT
-- Maintainer  : leonhard.markert@tweag.io
-- Stability   : experimental
-- Portability : GHC
--
-- 'Population' turns a single sample into a collection of weighted samples.
module Control.Monad.Bayes.Population
  ( Population,
    population,
    runPopulation,
    explicitPopulation,
    fromWeightedList,
    spawn,
    multinomial,
    resampleMultinomial,
    systematic,
    resampleSystematic,
    stratified,
    resampleStratified,
    extractEvidence,
    pushEvidence,
    proper,
    evidence,
    hoist,
    collapse,
    popAvg,
    withParticles,
  )
where

import Control.Arrow (second)
import Control.Monad (replicateM)
import Control.Monad.Bayes.Class
  ( MonadDistribution (categorical, logCategorical, random, uniform),
    MonadFactor,
    MonadMeasure,
    factor,
  )
import Control.Monad.Bayes.Weighted
  ( Weighted,
    applyWeight,
    extractWeight,
    weighted,
    withWeight,
  )
import Control.Monad.List (ListT (..), MonadIO, MonadTrans (..))
import Data.List (unfoldr)
import Data.List qualified
import Data.Maybe (catMaybes)
import Data.Vector ((!))
import Data.Vector qualified as V
import Numeric.Log (Log, ln, sum)
import Numeric.Log qualified as Log
import Prelude hiding (all, sum)

-- | A collection of weighted samples, or particles.
newtype Population m a = Population (Weighted (ListT m) a)
  deriving newtype (forall a b. a -> Population m b -> Population m a
forall a b. (a -> b) -> Population m a -> Population m b
forall (m :: * -> *) a b.
Functor m =>
a -> Population m b -> Population m a
forall (m :: * -> *) a b.
Functor m =>
(a -> b) -> Population m a -> Population m b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> Population m b -> Population m a
$c<$ :: forall (m :: * -> *) a b.
Functor m =>
a -> Population m b -> Population m a
fmap :: forall a b. (a -> b) -> Population m a -> Population m b
$cfmap :: forall (m :: * -> *) a b.
Functor m =>
(a -> b) -> Population m a -> Population m b
Functor, forall a. a -> Population m a
forall a b. Population m a -> Population m b -> Population m a
forall a b. Population m a -> Population m b -> Population m b
forall a b.
Population m (a -> b) -> Population m a -> Population m b
forall a b c.
(a -> b -> c) -> Population m a -> Population m b -> Population m c
forall {m :: * -> *}. Monad m => Functor (Population m)
forall (m :: * -> *) a. Monad m => a -> Population m a
forall (m :: * -> *) a b.
Monad m =>
Population m a -> Population m b -> Population m a
forall (m :: * -> *) a b.
Monad m =>
Population m a -> Population m b -> Population m b
forall (m :: * -> *) a b.
Monad m =>
Population m (a -> b) -> Population m a -> Population m b
forall (m :: * -> *) a b c.
Monad m =>
(a -> b -> c) -> Population m a -> Population m b -> Population m c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
<* :: forall a b. Population m a -> Population m b -> Population m a
$c<* :: forall (m :: * -> *) a b.
Monad m =>
Population m a -> Population m b -> Population m a
*> :: forall a b. Population m a -> Population m b -> Population m b
$c*> :: forall (m :: * -> *) a b.
Monad m =>
Population m a -> Population m b -> Population m b
liftA2 :: forall a b c.
(a -> b -> c) -> Population m a -> Population m b -> Population m c
$cliftA2 :: forall (m :: * -> *) a b c.
Monad m =>
(a -> b -> c) -> Population m a -> Population m b -> Population m c
<*> :: forall a b.
Population m (a -> b) -> Population m a -> Population m b
$c<*> :: forall (m :: * -> *) a b.
Monad m =>
Population m (a -> b) -> Population m a -> Population m b
pure :: forall a. a -> Population m a
$cpure :: forall (m :: * -> *) a. Monad m => a -> Population m a
Applicative, forall a. a -> Population m a
forall a b. Population m a -> Population m b -> Population m b
forall a b.
Population m a -> (a -> Population m b) -> Population m b
forall (m :: * -> *). Monad m => Applicative (Population m)
forall (m :: * -> *) a. Monad m => a -> Population m a
forall (m :: * -> *) a b.
Monad m =>
Population m a -> Population m b -> Population m b
forall (m :: * -> *) a b.
Monad m =>
Population m a -> (a -> Population m b) -> Population m b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
return :: forall a. a -> Population m a
$creturn :: forall (m :: * -> *) a. Monad m => a -> Population m a
>> :: forall a b. Population m a -> Population m b -> Population m b
$c>> :: forall (m :: * -> *) a b.
Monad m =>
Population m a -> Population m b -> Population m b
>>= :: forall a b.
Population m a -> (a -> Population m b) -> Population m b
$c>>= :: forall (m :: * -> *) a b.
Monad m =>
Population m a -> (a -> Population m b) -> Population m b
Monad, forall a. IO a -> Population m a
forall (m :: * -> *).
Monad m -> (forall a. IO a -> m a) -> MonadIO m
forall {m :: * -> *}. MonadIO m => Monad (Population m)
forall (m :: * -> *) a. MonadIO m => IO a -> Population m a
liftIO :: forall a. IO a -> Population m a
$cliftIO :: forall (m :: * -> *) a. MonadIO m => IO a -> Population m a
MonadIO, Population m Double
Double -> Population m Bool
Double -> Population m Int
Double -> Double -> Population m Double
forall a. [a] -> Population m a
forall (m :: * -> *).
Monad m
-> m Double
-> (Double -> Double -> m Double)
-> (Double -> Double -> m Double)
-> (Double -> Double -> m Double)
-> (Double -> Double -> m Double)
-> (Double -> m Bool)
-> (forall (v :: * -> *). Vector v Double => v Double -> m Int)
-> (forall (v :: * -> *).
    (Vector v (Log Double), Vector v Double) =>
    v (Log Double) -> m Int)
-> (forall a. [a] -> m a)
-> (Double -> m Int)
-> (Double -> m Int)
-> (forall (v :: * -> *).
    Vector v Double =>
    v Double -> m (v Double))
-> MonadDistribution m
forall (v :: * -> *).
Vector v Double =>
v Double -> Population m (v Double)
forall (v :: * -> *).
Vector v Double =>
v Double -> Population m Int
forall (v :: * -> *).
(Vector v (Log Double), Vector v Double) =>
v (Log Double) -> Population m Int
forall {m :: * -> *}. MonadDistribution m => Monad (Population m)
forall (m :: * -> *). MonadDistribution m => Population m Double
forall (m :: * -> *).
MonadDistribution m =>
Double -> Population m Bool
forall (m :: * -> *).
MonadDistribution m =>
Double -> Population m Int
forall (m :: * -> *).
MonadDistribution m =>
Double -> Double -> Population m Double
forall (m :: * -> *) a.
MonadDistribution m =>
[a] -> Population m a
forall (m :: * -> *) (v :: * -> *).
(MonadDistribution m, Vector v Double) =>
v Double -> Population m (v Double)
forall (m :: * -> *) (v :: * -> *).
(MonadDistribution m, Vector v Double) =>
v Double -> Population m Int
forall (m :: * -> *) (v :: * -> *).
(MonadDistribution m, Vector v (Log Double), Vector v Double) =>
v (Log Double) -> Population m Int
dirichlet :: forall (v :: * -> *).
Vector v Double =>
v Double -> Population m (v Double)
$cdirichlet :: forall (m :: * -> *) (v :: * -> *).
(MonadDistribution m, Vector v Double) =>
v Double -> Population m (v Double)
poisson :: Double -> Population m Int
$cpoisson :: forall (m :: * -> *).
MonadDistribution m =>
Double -> Population m Int
geometric :: Double -> Population m Int
$cgeometric :: forall (m :: * -> *).
MonadDistribution m =>
Double -> Population m Int
uniformD :: forall a. [a] -> Population m a
$cuniformD :: forall (m :: * -> *) a.
MonadDistribution m =>
[a] -> Population m a
logCategorical :: forall (v :: * -> *).
(Vector v (Log Double), Vector v Double) =>
v (Log Double) -> Population m Int
$clogCategorical :: forall (m :: * -> *) (v :: * -> *).
(MonadDistribution m, Vector v (Log Double), Vector v Double) =>
v (Log Double) -> Population m Int
categorical :: forall (v :: * -> *).
Vector v Double =>
v Double -> Population m Int
$ccategorical :: forall (m :: * -> *) (v :: * -> *).
(MonadDistribution m, Vector v Double) =>
v Double -> Population m Int
bernoulli :: Double -> Population m Bool
$cbernoulli :: forall (m :: * -> *).
MonadDistribution m =>
Double -> Population m Bool
beta :: Double -> Double -> Population m Double
$cbeta :: forall (m :: * -> *).
MonadDistribution m =>
Double -> Double -> Population m Double
gamma :: Double -> Double -> Population m Double
$cgamma :: forall (m :: * -> *).
MonadDistribution m =>
Double -> Double -> Population m Double
normal :: Double -> Double -> Population m Double
$cnormal :: forall (m :: * -> *).
MonadDistribution m =>
Double -> Double -> Population m Double
uniform :: Double -> Double -> Population m Double
$cuniform :: forall (m :: * -> *).
MonadDistribution m =>
Double -> Double -> Population m Double
random :: Population m Double
$crandom :: forall (m :: * -> *). MonadDistribution m => Population m Double
MonadDistribution, Log Double -> Population m ()
forall (m :: * -> *). Monad m => Monad (Population m)
forall (m :: * -> *). Monad m => Log Double -> Population m ()
forall (m :: * -> *).
Monad m -> (Log Double -> m ()) -> MonadFactor m
score :: Log Double -> Population m ()
$cscore :: forall (m :: * -> *). Monad m => Log Double -> Population m ()
MonadFactor, forall {m :: * -> *}.
MonadDistribution m =>
MonadFactor (Population m)
forall (m :: * -> *).
MonadDistribution m =>
MonadDistribution (Population m)
forall (m :: * -> *).
MonadDistribution m -> MonadFactor m -> MonadMeasure m
MonadMeasure)

instance MonadTrans Population where
  lift :: forall (m :: * -> *) a. Monad m => m a -> Population m a
lift = forall (m :: * -> *) a. Weighted (ListT m) a -> Population m a
Population forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift

-- | Explicit representation of the weighted sample with weights in the log
-- domain.
population, runPopulation :: Population m a -> m [(a, Log Double)]
population :: forall (m :: * -> *) a. Population m a -> m [(a, Log Double)]
population (Population Weighted (ListT m) a
m) = forall (m :: * -> *) a. ListT m a -> m [a]
runListT forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Weighted m a -> m (a, Log Double)
weighted Weighted (ListT m) a
m

-- | deprecated synonym
runPopulation :: forall (m :: * -> *) a. Population m a -> m [(a, Log Double)]
runPopulation = forall (m :: * -> *) a. Population m a -> m [(a, Log Double)]
population

-- | Explicit representation of the weighted sample.
explicitPopulation :: Functor m => Population m a -> m [(a, Double)]
explicitPopulation :: forall (m :: * -> *) a.
Functor m =>
Population m a -> m [(a, Double)]
explicitPopulation = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a b. (a -> b) -> [a] -> [b]
map (forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (forall a. Floating a => a -> a
exp forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Log a -> a
ln))) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. Population m a -> m [(a, Log Double)]
population

-- | Initialize 'Population' with a concrete weighted sample.
fromWeightedList :: Monad m => m [(a, Log Double)] -> Population m a
fromWeightedList :: forall (m :: * -> *) a.
Monad m =>
m [(a, Log Double)] -> Population m a
fromWeightedList = forall (m :: * -> *) a. Weighted (ListT m) a -> Population m a
Population forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a.
Monad m =>
m (a, Log Double) -> Weighted m a
withWeight forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. m [a] -> ListT m a
ListT

-- | Increase the sample size by a given factor.
-- The weights are adjusted such that their sum is preserved.
-- It is therefore safe to use 'spawn' in arbitrary places in the program
-- without introducing bias.
spawn :: Monad m => Int -> Population m ()
spawn :: forall (m :: * -> *). Monad m => Int -> Population m ()
spawn Int
n = forall (m :: * -> *) a.
Monad m =>
m [(a, Log Double)] -> Population m a
fromWeightedList forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. Int -> a -> [a]
replicate Int
n ((), Log Double
1 forall a. Fractional a => a -> a -> a
/ forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n)

withParticles :: Monad m => Int -> Population m a -> Population m a
withParticles :: forall (m :: * -> *) a.
Monad m =>
Int -> Population m a -> Population m a
withParticles Int
n = (forall (m :: * -> *). Monad m => Int -> Population m ()
spawn Int
n forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>>)

resampleGeneric ::
  MonadDistribution m =>
  -- | resampler
  (V.Vector Double -> m [Int]) ->
  Population m a ->
  Population m a
resampleGeneric :: forall (m :: * -> *) a.
MonadDistribution m =>
(Vector Double -> m [Int]) -> Population m a -> Population m a
resampleGeneric Vector Double -> m [Int]
resampler Population m a
m = forall (m :: * -> *) a.
Monad m =>
m [(a, Log Double)] -> Population m a
fromWeightedList forall a b. (a -> b) -> a -> b
$ do
  [(a, Log Double)]
pop <- forall (m :: * -> *) a. Population m a -> m [(a, Log Double)]
population Population m a
m
  let ([a]
xs, [Log Double]
ps) = forall a b. [(a, b)] -> ([a], [b])
unzip [(a, Log Double)]
pop
  let n :: Int
n = forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs
  let z :: Log Double
z = forall a (f :: * -> *).
(RealFloat a, Foldable f) =>
f (Log a) -> Log a
Log.sum [Log Double]
ps
  if Log Double
z forall a. Ord a => a -> a -> Bool
> Log Double
0
    then do
      let weights :: Vector Double
weights = forall a. [a] -> Vector a
V.fromList (forall a b. (a -> b) -> [a] -> [b]
map (forall a. Floating a => a -> a
exp forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Log a -> a
ln forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a. Fractional a => a -> a -> a
/ Log Double
z)) [Log Double]
ps)
      [Int]
ancestors <- Vector Double -> m [Int]
resampler Vector Double
weights
      let xvec :: Vector a
xvec = forall a. [a] -> Vector a
V.fromList [a]
xs
      let offsprings :: [a]
offsprings = forall a b. (a -> b) -> [a] -> [b]
map (Vector a
xvec forall a. Vector a -> Int -> a
V.!) [Int]
ancestors
      forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (,Log Double
z forall a. Fractional a => a -> a -> a
/ forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n) [a]
offsprings
    else -- if all weights are zero do not resample
      forall (m :: * -> *) a. Monad m => a -> m a
return [(a, Log Double)]
pop

-- | Systematic sampler.
-- Sample \(n\) values from \((0,1]\) as follows
-- \[
-- \begin{aligned}
-- u^{(1)} &\sim U\left(0, \frac{1}{n}\right] \\
-- u^{(i)} &=u^{(1)}+\frac{i-1}{n}, \quad i=2,3, \ldots, n
-- \end{aligned}
-- \]
-- and then pick integers \(m\) according to
-- \[
-- Q^{(m-1)}<u^{(n)} \leq Q^{(m)}
-- \]
-- where
-- \[
-- Q^{(m)}=\sum_{k=1}^{m} w^{(k)}
-- \]
-- and \(w^{(k)}\) are the weights. See also [Comparison of Resampling Schemes for Particle Filtering](https://arxiv.org/abs/cs/0507025).
systematic :: Double -> V.Vector Double -> [Int]
systematic :: Double -> Vector Double -> [Int]
systematic Double
u Vector Double
ps = Int -> Double -> Int -> Double -> [Int] -> [Int]
f Int
0 (Double
u forall a. Fractional a => a -> a -> a
/ forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n) Int
0 Double
0 []
  where
    prob :: Int -> Double
prob Int
i = Vector Double
ps forall a. Vector a -> Int -> a
V.! Int
i
    n :: Int
n = forall (t :: * -> *) a. Foldable t => t a -> Int
length Vector Double
ps
    inc :: Double
inc = Double
1 forall a. Fractional a => a -> a -> a
/ forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n
    f :: Int -> Double -> Int -> Double -> [Int] -> [Int]
f Int
i Double
_ Int
_ Double
_ [Int]
acc | Int
i forall a. Eq a => a -> a -> Bool
== Int
n = [Int]
acc
    f Int
i Double
v Int
j Double
q [Int]
acc =
      if Double
v forall a. Ord a => a -> a -> Bool
< Double
q
        then Int -> Double -> Int -> Double -> [Int] -> [Int]
f (Int
i forall a. Num a => a -> a -> a
+ Int
1) (Double
v forall a. Num a => a -> a -> a
+ Double
inc) Int
j Double
q (Int
j forall a. Num a => a -> a -> a
- Int
1 forall a. a -> [a] -> [a]
: [Int]
acc)
        else Int -> Double -> Int -> Double -> [Int] -> [Int]
f Int
i Double
v (Int
j forall a. Num a => a -> a -> a
+ Int
1) (Double
q forall a. Num a => a -> a -> a
+ Int -> Double
prob Int
j) [Int]
acc

-- | Resample the population using the underlying monad and a systematic resampling scheme.
-- The total weight is preserved.
resampleSystematic ::
  (MonadDistribution m) =>
  Population m a ->
  Population m a
resampleSystematic :: forall (m :: * -> *) a.
MonadDistribution m =>
Population m a -> Population m a
resampleSystematic = forall (m :: * -> *) a.
MonadDistribution m =>
(Vector Double -> m [Int]) -> Population m a -> Population m a
resampleGeneric (\Vector Double
ps -> (Double -> Vector Double -> [Int]
`systematic` Vector Double
ps) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *). MonadDistribution m => m Double
random)

-- | Stratified sampler.
--
-- Sample \(n\) values from \((0,1]\) as follows
-- \[
-- u^{(i)} \sim U\left(\frac{i-1}{n}, \frac{i}{n}\right], \quad i=1,2, \ldots, n
-- \]
-- and then pick integers \(m\) according to
-- \[
-- Q^{(m-1)}<u^{(n)} \leq Q^{(m)}
-- \]
-- where
-- \[
-- Q^{(m)}=\sum_{k=1}^{m} w^{(k)}
-- \]
-- and \(w^{(k)}\) are the weights.
--
-- The conditional variance of stratified sampling is always smaller than that of multinomial sampling and it is also unbiased - see  [Comparison of Resampling Schemes for Particle Filtering](https://arxiv.org/abs/cs/0507025).
stratified :: MonadDistribution m => V.Vector Double -> m [Int]
stratified :: forall (m :: * -> *).
MonadDistribution m =>
Vector Double -> m [Int]
stratified Vector Double
weights = do
  let bigN :: Int
bigN = forall a. Vector a -> Int
V.length Vector Double
weights
  Vector Double
dithers <- forall (m :: * -> *) a. Monad m => Int -> m a -> m (Vector a)
V.replicateM Int
bigN (forall (m :: * -> *).
MonadDistribution m =>
Double -> Double -> m Double
uniform Double
0.0 Double
1.0)
  let positions :: Vector Double
positions =
        forall a b. (a -> b) -> Vector a -> Vector b
V.map (forall a. Fractional a => a -> a -> a
/ forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
bigN) forall a b. (a -> b) -> a -> b
$
          forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
V.zipWith forall a. Num a => a -> a -> a
(+) Vector Double
dithers (forall a b. (a -> b) -> Vector a -> Vector b
V.map forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ forall a. [a] -> Vector a
V.fromList [Int
0 .. Int
bigN forall a. Num a => a -> a -> a
- Int
1])
      cumulativeSum :: Vector Double
cumulativeSum = forall a b. (a -> b -> a) -> a -> Vector b -> Vector a
V.scanl forall a. Num a => a -> a -> a
(+) Double
0.0 Vector Double
weights
      coalg :: (Int, Int) -> Maybe (Maybe Int, (Int, Int))
coalg (Int
i, Int
j)
        | Int
i forall a. Ord a => a -> a -> Bool
< Int
bigN =
            if (Vector Double
positions forall a. Vector a -> Int -> a
! Int
i) forall a. Ord a => a -> a -> Bool
< (Vector Double
cumulativeSum forall a. Vector a -> Int -> a
! Int
j)
              then forall a. a -> Maybe a
Just (forall a. a -> Maybe a
Just Int
j, (Int
i forall a. Num a => a -> a -> a
+ Int
1, Int
j))
              else forall a. a -> Maybe a
Just (forall a. Maybe a
Nothing, (Int
i, Int
j forall a. Num a => a -> a -> a
+ Int
1))
        | Bool
otherwise =
            forall a. Maybe a
Nothing
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (\Int
i -> Int
i forall a. Num a => a -> a -> a
- Int
1) forall a b. (a -> b) -> a -> b
$ forall a. [Maybe a] -> [a]
catMaybes forall a b. (a -> b) -> a -> b
$ forall b a. (b -> Maybe (a, b)) -> b -> [a]
unfoldr (Int, Int) -> Maybe (Maybe Int, (Int, Int))
coalg (Int
0, Int
0)

-- | Resample the population using the underlying monad and a stratified resampling scheme.
-- The total weight is preserved.
resampleStratified ::
  (MonadDistribution m) =>
  Population m a ->
  Population m a
resampleStratified :: forall (m :: * -> *) a.
MonadDistribution m =>
Population m a -> Population m a
resampleStratified = forall (m :: * -> *) a.
MonadDistribution m =>
(Vector Double -> m [Int]) -> Population m a -> Population m a
resampleGeneric forall (m :: * -> *).
MonadDistribution m =>
Vector Double -> m [Int]
stratified

-- | Multinomial sampler.  Sample from \(0, \ldots, n - 1\) \(n\)
-- times drawn at random according to the weights where \(n\) is the
-- length of vector of weights.
multinomial :: MonadDistribution m => V.Vector Double -> m [Int]
multinomial :: forall (m :: * -> *).
MonadDistribution m =>
Vector Double -> m [Int]
multinomial Vector Double
ps = forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM (forall a. Vector a -> Int
V.length Vector Double
ps) (forall (m :: * -> *) (v :: * -> *).
(MonadDistribution m, Vector v Double) =>
v Double -> m Int
categorical Vector Double
ps)

-- | Resample the population using the underlying monad and a multinomial resampling scheme.
-- The total weight is preserved.
resampleMultinomial ::
  (MonadDistribution m) =>
  Population m a ->
  Population m a
resampleMultinomial :: forall (m :: * -> *) a.
MonadDistribution m =>
Population m a -> Population m a
resampleMultinomial = forall (m :: * -> *) a.
MonadDistribution m =>
(Vector Double -> m [Int]) -> Population m a -> Population m a
resampleGeneric forall (m :: * -> *).
MonadDistribution m =>
Vector Double -> m [Int]
multinomial

-- | Separate the sum of weights into the 'Weighted' transformer.
-- Weights are normalized after this operation.
extractEvidence ::
  Monad m =>
  Population m a ->
  Population (Weighted m) a
extractEvidence :: forall (m :: * -> *) a.
Monad m =>
Population m a -> Population (Weighted m) a
extractEvidence Population m a
m = forall (m :: * -> *) a.
Monad m =>
m [(a, Log Double)] -> Population m a
fromWeightedList forall a b. (a -> b) -> a -> b
$ do
  [(a, Log Double)]
pop <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Population m a -> m [(a, Log Double)]
population Population m a
m
  let ([a]
xs, [Log Double]
ps) = forall a b. [(a, b)] -> ([a], [b])
unzip [(a, Log Double)]
pop
  let z :: Log Double
z = forall a (f :: * -> *).
(RealFloat a, Foldable f) =>
f (Log a) -> Log a
sum [Log Double]
ps
  let ws :: [Log Double]
ws = forall a b. (a -> b) -> [a] -> [b]
map (if Log Double
z forall a. Ord a => a -> a -> Bool
> Log Double
0 then (forall a. Fractional a => a -> a -> a
/ Log Double
z) else forall a b. a -> b -> a
const (Log Double
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 [Log Double]
ps))) [Log Double]
ps
  forall (m :: * -> *). MonadFactor m => Log Double -> m ()
factor Log Double
z
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip [a]
xs [Log Double]
ws

-- | Push the evidence estimator as a score to the transformed monad.
-- Weights are normalized after this operation.
pushEvidence ::
  MonadFactor m =>
  Population m a ->
  Population m a
pushEvidence :: forall (m :: * -> *) a.
MonadFactor m =>
Population m a -> Population m a
pushEvidence = forall (n :: * -> *) (m :: * -> *) a.
Monad n =>
(forall x. m x -> n x) -> Population m a -> Population n a
hoist forall (m :: * -> *) a. MonadFactor m => Weighted m a -> m a
applyWeight forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a.
Monad m =>
Population m a -> Population (Weighted m) a
extractEvidence

-- | A properly weighted single sample, that is one picked at random according
-- to the weights, with the sum of all weights.
proper ::
  (MonadDistribution m) =>
  Population m a ->
  Weighted m a
proper :: forall (m :: * -> *) a.
MonadDistribution m =>
Population m a -> Weighted m a
proper Population m a
m = do
  [(a, Log Double)]
pop <- forall (m :: * -> *) a. Population m a -> m [(a, Log Double)]
population forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
Monad m =>
Population m a -> Population (Weighted m) a
extractEvidence Population m a
m
  let ([a]
xs, [Log Double]
ps) = forall a b. [(a, b)] -> ([a], [b])
unzip [(a, Log Double)]
pop
  Int
index <- forall (m :: * -> *) (v :: * -> *).
(MonadDistribution m, Vector v (Log Double), Vector v Double) =>
v (Log Double) -> m Int
logCategorical forall a b. (a -> b) -> a -> b
$ forall a. [a] -> Vector a
V.fromList [Log Double]
ps
  let x :: a
x = [a]
xs forall a. [a] -> Int -> a
!! Int
index
  forall (m :: * -> *) a. Monad m => a -> m a
return a
x

-- | Model evidence estimator, also known as pseudo-marginal likelihood.
evidence :: (Monad m) => Population m a -> m (Log Double)
evidence :: forall (m :: * -> *) a. Monad m => Population m a -> m (Log Double)
evidence = forall (m :: * -> *) a. Functor m => Weighted m a -> m (Log Double)
extractWeight forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. Population m a -> m [(a, Log Double)]
population forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a.
Monad m =>
Population m a -> Population (Weighted m) a
extractEvidence

-- | Picks one point from the population and uses model evidence as a 'score'
-- in the transformed monad.
-- This way a single sample can be selected from a population without
-- introducing bias.
collapse ::
  (MonadMeasure m) =>
  Population m a ->
  m a
collapse :: forall (m :: * -> *) a. MonadMeasure m => Population m a -> m a
collapse = forall (m :: * -> *) a. MonadFactor m => Weighted m a -> m a
applyWeight forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a.
MonadDistribution m =>
Population m a -> Weighted m a
proper

-- | Population average of a function, computed using unnormalized weights.
popAvg :: (Monad m) => (a -> Double) -> Population m a -> m Double
popAvg :: forall (m :: * -> *) a.
Monad m =>
(a -> Double) -> Population m a -> m Double
popAvg a -> Double
f Population m a
p = do
  [(a, Double)]
xs <- forall (m :: * -> *) a.
Functor m =>
Population m a -> m [(a, Double)]
explicitPopulation Population m a
p
  let ys :: [Double]
ys = forall a b. (a -> b) -> [a] -> [b]
map (\(a
x, Double
w) -> a -> Double
f a
x forall a. Num a => a -> a -> a
* Double
w) [(a, Double)]
xs
  let t :: Double
t = forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
Data.List.sum [Double]
ys
  forall (m :: * -> *) a. Monad m => a -> m a
return Double
t

-- | Applies a transformation to the inner monad.
hoist ::
  Monad n =>
  (forall x. m x -> n x) ->
  Population m a ->
  Population n a
hoist :: forall (n :: * -> *) (m :: * -> *) a.
Monad n =>
(forall x. m x -> n x) -> Population m a -> Population n a
hoist forall x. m x -> n x
f = forall (m :: * -> *) a.
Monad m =>
m [(a, Log Double)] -> Population m a
fromWeightedList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall x. m x -> n x
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. Population m a -> m [(a, Log Double)]
population