diff --git a/src/Control/Monad/Bayes/Population.hs b/src/Control/Monad/Bayes/Population.hs index be670df2..d97201a3 100644 --- a/src/Control/Monad/Bayes/Population.hs +++ b/src/Control/Monad/Bayes/Population.hs @@ -27,6 +27,8 @@ module Control.Monad.Bayes.Population resampleSystematic, stratified, resampleStratified, + onlyBelowEffectiveSampleSize, + effectiveSampleSize, extractEvidence, pushEvidence, proper, @@ -210,6 +212,32 @@ resampleMultinomial :: Population m a resampleMultinomial = resampleGeneric multinomial +-- | Only use the given resampler when the effective sample size is below a certain threshold +onlyBelowEffectiveSampleSize :: + MonadDistribution m => + -- | The threshold under which the effective sample size must fall before the resampler is used. + -- For example, this may be half of the number of particles. + Double -> + -- | The resampler to user under the threshold + (MonadDistribution m => Population m a -> Population m a) -> + -- | The new resampler + (Population m a -> Population m a) +onlyBelowEffectiveSampleSize threshold resampler pop = do + ess <- lift $ effectiveSampleSize pop + if ess < threshold then resampler pop else pop + +{- | Compute the effective sample size of a population from the weights. + + See https://en.wikipedia.org/wiki/Design_effect#Effective_sample_size +-} +effectiveSampleSize :: Functor m => Population m a -> m Double +effectiveSampleSize = fmap (effectiveSampleSizeKish . map (exp . ln . snd)) . runPopulation + where + effectiveSampleSizeKish :: [Double] -> Double + effectiveSampleSizeKish weights = square (Data.List.sum weights) / Data.List.sum (square <$> weights) + square :: Double -> Double + square x = x * x + -- | Separate the sum of weights into the 'Weighted' transformer. -- Weights are normalized after this operation. extractEvidence ::