-- |
-- Module      : Control.Monad.Bayes.Inference.RMSMC
-- Description : Resample-Move Sequential Monte Carlo (RM-SMC)
-- Copyright   : (c) Adam Scibior, 2015-2020
-- License     : MIT
-- Maintainer  : leonhard.markert@tweag.io
-- Stability   : experimental
-- Portability : GHC
--
-- Resample-move Sequential Monte Carlo (RM-SMC) sampling.
--
-- Walter Gilks and Carlo Berzuini. 2001. Following a moving target - Monte Carlo inference for dynamic Bayesian models. /Journal of the Royal Statistical Society/ 63 (2001), 127-146. <http://www.mathcs.emory.edu/~whalen/Papers/BNs/MonteCarlo-DBNs.pdf>
module Control.Monad.Bayes.Inference.RMSMC
  ( rmsmc,
    rmsmcLocal,
    rmsmcBasic,
  )
where

import Control.Monad.Bayes.Class
import Control.Monad.Bayes.Helpers
import Control.Monad.Bayes.Population
import Control.Monad.Bayes.Sequential as Seq
import Control.Monad.Bayes.Traced as Tr
import qualified Control.Monad.Bayes.Traced.Basic as TrBas
import qualified Control.Monad.Bayes.Traced.Dynamic as TrDyn

-- | Resample-move Sequential Monte Carlo.
rmsmc ::
  MonadSample m =>
  -- | number of timesteps
  Int ->
  -- | number of particles
  Int ->
  -- | number of Metropolis-Hastings transitions after each resampling
  Int ->
  -- | model
  Sequential (Traced (Population m)) a ->
  Population m a
rmsmc :: Int
-> Int
-> Int
-> Sequential (Traced (Population m)) a
-> Population m a
rmsmc k :: Int
k n :: Int
n t :: Int
t =
  Traced (Population m) a -> Population m a
forall (m :: * -> *) a. Monad m => Traced m a -> m a
marginal
    (Traced (Population m) a -> Population m a)
-> (Sequential (Traced (Population m)) a
    -> Traced (Population m) a)
-> Sequential (Traced (Population m)) a
-> Population m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall x. Traced (Population m) x -> Traced (Population m) x)
-> Int
-> Sequential (Traced (Population m)) a
-> Traced (Population m) a
forall (m :: * -> *) a.
Monad m =>
(forall x. m x -> m x) -> Int -> Sequential m a -> m a
sis (Int
-> (Traced (Population m) x -> Traced (Population m) x)
-> Traced (Population m) x
-> Traced (Population m) x
forall a. Int -> (a -> a) -> a -> a
composeCopies Int
t Traced (Population m) x -> Traced (Population m) x
forall (m :: * -> *) a. MonadSample m => Traced m a -> Traced m a
mhStep (Traced (Population m) x -> Traced (Population m) x)
-> (Traced (Population m) x -> Traced (Population m) x)
-> Traced (Population m) x
-> Traced (Population m) x
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall x. Population m x -> Population m x)
-> Traced (Population m) x -> Traced (Population m) x
forall (m :: * -> *) a.
(forall x. m x -> m x) -> Traced m a -> Traced m a
hoistT forall x. Population m x -> Population m x
forall (m :: * -> *) a.
MonadSample m =>
Population m a -> Population m a
resampleSystematic) Int
k
    (Sequential (Traced (Population m)) a -> Traced (Population m) a)
-> (Sequential (Traced (Population m)) a
    -> Sequential (Traced (Population m)) a)
-> Sequential (Traced (Population m)) a
-> Traced (Population m) a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall x. Traced (Population m) x -> Traced (Population m) x)
-> Sequential (Traced (Population m)) a
-> Sequential (Traced (Population m)) a
forall (m :: * -> *) a. (forall x. m x -> m x) -> S m a -> S m a
hoistS ((forall x. Population m x -> Population m x)
-> Traced (Population m) x -> Traced (Population m) x
forall (m :: * -> *) a.
(forall x. m x -> m x) -> Traced m a -> Traced m a
hoistT (Int -> Population m ()
forall (m :: * -> *). Monad m => Int -> Population m ()
spawn Int
n Population m () -> Population m x -> Population m x
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>>))

-- | Resample-move Sequential Monte Carlo with a more efficient
-- tracing representation.
rmsmcBasic ::
  MonadSample m =>
  -- | number of timesteps
  Int ->
  -- | number of particles
  Int ->
  -- | number of Metropolis-Hastings transitions after each resampling
  Int ->
  -- | model
  Sequential (TrBas.Traced (Population m)) a ->
  Population m a
rmsmcBasic :: Int
-> Int
-> Int
-> Sequential (Traced (Population m)) a
-> Population m a
rmsmcBasic k :: Int
k n :: Int
n t :: Int
t =
  Traced (Population m) a -> Population m a
forall (m :: * -> *) a. Monad m => Traced m a -> m a
TrBas.marginal
    (Traced (Population m) a -> Population m a)
-> (Sequential (Traced (Population m)) a
    -> Traced (Population m) a)
-> Sequential (Traced (Population m)) a
-> Population m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall x. Traced (Population m) x -> Traced (Population m) x)
-> Int
-> Sequential (Traced (Population m)) a
-> Traced (Population m) a
forall (m :: * -> *) a.
Monad m =>
(forall x. m x -> m x) -> Int -> Sequential m a -> m a
sis (Int
-> (Traced (Population m) x -> Traced (Population m) x)
-> Traced (Population m) x
-> Traced (Population m) x
forall a. Int -> (a -> a) -> a -> a
composeCopies Int
t Traced (Population m) x -> Traced (Population m) x
forall (m :: * -> *) a. MonadSample m => Traced m a -> Traced m a
TrBas.mhStep (Traced (Population m) x -> Traced (Population m) x)
-> (Traced (Population m) x -> Traced (Population m) x)
-> Traced (Population m) x
-> Traced (Population m) x
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall x. Population m x -> Population m x)
-> Traced (Population m) x -> Traced (Population m) x
forall (m :: * -> *) a.
(forall x. m x -> m x) -> Traced m a -> Traced m a
TrBas.hoistT forall x. Population m x -> Population m x
forall (m :: * -> *) a.
MonadSample m =>
Population m a -> Population m a
resampleSystematic) Int
k
    (Sequential (Traced (Population m)) a -> Traced (Population m) a)
-> (Sequential (Traced (Population m)) a
    -> Sequential (Traced (Population m)) a)
-> Sequential (Traced (Population m)) a
-> Traced (Population m) a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall x. Traced (Population m) x -> Traced (Population m) x)
-> Sequential (Traced (Population m)) a
-> Sequential (Traced (Population m)) a
forall (m :: * -> *) a. (forall x. m x -> m x) -> S m a -> S m a
hoistS ((forall x. Population m x -> Population m x)
-> Traced (Population m) x -> Traced (Population m) x
forall (m :: * -> *) a.
(forall x. m x -> m x) -> Traced m a -> Traced m a
TrBas.hoistT (Int -> Population m ()
forall (m :: * -> *). Monad m => Int -> Population m ()
spawn Int
n Population m () -> Population m x -> Population m x
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>>))

-- | A variant of resample-move Sequential Monte Carlo
-- where only random variables since last resampling are considered
-- for rejuvenation.
rmsmcLocal ::
  MonadSample m =>
  -- | number of timesteps
  Int ->
  -- | number of particles
  Int ->
  -- | number of Metropolis-Hastings transitions after each resampling
  Int ->
  -- | model
  Sequential (TrDyn.Traced (Population m)) a ->
  Population m a
rmsmcLocal :: Int
-> Int
-> Int
-> Sequential (Traced (Population m)) a
-> Population m a
rmsmcLocal k :: Int
k n :: Int
n t :: Int
t =
  Traced (Population m) a -> Population m a
forall (m :: * -> *) a. Monad m => Traced m a -> m a
TrDyn.marginal
    (Traced (Population m) a -> Population m a)
-> (Sequential (Traced (Population m)) a
    -> Traced (Population m) a)
-> Sequential (Traced (Population m)) a
-> Population m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall x. Traced (Population m) x -> Traced (Population m) x)
-> Int
-> Sequential (Traced (Population m)) a
-> Traced (Population m) a
forall (m :: * -> *) a.
Monad m =>
(forall x. m x -> m x) -> Int -> Sequential m a -> m a
sis (Traced (Population m) x -> Traced (Population m) x
forall (m :: * -> *) a. Monad m => Traced m a -> Traced m a
TrDyn.freeze (Traced (Population m) x -> Traced (Population m) x)
-> (Traced (Population m) x -> Traced (Population m) x)
-> Traced (Population m) x
-> Traced (Population m) x
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int
-> (Traced (Population m) x -> Traced (Population m) x)
-> Traced (Population m) x
-> Traced (Population m) x
forall a. Int -> (a -> a) -> a -> a
composeCopies Int
t Traced (Population m) x -> Traced (Population m) x
forall (m :: * -> *) a. MonadSample m => Traced m a -> Traced m a
TrDyn.mhStep (Traced (Population m) x -> Traced (Population m) x)
-> (Traced (Population m) x -> Traced (Population m) x)
-> Traced (Population m) x
-> Traced (Population m) x
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall x. Population m x -> Population m x)
-> Traced (Population m) x -> Traced (Population m) x
forall (m :: * -> *) a.
(forall x. m x -> m x) -> Traced m a -> Traced m a
TrDyn.hoistT forall x. Population m x -> Population m x
forall (m :: * -> *) a.
MonadSample m =>
Population m a -> Population m a
resampleSystematic) Int
k
    (Sequential (Traced (Population m)) a -> Traced (Population m) a)
-> (Sequential (Traced (Population m)) a
    -> Sequential (Traced (Population m)) a)
-> Sequential (Traced (Population m)) a
-> Traced (Population m) a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall x. Traced (Population m) x -> Traced (Population m) x)
-> Sequential (Traced (Population m)) a
-> Sequential (Traced (Population m)) a
forall (m :: * -> *) a. (forall x. m x -> m x) -> S m a -> S m a
hoistS ((forall x. Population m x -> Population m x)
-> Traced (Population m) x -> Traced (Population m) x
forall (m :: * -> *) a.
(forall x. m x -> m x) -> Traced m a -> Traced m a
TrDyn.hoistT (Int -> Population m ()
forall (m :: * -> *). Monad m => Int -> Population m ()
spawn Int
n Population m () -> Population m x -> Population m x
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>>))

-- | Apply a function a given number of times.
composeCopies :: Int -> (a -> a) -> (a -> a)
composeCopies :: Int -> (a -> a) -> a -> a
composeCopies k :: Int
k f :: a -> a
f = ((a -> a) -> (a -> a) -> a -> a) -> (a -> a) -> [a -> a] -> a -> a
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (a -> a) -> (a -> a) -> a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
(.) a -> a
forall a. a -> a
id (Int -> (a -> a) -> [a -> a]
forall a. Int -> a -> [a]
replicate Int
k a -> a
f)