{-# 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 -> b) -> Population m a -> Population m b)
-> (forall a b. a -> Population m b -> Population m a)
-> Functor (Population m)
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
$cfmap :: forall (m :: * -> *) a b.
Functor m =>
(a -> b) -> Population m a -> Population m b
fmap :: forall a b. (a -> b) -> Population m a -> Population m b
$c<$ :: forall (m :: * -> *) a b.
Functor m =>
a -> Population m b -> Population m a
<$ :: forall a b. a -> Population m b -> Population m a
Functor, Functor (Population m)
Functor (Population m)
-> (forall a. a -> Population m a)
-> (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 a b. Population m a -> Population m b -> Population m b)
-> (forall a b. Population m a -> Population m b -> Population m a)
-> Applicative (Population m)
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
$cpure :: forall (m :: * -> *) a. Monad m => a -> Population m a
pure :: forall a. a -> Population m a
$c<*> :: forall (m :: * -> *) a b.
Monad m =>
Population m (a -> b) -> Population m a -> Population m b
<*> :: forall a b.
Population m (a -> b) -> Population m a -> Population m b
$cliftA2 :: forall (m :: * -> *) a b c.
Monad m =>
(a -> b -> c) -> Population m a -> Population m b -> Population m c
liftA2 :: forall a b c.
(a -> b -> c) -> Population m a -> Population m b -> Population m c
$c*> :: forall (m :: * -> *) a b.
Monad m =>
Population m a -> Population m b -> Population m b
*> :: 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 a
<* :: forall a b. Population m a -> Population m b -> Population m a
Applicative, Applicative (Population m)
Applicative (Population m)
-> (forall a b.
    Population m a -> (a -> Population m b) -> Population m b)
-> (forall a b. Population m a -> Population m b -> Population m b)
-> (forall a. a -> Population m a)
-> Monad (Population m)
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
$c>>= :: forall (m :: * -> *) a b.
Monad m =>
Population m a -> (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 -> Population m b -> Population m b
>> :: forall a b. Population m a -> Population m b -> Population m b
$creturn :: forall (m :: * -> *) a. Monad m => a -> Population m a
return :: forall a. a -> Population m a
Monad, Monad (Population m)
Monad (Population m)
-> (forall a. IO a -> Population m a) -> MonadIO (Population m)
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
$cliftIO :: forall (m :: * -> *) a. MonadIO m => IO a -> Population m a
liftIO :: forall a. IO a -> Population m a
MonadIO, Monad (Population m)
Population m Double
Monad (Population m)
-> Population m Double
-> (Double -> Double -> Population m Double)
-> (Double -> Double -> Population m Double)
-> (Double -> Double -> Population m Double)
-> (Double -> Double -> Population m Double)
-> (Double -> Population m Bool)
-> (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 a. [a] -> Population m a)
-> (Double -> Population m Int)
-> (Double -> Population m Int)
-> (forall (v :: * -> *).
    Vector v Double =>
    v Double -> Population m (v Double))
-> MonadDistribution (Population m)
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
$crandom :: forall (m :: * -> *). MonadDistribution m => Population m Double
random :: Population m Double
$cuniform :: forall (m :: * -> *).
MonadDistribution m =>
Double -> Double -> Population m Double
uniform :: Double -> Double -> Population m Double
$cnormal :: forall (m :: * -> *).
MonadDistribution m =>
Double -> Double -> Population m Double
normal :: Double -> Double -> Population m Double
$cgamma :: forall (m :: * -> *).
MonadDistribution m =>
Double -> Double -> Population m Double
gamma :: Double -> Double -> Population m Double
$cbeta :: forall (m :: * -> *).
MonadDistribution m =>
Double -> Double -> Population m Double
beta :: Double -> Double -> Population m Double
$cbernoulli :: forall (m :: * -> *).
MonadDistribution m =>
Double -> Population m Bool
bernoulli :: Double -> Population m Bool
$ccategorical :: forall (m :: * -> *) (v :: * -> *).
(MonadDistribution m, Vector v Double) =>
v Double -> Population m Int
categorical :: forall (v :: * -> *).
Vector v Double =>
v Double -> Population m Int
$clogCategorical :: forall (m :: * -> *) (v :: * -> *).
(MonadDistribution m, Vector v (Log Double), Vector v Double) =>
v (Log Double) -> Population m Int
logCategorical :: forall (v :: * -> *).
(Vector v (Log Double), Vector v Double) =>
v (Log Double) -> Population m Int
$cuniformD :: forall (m :: * -> *) a.
MonadDistribution m =>
[a] -> Population m a
uniformD :: forall a. [a] -> Population m a
$cgeometric :: forall (m :: * -> *).
MonadDistribution m =>
Double -> Population m Int
geometric :: Double -> Population m Int
$cpoisson :: forall (m :: * -> *).
MonadDistribution m =>
Double -> Population m Int
poisson :: Double -> Population m Int
$cdirichlet :: forall (m :: * -> *) (v :: * -> *).
(MonadDistribution m, Vector v Double) =>
v Double -> Population m (v Double)
dirichlet :: forall (v :: * -> *).
Vector v Double =>
v Double -> Population m (v Double)
MonadDistribution, Monad (Population m)
Monad (Population m)
-> (Log Double -> Population m ()) -> MonadFactor (Population m)
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
$cscore :: forall (m :: * -> *). Monad m => Log Double -> Population m ()
score :: Log Double -> Population m ()
MonadFactor, MonadFactor (Population m)
MonadDistribution (Population m)
MonadDistribution (Population m)
-> MonadFactor (Population m) -> MonadMeasure (Population m)
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 = Weighted (ListT m) a -> Population m a
forall (m :: * -> *) a. Weighted (ListT m) a -> Population m a
Population (Weighted (ListT m) a -> Population m a)
-> (m a -> Weighted (ListT m) a) -> m a -> Population m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ListT m a -> Weighted (ListT m) a
forall (m :: * -> *) a. Monad m => m a -> Weighted m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (ListT m a -> Weighted (ListT m) a)
-> (m a -> ListT m a) -> m a -> Weighted (ListT m) a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. m a -> ListT m a
forall (m :: * -> *) a. Monad m => m a -> ListT m a
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) = ListT m (a, Log Double) -> m [(a, Log Double)]
forall (m :: * -> *) a. ListT m a -> m [a]
runListT (ListT m (a, Log Double) -> m [(a, Log Double)])
-> ListT m (a, Log Double) -> m [(a, Log Double)]
forall a b. (a -> b) -> a -> b
$ Weighted (ListT m) a -> ListT m (a, Log Double)
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 = Population m a -> m [(a, Log Double)]
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 = ([(a, Log Double)] -> [(a, Double)])
-> m [(a, Log Double)] -> m [(a, Double)]
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (((a, Log Double) -> (a, Double))
-> [(a, Log Double)] -> [(a, Double)]
forall a b. (a -> b) -> [a] -> [b]
map ((Log Double -> Double) -> (a, Log Double) -> (a, Double)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (Double -> Double
forall a. Floating a => a -> a
exp (Double -> Double)
-> (Log Double -> Double) -> Log Double -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Log Double -> Double
forall a. Log a -> a
ln))) (m [(a, Log Double)] -> m [(a, Double)])
-> (Population m a -> m [(a, Log Double)])
-> Population m a
-> m [(a, Double)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Population m a -> m [(a, Log Double)]
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 = Weighted (ListT m) a -> Population m a
forall (m :: * -> *) a. Weighted (ListT m) a -> Population m a
Population (Weighted (ListT m) a -> Population m a)
-> (m [(a, Log Double)] -> Weighted (ListT m) a)
-> m [(a, Log Double)]
-> Population m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ListT m (a, Log Double) -> Weighted (ListT m) a
forall (m :: * -> *) a.
Monad m =>
m (a, Log Double) -> Weighted m a
withWeight (ListT m (a, Log Double) -> Weighted (ListT m) a)
-> (m [(a, Log Double)] -> ListT m (a, Log Double))
-> m [(a, Log Double)]
-> Weighted (ListT m) a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. m [(a, Log Double)] -> ListT m (a, Log Double)
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 = m [((), Log Double)] -> Population m ()
forall (m :: * -> *) a.
Monad m =>
m [(a, Log Double)] -> Population m a
fromWeightedList (m [((), Log Double)] -> Population m ())
-> m [((), Log Double)] -> Population m ()
forall a b. (a -> b) -> a -> b
$ [((), Log Double)] -> m [((), Log Double)]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([((), Log Double)] -> m [((), Log Double)])
-> [((), Log Double)] -> m [((), Log Double)]
forall a b. (a -> b) -> a -> b
$ Int -> ((), Log Double) -> [((), Log Double)]
forall a. Int -> a -> [a]
replicate Int
n ((), Log Double
1 Log Double -> Log Double -> Log Double
forall a. Fractional a => a -> a -> a
/ Int -> Log Double
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 = (Int -> Population m ()
forall (m :: * -> *). Monad m => Int -> Population m ()
spawn Int
n Population m () -> Population m a -> Population m a
forall a b. Population m a -> Population m b -> Population m b
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 = m [(a, Log Double)] -> Population m a
forall (m :: * -> *) a.
Monad m =>
m [(a, Log Double)] -> Population m a
fromWeightedList (m [(a, Log Double)] -> Population m a)
-> m [(a, Log Double)] -> Population m a
forall a b. (a -> b) -> a -> b
$ do
  [(a, Log Double)]
pop <- Population m a -> m [(a, Log Double)]
forall (m :: * -> *) a. Population m a -> m [(a, Log Double)]
population Population m a
m
  let ([a]
xs, [Log Double]
ps) = [(a, Log Double)] -> ([a], [Log Double])
forall a b. [(a, b)] -> ([a], [b])
unzip [(a, Log Double)]
pop
  let n :: Int
n = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs
  let z :: Log Double
z = [Log Double] -> Log Double
forall a (f :: * -> *).
(RealFloat a, Foldable f) =>
f (Log a) -> Log a
Log.sum [Log Double]
ps
  if Log Double
z Log Double -> Log Double -> Bool
forall a. Ord a => a -> a -> Bool
> Log Double
0
    then do
      let weights :: Vector Double
weights = [Double] -> Vector Double
forall a. [a] -> Vector a
V.fromList ((Log Double -> Double) -> [Log Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double -> Double
forall a. Floating a => a -> a
exp (Double -> Double)
-> (Log Double -> Double) -> Log Double -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Log Double -> Double
forall a. Log a -> a
ln (Log Double -> Double)
-> (Log Double -> Log Double) -> Log Double -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Log Double -> Log Double -> Log Double
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 = [a] -> Vector a
forall a. [a] -> Vector a
V.fromList [a]
xs
      let offsprings :: [a]
offsprings = (Int -> a) -> [Int] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (Vector a
xvec Vector a -> Int -> a
forall a. Vector a -> Int -> a
V.!) [Int]
ancestors
      [(a, Log Double)] -> m [(a, Log Double)]
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ([(a, Log Double)] -> m [(a, Log Double)])
-> [(a, Log Double)] -> m [(a, Log Double)]
forall a b. (a -> b) -> a -> b
$ (a -> (a, Log Double)) -> [a] -> [(a, Log Double)]
forall a b. (a -> b) -> [a] -> [b]
map (,Log Double
z Log Double -> Log Double -> Log Double
forall a. Fractional a => a -> a -> a
/ Int -> Log Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n) [a]
offsprings
    else -- if all weights are zero do not resample
      [(a, Log Double)] -> m [(a, Log Double)]
forall a. a -> m a
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 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n) Int
0 Double
0 []
  where
    prob :: Int -> Double
prob Int
i = Vector Double
ps Vector Double -> Int -> Double
forall a. Vector a -> Int -> a
V.! Int
i
    n :: Int
n = Vector Double -> Int
forall a. Vector a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Vector Double
ps
    inc :: Double
inc = Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
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 Int -> Int -> Bool
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 Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
q
        then Int -> Double -> Int -> Double -> [Int] -> [Int]
f (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Double
v Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
inc) Int
j Double
q (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
acc)
        else Int -> Double -> Int -> Double -> [Int] -> [Int]
f Int
i Double
v (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Double
q Double -> Double -> Double
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 = (Vector Double -> m [Int]) -> Population m a -> Population m a
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) (Double -> [Int]) -> m Double -> m [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m Double
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 = Vector Double -> Int
forall a. Vector a -> Int
V.length Vector Double
weights
  Vector Double
dithers <- Int -> m Double -> m (Vector Double)
forall (m :: * -> *) a. Monad m => Int -> m a -> m (Vector a)
V.replicateM Int
bigN (Double -> Double -> m Double
forall (m :: * -> *).
MonadDistribution m =>
Double -> Double -> m Double
uniform Double
0.0 Double
1.0)
  let positions :: Vector Double
positions =
        (Double -> Double) -> Vector Double -> Vector Double
forall a b. (a -> b) -> Vector a -> Vector b
V.map (Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
bigN) (Vector Double -> Vector Double) -> Vector Double -> Vector Double
forall a b. (a -> b) -> a -> b
$
          (Double -> Double -> Double)
-> Vector Double -> Vector Double -> Vector Double
forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
V.zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(+) Vector Double
dithers ((Int -> Double) -> Vector Int -> Vector Double
forall a b. (a -> b) -> Vector a -> Vector b
V.map Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Vector Int -> Vector Double) -> Vector Int -> Vector Double
forall a b. (a -> b) -> a -> b
$ [Int] -> Vector Int
forall a. [a] -> Vector a
V.fromList [Int
0 .. Int
bigN Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1])
      cumulativeSum :: Vector Double
cumulativeSum = (Double -> Double -> Double)
-> Double -> Vector Double -> Vector Double
forall a b. (a -> b -> a) -> a -> Vector b -> Vector a
V.scanl Double -> Double -> Double
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 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
bigN =
            if (Vector Double
positions Vector Double -> Int -> Double
forall a. Vector a -> Int -> a
! Int
i) Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< (Vector Double
cumulativeSum Vector Double -> Int -> Double
forall a. Vector a -> Int -> a
! Int
j)
              then (Maybe Int, (Int, Int)) -> Maybe (Maybe Int, (Int, Int))
forall a. a -> Maybe a
Just (Int -> Maybe Int
forall a. a -> Maybe a
Just Int
j, (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, Int
j))
              else (Maybe Int, (Int, Int)) -> Maybe (Maybe Int, (Int, Int))
forall a. a -> Maybe a
Just (Maybe Int
forall a. Maybe a
Nothing, (Int
i, Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
        | Bool
otherwise =
            Maybe (Maybe Int, (Int, Int))
forall a. Maybe a
Nothing
  [Int] -> m [Int]
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ([Int] -> m [Int]) -> [Int] -> m [Int]
forall a b. (a -> b) -> a -> b
$ (Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (\Int
i -> Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ [Maybe Int] -> [Int]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe Int] -> [Int]) -> [Maybe Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ ((Int, Int) -> Maybe (Maybe Int, (Int, Int)))
-> (Int, Int) -> [Maybe Int]
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 = (Vector Double -> m [Int]) -> Population m a -> Population m a
forall (m :: * -> *) a.
MonadDistribution m =>
(Vector Double -> m [Int]) -> Population m a -> Population m a
resampleGeneric Vector Double -> m [Int]
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 = Int -> m Int -> m [Int]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM (Vector Double -> Int
forall a. Vector a -> Int
V.length Vector Double
ps) (Vector Double -> m Int
forall (v :: * -> *). Vector v Double => v Double -> m Int
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 = (Vector Double -> m [Int]) -> Population m a -> Population m a
forall (m :: * -> *) a.
MonadDistribution m =>
(Vector Double -> m [Int]) -> Population m a -> Population m a
resampleGeneric Vector Double -> m [Int]
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 = Weighted m [(a, Log Double)] -> Population (Weighted m) a
forall (m :: * -> *) a.
Monad m =>
m [(a, Log Double)] -> Population m a
fromWeightedList (Weighted m [(a, Log Double)] -> Population (Weighted m) a)
-> Weighted m [(a, Log Double)] -> Population (Weighted m) a
forall a b. (a -> b) -> a -> b
$ do
  [(a, Log Double)]
pop <- m [(a, Log Double)] -> Weighted m [(a, Log Double)]
forall (m :: * -> *) a. Monad m => m a -> Weighted m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m [(a, Log Double)] -> Weighted m [(a, Log Double)])
-> m [(a, Log Double)] -> Weighted m [(a, Log Double)]
forall a b. (a -> b) -> a -> b
$ Population m a -> m [(a, Log Double)]
forall (m :: * -> *) a. Population m a -> m [(a, Log Double)]
population Population m a
m
  let ([a]
xs, [Log Double]
ps) = [(a, Log Double)] -> ([a], [Log Double])
forall a b. [(a, b)] -> ([a], [b])
unzip [(a, Log Double)]
pop
  let z :: Log Double
z = [Log Double] -> Log Double
forall a (f :: * -> *).
(RealFloat a, Foldable f) =>
f (Log a) -> Log a
sum [Log Double]
ps
  let ws :: [Log Double]
ws = (Log Double -> Log Double) -> [Log Double] -> [Log Double]
forall a b. (a -> b) -> [a] -> [b]
map (if Log Double
z Log Double -> Log Double -> Bool
forall a. Ord a => a -> a -> Bool
> Log Double
0 then (Log Double -> Log Double -> Log Double
forall a. Fractional a => a -> a -> a
/ Log Double
z) else Log Double -> Log Double -> Log Double
forall a b. a -> b -> a
const (Log Double
1 Log Double -> Log Double -> Log Double
forall a. Fractional a => a -> a -> a
/ Int -> Log Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Log Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Log Double]
ps))) [Log Double]
ps
  Log Double -> Weighted m ()
forall (m :: * -> *). MonadFactor m => Log Double -> m ()
factor Log Double
z
  [(a, Log Double)] -> Weighted m [(a, Log Double)]
forall a. a -> Weighted m a
forall (m :: * -> *) a. Monad m => a -> m a
return ([(a, Log Double)] -> Weighted m [(a, Log Double)])
-> [(a, Log Double)] -> Weighted m [(a, Log Double)]
forall a b. (a -> b) -> a -> b
$ [a] -> [Log Double] -> [(a, Log Double)]
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 x. Weighted m x -> m x)
-> Population (Weighted m) a -> Population m a
forall (n :: * -> *) (m :: * -> *) a.
Monad n =>
(forall x. m x -> n x) -> Population m a -> Population n a
hoist Weighted m x -> m x
forall x. Weighted m x -> m x
forall (m :: * -> *) a. MonadFactor m => Weighted m a -> m a
applyWeight (Population (Weighted m) a -> Population m a)
-> (Population m a -> Population (Weighted m) a)
-> Population m a
-> Population m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Population m a -> Population (Weighted m) a
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 <- Population (Weighted m) a -> Weighted m [(a, Log Double)]
forall (m :: * -> *) a. Population m a -> m [(a, Log Double)]
population (Population (Weighted m) a -> Weighted m [(a, Log Double)])
-> Population (Weighted m) a -> Weighted m [(a, Log Double)]
forall a b. (a -> b) -> a -> b
$ Population m a -> Population (Weighted m) a
forall (m :: * -> *) a.
Monad m =>
Population m a -> Population (Weighted m) a
extractEvidence Population m a
m
  let ([a]
xs, [Log Double]
ps) = [(a, Log Double)] -> ([a], [Log Double])
forall a b. [(a, b)] -> ([a], [b])
unzip [(a, Log Double)]
pop
  Int
index <- Vector (Log Double) -> Weighted m Int
forall (v :: * -> *).
(Vector v (Log Double), Vector v Double) =>
v (Log Double) -> Weighted m Int
forall (m :: * -> *) (v :: * -> *).
(MonadDistribution m, Vector v (Log Double), Vector v Double) =>
v (Log Double) -> m Int
logCategorical (Vector (Log Double) -> Weighted m Int)
-> Vector (Log Double) -> Weighted m Int
forall a b. (a -> b) -> a -> b
$ [Log Double] -> Vector (Log Double)
forall a. [a] -> Vector a
V.fromList [Log Double]
ps
  let x :: a
x = [a]
xs [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
index
  a -> Weighted m a
forall a. a -> Weighted m a
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 = Weighted m [(a, Log Double)] -> m (Log Double)
forall (m :: * -> *) a. Functor m => Weighted m a -> m (Log Double)
extractWeight (Weighted m [(a, Log Double)] -> m (Log Double))
-> (Population m a -> Weighted m [(a, Log Double)])
-> Population m a
-> m (Log Double)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Population (Weighted m) a -> Weighted m [(a, Log Double)]
forall (m :: * -> *) a. Population m a -> m [(a, Log Double)]
population (Population (Weighted m) a -> Weighted m [(a, Log Double)])
-> (Population m a -> Population (Weighted m) a)
-> Population m a
-> Weighted m [(a, Log Double)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Population m a -> Population (Weighted m) a
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 = Weighted m a -> m a
forall (m :: * -> *) a. MonadFactor m => Weighted m a -> m a
applyWeight (Weighted m a -> m a)
-> (Population m a -> Weighted m a) -> Population m a -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Population m a -> Weighted m a
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 <- Population m a -> m [(a, Double)]
forall (m :: * -> *) a.
Functor m =>
Population m a -> m [(a, Double)]
explicitPopulation Population m a
p
  let ys :: [Double]
ys = ((a, Double) -> Double) -> [(a, Double)] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (\(a
x, Double
w) -> a -> Double
f a
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
w) [(a, Double)]
xs
  let t :: Double
t = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
Data.List.sum [Double]
ys
  Double -> m Double
forall a. a -> m a
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 = n [(a, Log Double)] -> Population n a
forall (m :: * -> *) a.
Monad m =>
m [(a, Log Double)] -> Population m a
fromWeightedList (n [(a, Log Double)] -> Population n a)
-> (Population m a -> n [(a, Log Double)])
-> Population m a
-> Population n a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. m [(a, Log Double)] -> n [(a, Log Double)]
forall x. m x -> n x
f (m [(a, Log Double)] -> n [(a, Log Double)])
-> (Population m a -> m [(a, Log Double)])
-> Population m a
-> n [(a, Log Double)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Population m a -> m [(a, Log Double)]
forall (m :: * -> *) a. Population m a -> m [(a, Log Double)]
population