From 7811d676f7024ae8e3571ccb2a9e8e35b2adbe21 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 2a384f77..bab81a9d 100644
--- a/src/Control/Monad/Bayes/Population.hs
+++ b/src/Control/Monad/Bayes/Population.hs
@@ -26,6 +26,8 @@ module Control.Monad.Bayes.Population
     resampleSystematic,
     stratified,
     resampleStratified,
+    onlyBelowEffectiveSampleSize,
+    effectiveSampleSize,
     extractEvidence,
     pushEvidence,
     proper,
@@ -206,6 +208,31 @@ resampleMultinomial ::
   PopulationT 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) => PopulationT m a -> PopulationT m a) ->
+  -- | The new resampler
+  (PopulationT m a -> PopulationT 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) => PopulationT m a -> m Double
+effectiveSampleSize = fmap (effectiveSampleSizeKish . map (exp . ln . snd)) . runPopulationT
+  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 'WeightedT' transformer.
 -- Weights are normalized after this operation.
 extractEvidence ::