monad-bayes-0.1.1.0: test/TestPopulation.hs

module TestPopulation where

import Control.Monad.Bayes.Class
import Control.Monad.Bayes.Enumerator
import Control.Monad.Bayes.Population as Population
import Control.Monad.Bayes.Sampler
import Data.AEq
import Sprinkler

weightedSampleSize :: MonadSample m => Population m a -> m Int
weightedSampleSize = fmap length . runPopulation

popSize :: IO Int
popSize = sampleIOfixed $ weightedSampleSize $ spawn 5 >> sprinkler

manySize :: IO Int
manySize = sampleIOfixed $ weightedSampleSize $ spawn 5 >> sprinkler >> spawn 3

sprinkler :: MonadInfer m => m Bool
sprinkler = Sprinkler.soft

sprinklerExact :: [(Bool, Double)]
sprinklerExact = enumerate Sprinkler.soft

--all_check = (mass (Population.all id (spawn 2 >> sprinkler)) True) ~== 0.09

transCheck1 :: Bool
transCheck1 =
  enumerate (collapse sprinkler)
    ~== sprinklerExact

transCheck2 :: Bool
transCheck2 =
  enumerate (collapse (spawn 2 >> sprinkler))
    ~== sprinklerExact

resampleCheck :: Int -> Bool
resampleCheck n =
  (enumerate . collapse . resampleMultinomial) (spawn n >> sprinkler)
    ~== sprinklerExact

popAvgCheck :: Bool
popAvgCheck = expectation f Sprinkler.soft ~== expectation id (popAvg f $ pushEvidence Sprinkler.soft)
  where
    f True = 10
    f False = 4