From 358b7bd777eaea7a4cf1b3e8461d3328ba7d542c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Manuel=20B=C3=A4renz?= <programming@manuelbaerenz.de>
Date: Wed, 10 May 2023 12:17:52 +0200
Subject: [PATCH] Add effective sample size to Population

---
 src/Control/Monad/Bayes/Population.hs | 27 +++++++++++++++++++++++++++
 1 file changed, 27 insertions(+)

diff --git a/src/Control/Monad/Bayes/Population.hs b/src/Control/Monad/Bayes/Population.hs
index d937f67a..301a00b9 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,31 @@ 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 ::