Add effective sample size to Population

Author: turion

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 ::