A simple embedded probabilistic programming language

Prob.hs

{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE LambdaCase #-}

import Control.Monad
import Control.Monad.Free
import qualified System.Random.MWC.Probability as MWC

data ModelF r =
    BernoulliF Double (Bool -> r)
  | BetaF Double Double (Double -> r)
  deriving Functor

type Model = Free ModelF

-- terms

bernoulli :: Double -> Model Bool
bernoulli p = liftF (BernoulliF p id)

beta :: Double -> Double -> Model Double
beta a b = liftF (BetaF a b id)

binomial :: Int -> Double -> Model Int
binomial n p = fmap count coins where
  coins = replicateM n (bernoulli p)

uniform :: Model Double
uniform = beta 1 1

betaBinomial :: Int -> Double -> Double -> Model Int
betaBinomial n a b = do
  p <- beta a b
  binomial n p

-- simulation

toSampler :: Model a -> MWC.Prob IO a
toSampler = iterM $ \case
  BernoulliF p f -> MWC.bernoulli p >>= f
  BetaF a b f    -> MWC.beta a b >>= f

simulate :: Model a -> IO a
simulate model = MWC.withSystemRandom . MWC.asGenIO $
  MWC.sample (toSampler model)

-- conditioning

invert :: (Monad m, Eq b) => m a -> (a -> m b) -> [b] -> m a
invert proposal model observed = loop where
  len  = length observed
  loop = do
    parameters <- proposal
    generated  <- replicateM len (model parameters)
    if   generated == observed
    then return parameters
    else loop

invertWithAssistance
  :: (Monad m, Eq c) => ([a] -> c) -> m b -> (b -> m a) -> [a] -> m b
invertWithAssistance assister proposal model observed = loop where
  len  = length observed
  loop = do
    parameters <- proposal
    generated  <- replicateM len (model parameters)
    if   assister generated == assister observed
    then return parameters
    else loop

-- example

posterior :: Model Double
posterior = invert uniform bernoulli [True, True, False, True]

posterior0 :: Model Double
posterior0 = invertWithAssistance count uniform bernoulli obs where
  obs =
    [True, True, True, False, True, True, False, True, True, True, True, False]

count :: [Bool] -> Int
count = length . filter id