:e ImportQualifiedPost
:e FlexibleContexts
:e BlockArguments
:e TupleSections
:e FlexibleContexts
:e OverloadedStrings
:e LambdaCase

import Control.Monad.Bayes.Class
import Control.Monad.Bayes.Enumerator
import Control.Monad.Bayes.Weighted
import Control.Monad.Bayes.Sampler.Strict
import Control.Monad.Bayes.Density.Free
import Control.Monad.Bayes.Population
import Control.Monad.Bayes.Sequential.Coroutine
import Control.Monad.Bayes.Inference.SMC
import Control.Monad.Bayes.Inference.RMSMC
import Control.Monad.Bayes.Inference.PMMH
import Control.Monad.Bayes.Inference.MCMC
import qualified Data.Text as T
import Numeric.Log
import Control.Arrow (first,second)
import Control.Monad
import Control.Applicative
import Control.Monad.Bayes.Class 
import Control.Monad.Bayes.Traced
import Control.Monad.Bayes.Weighted
import Graphics.Vega.VegaLite hiding (density)
import qualified Graphics.Vega.VegaLite as VL
import IHaskell.Display.Hvega (vlShow)

import qualified Pipes.Prelude as P
import qualified Pipes as P
import Pipes (Producer, (>->), MonadTrans (lift))

:l ../plotting.hs

Combining SMC and MCMC

Monad-bayes offers several more advanced inference methods, which are modular combinations of SMC and MCMC.

Sampling, MCMC and SMC notebooks describe those methods in more detail.

Resample Move Sequential Monte Carlo (RMSMC)

RMSMC is fundamentally an SMC technique. It creates and updates a population of weighted samples. The clever part is that after resampling, the update step uses MCMC to perform a walk on each particle, which updates the population efficiently.

To motivate this more sophisticated inference method, let's pick a relatively hard inference problem: inferring the position of a moving point mass from measurements of its bearings.

Here it is in practice:

-- calculate the bearing (the angle) of a point
bearing (x, y) = atan2 y x

prior :: MonadDistribution m => Producer (Double,Double) m ()
prior = do
    let velocity_var = 1e-5
    initialPosition <- lift $ liftA2 (,) (normal 0.01 0.01) (normal 0.95 0.01)
    initialVelocity <- lift $ liftA2 (,) (normal 0.002 0.01) (normal (-0.013) 0.01)
    -- step is performed for a unit of time
    let step ((qx,qy), (vx,vy)) = do
        vxNew <- normal vx (sqrt velocity_var)
        vyNew <- normal vy (sqrt velocity_var)
        return ((qx + vx, qy + vy), (vxNew, vyNew))
    P.unfoldr (\x -> Right <$> dup (step x)) (initialPosition, initialVelocity) >-> P.map fst
    
    where 
    dup :: Monad m => m b -> m (b, b)
    dup mx = do
        x <- mx
        return (x,x)

measurementNoise = 0.005

-- observationModel :: MonadDistribution m => P.Pipe (Double, Double) Double m ()
observationModel position = normal (bearing position) measurementNoise


numPoints = 40

model :: MonadMeasure m => [Double] -> Producer (Double, Double) m ()
model observations = zipWithM likelihood prior observationStream
  where  
    observationStream = P.each observations
    measurement_noise = 0.005
    likelihood (pos, obs) = do
        o <- observationModel pos 
        factor $ normalPdf o measurement_noise obs
    zipWithM f p1 p2 = P.zip p1 p2 >-> P.chain f >-> P.map fst

positions <- sampleIOfixed $ P.toListM (prior >-> P.take numPoints)

plot ( zip positions $ Prelude.repeat $ T.pack "prior")

observations <- sampleIOfixed $ mapM observationModel positions

sMCMC <- sampleIOfixed $ unweighted $ mcmc MCMCConfig {numMCMCSteps = 10000, numBurnIn = 1000, proposal = SingleSiteMH} $ 
    P.toListM $ model observations

MCMC inference arrived returned a path which is divergent in the begining (starting at point (0,0.95)) from prior path.

plot $
  zip positions (Prelude.repeat $ T.pack "prior") <>  
  zip (head sMCMC) (Prelude.repeat $ T.pack "inferred: MCMC")

sSMC <- sampleIOfixed $ runPopulation $ smc 
    SMCConfig {numSteps = numPoints, numParticles = 10000, resampler = resampleSystematic} 
    $ P.toListM $ model observations

SMC on the ther hand provided many paths close to the prior at the path's beginning and diverging as the path progresses.

a = (\(x,d) -> zip x (Prelude.repeat $ T.pack "inferred: SMC")) =<< sSMC

plot $ zip positions (Prelude.repeat $ T.pack "prior") <> take 10000 a

sRMSMC <- sampleIOfixed $ runPopulation $ rmsmcBasic
    MCMCConfig {numMCMCSteps = 50, numBurnIn = 0, proposal = SingleSiteMH}
    SMCConfig {numSteps = numPoints, numParticles = 10, resampler = resampleSystematic} 
    $ P.toListM $ model observations

a = (\(x,d) -> zip x (Prelude.repeat $ T.pack "inferred: RMSMC")) =<< sRMSMC

plot $ zip positions (Prelude.repeat $ T.pack "prior") <> take 10000 a

The good results are if the inferred distribution is close to prior, as observed samples are only modified by a small noise from prior.

And it works! Look how few particles it needed.

Particle Marginal Metropolis Hastings (PMMH)

PMMH is fundamentally an MCMC technique. It performs a random walk through parameter space, but to estimate the likelihood at each step, it uses an unbiased estimate from a population of samples. We'll reuse the regression with outliers example from tutorials/MCMC.ipynb.

paramPrior = do
    slope <- normal 0 2
    intercept <- normal 0 2
    noise <- gamma 7 7
    prob_outlier <- uniform 0 0.5 
    return (slope, intercept, noise, prob_outlier)

forward (slope, intercept, noise, probOutlier) x = do
    isOutlier <- bernoulli probOutlier
    let meanParams = if isOutlier
                    then (0, 20)
                    else (x*slope + intercept, sqrt noise)
    return (meanParams, isOutlier)

regressionWithOutliersData :: (MonadDistribution m, Traversable t) => t Double -> m (t ((Double, Double), Bool))
regressionWithOutliersData xs = do
    params <- paramPrior

    forM xs \x -> do
        ((mu, std), isOutlier) <- forward params x
        y <- normal mu std
        return ((x, y), isOutlier)
        
range = [-10,-9.9..10] :: [Double]
samples <- sampleIOfixed $ regressionWithOutliersData range
plot (fmap (second (T.pack . show)) samples)

regressionWithOutliers xs ys params = do
    
    outliers <- forM (zip xs ys) \(x, y) -> do
        ((mu, std), isOutlier) <- forward params x
        factor $ normalPdf mu std y
        return isOutlier
    return (params, outliers)

mhRuns <- sampleIOfixed $ pmmh 
    MCMCConfig {numMCMCSteps = 20, numBurnIn = 10, proposal = SingleSiteMH}
    SMCConfig {numSteps = 200, numParticles = 100, resampler = resampleSystematic}
    paramPrior
    (regressionWithOutliers range (snd . fst <$> samples))

m = (\((_,a), b) -> (a,b)) <$> head mhRuns

outlierProb s = (\(x, y) ->   ln (exp (x / (x+y))) )
        <$> foldr
              (\(lb, d) li -> [ if b then (num1+d, num2) else (num1,num2+d) | (b,(num1, num2)) <- zip lb li])
              (Prelude.repeat (0 :: Log Double,0 :: Log Double)) 
              s

plot $ take 1000 (zip (fst <$> samples) (outlierProb m))

As the above plot shows, this works nicely: the slope, intercept, noise and prob_outlier variables are inferred by a random walk through the space, while the score to determine whether to accept a new proposed step in this walk is determined by a particle filter which guesses which points are outliers after each observation.

$SMC^2$

The culmination of monad-bayes is $SMC^2$, which combines all the previous piece together into one inference algorithm.

It is $RMSMC$, but with the $MCMC$ walk performed using $PMMH$.