From ecbff57fea39f327deed033e677144ad5f375fd6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Manuel=20B=C3=A4renz?= <programming@manuelbaerenz.de>
Date: Thu, 9 Mar 2023 14:58:42 +0100
Subject: [PATCH 1/9] WIP

---
 rhine-bayes/app/Main.hs                       | 28 +++++++++++++++
 .../src/Data/MonadicStreamFunction/Bayes.hs   | 34 +++++++++++++++++++
 rhine-bayes/src/FRP/Rhine/Bayes.hs            | 18 ++++++++++
 3 files changed, 80 insertions(+)

diff --git a/rhine-bayes/app/Main.hs b/rhine-bayes/app/Main.hs
index 001f35c4..7704d502 100644
--- a/rhine-bayes/app/Main.hs
+++ b/rhine-bayes/app/Main.hs
@@ -234,6 +234,7 @@ mains :: [(String, IO ())]
 mains =
   [ ("single rate", mainSingleRate)
   , ("multi rate, temperature process", mainMultiRate)
+  , ("multi rate, temperature process, RMSMC", mainMultiRateRMSMC)
   ]
 
 main :: IO ()
@@ -321,6 +322,17 @@ inference = hoistClSF sampleIOGloss inferenceBehaviour @@ liftClock Busy
     particles <- runPopulationCl nParticles resampleSystematic posteriorTemperatureProcess -< measured
     returnA -< Result{temperature, measured, latent, particles}
 
+{- | This part performs the inference (and passes along temperature, sensor and position simulations).
+   It runs as fast as possible, so this will potentially drain the CPU.
+-}
+inferenceRMSMC :: Rhine (GlossConcT IO) (LiftClock IO GlossConcT Busy) (Temperature, (Sensor, Pos)) Result
+inferenceRMSMC = hoistClSF sampleIOGloss inferenceBehaviour @@ liftClock Busy
+ where
+  inferenceBehaviour :: (MonadDistribution m, Diff td ~ Double, MonadIO m) => BehaviourF m td (Temperature, (Sensor, Pos)) Result
+  inferenceBehaviour = proc (temperature, (measured, latent)) -> do
+    particles <- resampleMoveSequentialMonteCarloCl nParticles 10 resampleSystematic posteriorTemperatureProcess -< measured
+    returnA -< Result{temperature, measured, latent, particles}
+
 -- | Visualize the current 'Result' at a rate controlled by the @gloss@ backend, usually 30 FPS.
 visualisationRhine :: Rhine (GlossConcT IO) (GlossClockUTC GlossSimClockIO) Result ()
 visualisationRhine = hoistClSF sampleIOGloss visualisation @@ glossClockUTC GlossSimClockIO
@@ -336,12 +348,28 @@ mainRhineMultiRate =
             >-- keepLast Result{temperature = initialTemperature, measured = zeroVector, latent = zeroVector, particles = []} -@- glossConcurrently -->
               visualisationRhine
 
+mainRhineMultiRateRMSMC =
+  userTemperature
+    @@ glossClockUTC GlossEventClockIO
+      >-- keepLast initialTemperature -@- glossConcurrently -->
+        modelRhine
+        >-- keepLast (initialTemperature, (zeroVector, zeroVector)) -@- glossConcurrently -->
+          inferenceRMSMC
+            >-- keepLast Result{temperature = initialTemperature, measured = zeroVector, latent = zeroVector, particles = []} -@- glossConcurrently -->
+              visualisationRhine
+
 mainMultiRate :: IO ()
 mainMultiRate =
   void $
     launchGlossThread glossSettings $
       flow mainRhineMultiRate
 
+mainMultiRateRMSMC :: IO ()
+mainMultiRateRMSMC =
+  void $
+    launchGlossThread glossSettings $
+      flow mainRhineMultiRateRMSMC
+
 -- * Utilities
 
 instance MonadDistribution m => MonadDistribution (GlossConcT m) where
diff --git a/rhine-bayes/src/Data/MonadicStreamFunction/Bayes.hs b/rhine-bayes/src/Data/MonadicStreamFunction/Bayes.hs
index a33149c8..dc8bebe3 100644
--- a/rhine-bayes/src/Data/MonadicStreamFunction/Bayes.hs
+++ b/rhine-bayes/src/Data/MonadicStreamFunction/Bayes.hs
@@ -12,11 +12,15 @@ import Numeric.Log hiding (sum)
 
 -- monad-bayes
 import Control.Monad.Bayes.Population
+import Control.Monad.Bayes.Traced.Class
 
 -- dunai
 import Data.MonadicStreamFunction
 import Data.MonadicStreamFunction.InternalCore (MSF (..))
+import Control.Monad.Trans.Class
+import Control.Monad.Bayes.Class (MonadDistribution)
 
+-- FIXME rename to sequentialMonteCarlo or smc?
 -- | Run the Sequential Monte Carlo algorithm continuously on an 'MSF'
 runPopulationS ::
   forall m a b.
@@ -48,6 +52,36 @@ runPopulationsS resampler = go
         unzip $
           (swap . fmap fst &&& swap . fmap snd) . swap <$> bAndMSFs
 
+resampleMoveSequentialMonteCarlo ::
+  forall t m a b.
+  (MonadDistribution m, HasTraced t, MonadTrans t) =>
+  -- | Number of particles
+  Int ->
+  -- | Number of MC steps
+  Int ->
+  -- | Resampler
+  (forall x. Population m x -> Population m x) ->
+  MSF (t (Population m)) a b ->
+  -- FIXME Why not MSF m a (Population b)
+  MSF m a [(b, Log Double)]
+resampleMoveSequentialMonteCarlo nParticles nMC resampler = go . (spawn nParticles $>)
+  where
+    go ::
+      Monad m =>
+      Population m (MSF (t (Population m)) a b) ->
+      MSF m a [(b, Log Double)]
+    go msfs = MSF $ \a -> do
+      -- TODO This is quite different than the dunai version now. Maybe it's right nevertheless.
+      -- FIXME This normalizes, which introduces bias, whatever that means
+      bAndMSFs <- runPopulation $ normalize $ marginal $ freeze $ composeCopies nMC mhStep $ hoistTrace resampler $ flip unMSF a =<< lift msfs
+      return $
+        second (go . fromWeightedList . return) $
+          unzip $
+            (swap . fmap fst &&& swap . fmap snd) . swap <$> bAndMSFs
+
+
+-- resampleMoveSequentialMonteCarlo nParticles nMC resampler = morphS marginal $ runPopulationS nParticles $ freeze . composeCopies nMC mhStep . hoistTrace resampler
+
 -- FIXME see PR re-adding this to monad-bayes
 normalize :: Monad m => Population m a -> Population m a
 normalize = fromWeightedList . fmap (\particles -> second (/ (sum $ snd <$> particles)) <$> particles) . runPopulation
diff --git a/rhine-bayes/src/FRP/Rhine/Bayes.hs b/rhine-bayes/src/FRP/Rhine/Bayes.hs
index 5422a0c3..e6200bed 100644
--- a/rhine-bayes/src/FRP/Rhine/Bayes.hs
+++ b/rhine-bayes/src/FRP/Rhine/Bayes.hs
@@ -15,6 +15,8 @@ import qualified Data.MonadicStreamFunction.Bayes as DunaiBayes
 
 -- rhine
 import FRP.Rhine
+import Control.Monad.Bayes.Traced.Class (HasTraced)
+import Control.Monad.Trans.Class (MonadTrans)
 
 -- * Inference methods
 
@@ -32,6 +34,22 @@ runPopulationCl :: forall m cl a b . Monad m =>
   -> ClSF m cl a [(b, Log Double)]
 runPopulationCl nParticles resampler = DunaiReader.readerS . DunaiBayes.runPopulationS nParticles resampler . DunaiReader.runReaderS
 
+-- | Run the Resample Move Sequential Monte Carlo algorithm continuously on a 'ClSF'.
+resampleMoveSequentialMonteCarloCl :: forall t m cl a b . (MonadDistribution m, HasTraced t, MonadTrans t) =>
+  -- | Number of particles
+  Int ->
+  -- | Number of MC steps
+  Int ->
+  -- | Resampler (see 'Control.Monad.Bayes.Population' for some standard choices)
+  (forall x . Population m x -> Population m x)
+  -- | A signal function modelling the stochastic process on which to perform inference.
+  --   @a@ represents observations upon which the model should condition, using e.g. 'score'.
+  --   It can also additionally contain hyperparameters.
+  --   @b@ is the type of estimated current state.
+  -> ClSF (t (Population m)) cl a b
+  -> ClSF m cl a [(b, Log Double)]
+resampleMoveSequentialMonteCarloCl nParticles nMC resampler = DunaiReader.readerS . DunaiBayes.resampleMoveSequentialMonteCarlo nParticles nMC resampler . DunaiReader.runReaderS
+
 -- * Short standard library of stochastic processes
 
 -- | White noise, that is, an independent normal distribution at every time step.

From 73d2a4fe3300c0e012931778227a320dd750a6fe Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Manuel=20B=C3=A4renz?= <programming@manuelbaerenz.de>
Date: Tue, 2 May 2023 10:39:27 +0200
Subject: [PATCH 2/9] Only use 2 particles for debugging

---
 rhine-bayes/app/Main.hs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/rhine-bayes/app/Main.hs b/rhine-bayes/app/Main.hs
index 7704d502..d4786603 100644
--- a/rhine-bayes/app/Main.hs
+++ b/rhine-bayes/app/Main.hs
@@ -330,7 +330,7 @@ inferenceRMSMC = hoistClSF sampleIOGloss inferenceBehaviour @@ liftClock Busy
  where
   inferenceBehaviour :: (MonadDistribution m, Diff td ~ Double, MonadIO m) => BehaviourF m td (Temperature, (Sensor, Pos)) Result
   inferenceBehaviour = proc (temperature, (measured, latent)) -> do
-    particles <- resampleMoveSequentialMonteCarloCl nParticles 10 resampleSystematic posteriorTemperatureProcess -< measured
+    particles <- resampleMoveSequentialMonteCarloCl 2 1 resampleSystematic posteriorTemperatureProcess -< measured
     returnA -< Result{temperature, measured, latent, particles}
 
 -- | Visualize the current 'Result' at a rate controlled by the @gloss@ backend, usually 30 FPS.

From 103073a12170ec3820d0ec5ecc997aa99786bd3e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Manuel=20B=C3=A4renz?= <programming@manuelbaerenz.de>
Date: Tue, 2 May 2023 10:40:19 +0200
Subject: [PATCH 3/9] WIP Monomorphize to Static.Traced to debug

---
 .../src/Data/MonadicStreamFunction/Bayes.hs   | 35 ++++++++++++++-----
 rhine-bayes/src/FRP/Rhine/Bayes.hs            |  8 ++---
 2 files changed, 31 insertions(+), 12 deletions(-)

diff --git a/rhine-bayes/src/Data/MonadicStreamFunction/Bayes.hs b/rhine-bayes/src/Data/MonadicStreamFunction/Bayes.hs
index dc8bebe3..4b652b97 100644
--- a/rhine-bayes/src/Data/MonadicStreamFunction/Bayes.hs
+++ b/rhine-bayes/src/Data/MonadicStreamFunction/Bayes.hs
@@ -12,13 +12,14 @@ import Numeric.Log hiding (sum)
 
 -- monad-bayes
 import Control.Monad.Bayes.Population
-import Control.Monad.Bayes.Traced.Class
+import Control.Monad.Bayes.Traced
 
 -- dunai
 import Data.MonadicStreamFunction
 import Data.MonadicStreamFunction.InternalCore (MSF (..))
 import Control.Monad.Trans.Class
 import Control.Monad.Bayes.Class (MonadDistribution)
+import qualified Control.Monad.Bayes.Traced.Static as Static
 
 -- FIXME rename to sequentialMonteCarlo or smc?
 -- | Run the Sequential Monte Carlo algorithm continuously on an 'MSF'
@@ -53,35 +54,53 @@ runPopulationsS resampler = go
           (swap . fmap fst &&& swap . fmap snd) . swap <$> bAndMSFs
 
 resampleMoveSequentialMonteCarlo ::
-  forall t m a b.
-  (MonadDistribution m, HasTraced t, MonadTrans t) =>
+  forall m a b.
+  MonadDistribution m =>
+  -- (MonadDistribution m, HasTraced t, MonadTrans t) =>
   -- | Number of particles
   Int ->
   -- | Number of MC steps
   Int ->
   -- | Resampler
   (forall x. Population m x -> Population m x) ->
-  MSF (t (Population m)) a b ->
+  MSF (Static.Traced (Population m)) a b ->
+  -- MSF (t (Population m)) a b ->
   -- FIXME Why not MSF m a (Population b)
   MSF m a [(b, Log Double)]
 resampleMoveSequentialMonteCarlo nParticles nMC resampler = go . (spawn nParticles $>)
   where
     go ::
       Monad m =>
-      Population m (MSF (t (Population m)) a b) ->
+      Population m (MSF (Static.Traced (Population m)) a b) ->
+      -- Population m (MSF (t (Population m)) a b) ->
       MSF m a [(b, Log Double)]
     go msfs = MSF $ \a -> do
       -- TODO This is quite different than the dunai version now. Maybe it's right nevertheless.
       -- FIXME This normalizes, which introduces bias, whatever that means
-      bAndMSFs <- runPopulation $ normalize $ marginal $ freeze $ composeCopies nMC mhStep $ hoistTrace resampler $ flip unMSF a =<< lift msfs
+      bAndMSFs <- runPopulation
+        $ normalize
+        $ Debug.Trace.trace "4"
+        $ marginal
+        $ Debug.Trace.trace "3"
+        $ composeCopies nMC (Debug.Trace.trace "mhStep" mhStep)
+        $ Debug.Trace.trace "2"
+        $ Control.Monad.Bayes.Traced.hoist (Debug.Trace.trace "resampler" resampler)
+        $ Debug.Trace.trace "1"
+        $ flip unMSF (Debug.Trace.trace "a" a) =<< lift (tracePop "msfs" msfs)
       return $
         second (go . fromWeightedList . return) $
           unzip $
-            (swap . fmap fst &&& swap . fmap snd) . swap <$> bAndMSFs
+            (swap . fmap fst &&& swap . fmap snd) . swap <$> Debug.Trace.trace "bAndMSFs" bAndMSFs
+
+-- | Apply a function a given number of times.
+composeCopies :: Int -> (a -> a) -> (a -> a)
+composeCopies k f = foldr (.) id (replicate k f)
 
+tracePop :: Monad m => String -> Population m a -> Population m a
+tracePop msg = fromWeightedList . fmap (\pop -> Debug.Trace.traceShow (msg, length pop) pop) . runPopulation
 
 -- resampleMoveSequentialMonteCarlo nParticles nMC resampler = morphS marginal $ runPopulationS nParticles $ freeze . composeCopies nMC mhStep . hoistTrace resampler
 
 -- FIXME see PR re-adding this to monad-bayes
 normalize :: Monad m => Population m a -> Population m a
-normalize = fromWeightedList . fmap (\particles -> second (/ (sum $ snd <$> particles)) <$> particles) . runPopulation
+normalize = fromWeightedList . fmap (\particles -> Debug.Trace.traceShow ("length particles", length particles) $ second (/ (sum $ snd <$> particles)) <$> particles) . runPopulation
diff --git a/rhine-bayes/src/FRP/Rhine/Bayes.hs b/rhine-bayes/src/FRP/Rhine/Bayes.hs
index e6200bed..5974a4bc 100644
--- a/rhine-bayes/src/FRP/Rhine/Bayes.hs
+++ b/rhine-bayes/src/FRP/Rhine/Bayes.hs
@@ -15,8 +15,7 @@ import qualified Data.MonadicStreamFunction.Bayes as DunaiBayes
 
 -- rhine
 import FRP.Rhine
-import Control.Monad.Bayes.Traced.Class (HasTraced)
-import Control.Monad.Trans.Class (MonadTrans)
+import qualified Control.Monad.Bayes.Traced.Static as Static
 
 -- * Inference methods
 
@@ -35,7 +34,8 @@ runPopulationCl :: forall m cl a b . Monad m =>
 runPopulationCl nParticles resampler = DunaiReader.readerS . DunaiBayes.runPopulationS nParticles resampler . DunaiReader.runReaderS
 
 -- | Run the Resample Move Sequential Monte Carlo algorithm continuously on a 'ClSF'.
-resampleMoveSequentialMonteCarloCl :: forall t m cl a b . (MonadDistribution m, HasTraced t, MonadTrans t) =>
+resampleMoveSequentialMonteCarloCl :: forall m cl a b . (MonadDistribution m) =>
+-- resampleMoveSequentialMonteCarloCl :: forall t m cl a b . (MonadDistribution m, HasTraced t, MonadTrans t) =>
   -- | Number of particles
   Int ->
   -- | Number of MC steps
@@ -46,7 +46,7 @@ resampleMoveSequentialMonteCarloCl :: forall t m cl a b . (MonadDistribution m,
   --   @a@ represents observations upon which the model should condition, using e.g. 'score'.
   --   It can also additionally contain hyperparameters.
   --   @b@ is the type of estimated current state.
-  -> ClSF (t (Population m)) cl a b
+  -> ClSF (Static.Traced (Population m)) cl a b
   -> ClSF m cl a [(b, Log Double)]
 resampleMoveSequentialMonteCarloCl nParticles nMC resampler = DunaiReader.readerS . DunaiBayes.resampleMoveSequentialMonteCarlo nParticles nMC resampler . DunaiReader.runReaderS
 

From 87d9bf01d7610be0fbe6981421be49f17435fa51 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Manuel=20B=C3=A4renz?= <programming@manuelbaerenz.de>
Date: Tue, 2 May 2023 12:21:35 +0200
Subject: [PATCH 4/9] WIP maybe this is RMSMC? Still a time problem

---
 rhine-bayes/app/Main.hs                       |  2 +-
 .../src/Data/MonadicStreamFunction/Bayes.hs   | 32 ++++++++-----------
 2 files changed, 15 insertions(+), 19 deletions(-)

diff --git a/rhine-bayes/app/Main.hs b/rhine-bayes/app/Main.hs
index d4786603..36b14453 100644
--- a/rhine-bayes/app/Main.hs
+++ b/rhine-bayes/app/Main.hs
@@ -330,7 +330,7 @@ inferenceRMSMC = hoistClSF sampleIOGloss inferenceBehaviour @@ liftClock Busy
  where
   inferenceBehaviour :: (MonadDistribution m, Diff td ~ Double, MonadIO m) => BehaviourF m td (Temperature, (Sensor, Pos)) Result
   inferenceBehaviour = proc (temperature, (measured, latent)) -> do
-    particles <- resampleMoveSequentialMonteCarloCl 2 1 resampleSystematic posteriorTemperatureProcess -< measured
+    particles <- resampleMoveSequentialMonteCarloCl 10 1 resampleSystematic posteriorTemperatureProcess -< measured
     returnA -< Result{temperature, measured, latent, particles}
 
 -- | Visualize the current 'Result' at a rate controlled by the @gloss@ backend, usually 30 FPS.
diff --git a/rhine-bayes/src/Data/MonadicStreamFunction/Bayes.hs b/rhine-bayes/src/Data/MonadicStreamFunction/Bayes.hs
index 4b652b97..690733ce 100644
--- a/rhine-bayes/src/Data/MonadicStreamFunction/Bayes.hs
+++ b/rhine-bayes/src/Data/MonadicStreamFunction/Bayes.hs
@@ -20,6 +20,8 @@ import Data.MonadicStreamFunction.InternalCore (MSF (..))
 import Control.Monad.Trans.Class
 import Control.Monad.Bayes.Class (MonadDistribution)
 import qualified Control.Monad.Bayes.Traced.Static as Static
+import Control.Monad.Bayes.Sequential.Coroutine (hoistFirst)
+import Control.Monad.Trans.MSF (performOnFirstSample)
 
 -- FIXME rename to sequentialMonteCarlo or smc?
 -- | Run the Sequential Monte Carlo algorithm continuously on an 'MSF'
@@ -67,40 +69,34 @@ resampleMoveSequentialMonteCarlo ::
   -- MSF (t (Population m)) a b ->
   -- FIXME Why not MSF m a (Population b)
   MSF m a [(b, Log Double)]
-resampleMoveSequentialMonteCarlo nParticles nMC resampler = go . (spawn nParticles $>)
+resampleMoveSequentialMonteCarlo nParticles nMC resampler = go . Control.Monad.Bayes.Traced.hoist (spawn nParticles >>) . pure
   where
     go ::
       Monad m =>
-      Population m (MSF (Static.Traced (Population m)) a b) ->
+      Static.Traced (Population m) (MSF (Static.Traced (Population m)) a b) ->
       -- Population m (MSF (t (Population m)) a b) ->
       MSF m a [(b, Log Double)]
     go msfs = MSF $ \a -> do
       -- TODO This is quite different than the dunai version now. Maybe it's right nevertheless.
       -- FIXME This normalizes, which introduces bias, whatever that means
-      bAndMSFs <- runPopulation
-        $ normalize
-        $ Debug.Trace.trace "4"
-        $ marginal
-        $ Debug.Trace.trace "3"
-        $ composeCopies nMC (Debug.Trace.trace "mhStep" mhStep)
-        $ Debug.Trace.trace "2"
-        $ Control.Monad.Bayes.Traced.hoist (Debug.Trace.trace "resampler" resampler)
-        $ Debug.Trace.trace "1"
-        $ flip unMSF (Debug.Trace.trace "a" a) =<< lift (tracePop "msfs" msfs)
-      return $
-        second (go . fromWeightedList . return) $
-          unzip $
-            (swap . fmap fst &&& swap . fmap snd) . swap <$> Debug.Trace.trace "bAndMSFs" bAndMSFs
+      let bAndMSFs = composeCopies nMC mhStep
+            -- $ Debug.Trace.trace "2"
+            $ Control.Monad.Bayes.Traced.hoist (tracePop "after resampler" . Debug.Trace.trace "resampler" resampler . tracePop "before resampler")
+            -- $ Debug.Trace.trace "1"
+            $ flip unMSF (Debug.Trace.trace "a" a) =<< Control.Monad.Bayes.Traced.hoist (tracePop "msfs") msfs
+      bs <- runPopulation $ marginal $ fst <$> bAndMSFs
+      return (bs, go $ snd <$> bAndMSFs)
 
 -- | Apply a function a given number of times.
 composeCopies :: Int -> (a -> a) -> (a -> a)
 composeCopies k f = foldr (.) id (replicate k f)
 
 tracePop :: Monad m => String -> Population m a -> Population m a
-tracePop msg = fromWeightedList . fmap (\pop -> Debug.Trace.traceShow (msg, length pop) pop) . runPopulation
+-- tracePop msg = fromWeightedList . fmap (\pop -> Debug.Trace.traceShow (msg, length pop) pop) . runPopulation
+tracePop _ = id
 
 -- resampleMoveSequentialMonteCarlo nParticles nMC resampler = morphS marginal $ runPopulationS nParticles $ freeze . composeCopies nMC mhStep . hoistTrace resampler
 
 -- FIXME see PR re-adding this to monad-bayes
 normalize :: Monad m => Population m a -> Population m a
-normalize = fromWeightedList . fmap (\particles -> Debug.Trace.traceShow ("length particles", length particles) $ second (/ (sum $ snd <$> particles)) <$> particles) . runPopulation
+normalize = fromWeightedList . fmap (\particles -> second (/ (sum $ snd <$> particles)) <$> particles) . runPopulation

From e7ceea167c6aa3d5425d64b29159f29d00002d55 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Manuel=20B=C3=A4renz?= <programming@manuelbaerenz.de>
Date: Tue, 2 May 2023 12:24:36 +0200
Subject: [PATCH 5/9] FIXUP? trace

---
 rhine-bayes/src/Data/MonadicStreamFunction/Bayes.hs | 1 +
 1 file changed, 1 insertion(+)

diff --git a/rhine-bayes/src/Data/MonadicStreamFunction/Bayes.hs b/rhine-bayes/src/Data/MonadicStreamFunction/Bayes.hs
index 690733ce..b1ca3985 100644
--- a/rhine-bayes/src/Data/MonadicStreamFunction/Bayes.hs
+++ b/rhine-bayes/src/Data/MonadicStreamFunction/Bayes.hs
@@ -4,6 +4,7 @@ module Data.MonadicStreamFunction.Bayes where
 import Control.Arrow
 import Data.Functor (($>))
 import Data.Tuple (swap)
+import Debug.Trace
 
 -- transformers
 

From cb1795384f2c582c0fb74f7fb89e36781a5615ae Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Manuel=20B=C3=A4renz?= <programming@manuelbaerenz.de>
Date: Wed, 10 May 2023 14:34:36 +0200
Subject: [PATCH 6/9] FIXUP Remove some debug trace

---
 rhine-bayes/src/Data/MonadicStreamFunction/Bayes.hs | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/rhine-bayes/src/Data/MonadicStreamFunction/Bayes.hs b/rhine-bayes/src/Data/MonadicStreamFunction/Bayes.hs
index b1ca3985..6c9a4762 100644
--- a/rhine-bayes/src/Data/MonadicStreamFunction/Bayes.hs
+++ b/rhine-bayes/src/Data/MonadicStreamFunction/Bayes.hs
@@ -81,10 +81,8 @@ resampleMoveSequentialMonteCarlo nParticles nMC resampler = go . Control.Monad.B
       -- TODO This is quite different than the dunai version now. Maybe it's right nevertheless.
       -- FIXME This normalizes, which introduces bias, whatever that means
       let bAndMSFs = composeCopies nMC mhStep
-            -- $ Debug.Trace.trace "2"
-            $ Control.Monad.Bayes.Traced.hoist (tracePop "after resampler" . Debug.Trace.trace "resampler" resampler . tracePop "before resampler")
-            -- $ Debug.Trace.trace "1"
-            $ flip unMSF (Debug.Trace.trace "a" a) =<< Control.Monad.Bayes.Traced.hoist (tracePop "msfs") msfs
+            $ Control.Monad.Bayes.Traced.hoist resampler
+            $ flip unMSF a =<< msfs
       bs <- runPopulation $ marginal $ fst <$> bAndMSFs
       return (bs, go $ snd <$> bAndMSFs)
 

From 68baa53b157adc8f731d6cd904488ecf0c9d1b4c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Manuel=20B=C3=A4renz?= <programming@manuelbaerenz.de>
Date: Fri, 12 May 2023 16:07:46 +0200
Subject: [PATCH 7/9] Add RMSMC with Dynamic trace

---
 rhine-bayes/app/Main.hs                       | 20 +++++++++++
 .../src/Data/MonadicStreamFunction/Bayes.hs   | 34 ++++++++++++++++++-
 rhine-bayes/src/FRP/Rhine/Bayes.hs            | 17 ++++++++++
 3 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/rhine-bayes/app/Main.hs b/rhine-bayes/app/Main.hs
index 36b14453..92c2819f 100644
--- a/rhine-bayes/app/Main.hs
+++ b/rhine-bayes/app/Main.hs
@@ -235,6 +235,7 @@ mains =
   [ ("single rate", mainSingleRate)
   , ("multi rate, temperature process", mainMultiRate)
   , ("multi rate, temperature process, RMSMC", mainMultiRateRMSMC)
+  , ("multi rate, temperature process, RMSMC dynamic", mainMultiRateRMSMCDyn)
   ]
 
 main :: IO ()
@@ -333,6 +334,14 @@ inferenceRMSMC = hoistClSF sampleIOGloss inferenceBehaviour @@ liftClock Busy
     particles <- resampleMoveSequentialMonteCarloCl 10 1 resampleSystematic posteriorTemperatureProcess -< measured
     returnA -< Result{temperature, measured, latent, particles}
 
+inferenceRMSMCDyn :: Rhine (GlossConcT IO) (LiftClock IO GlossConcT Busy) (Temperature, (Sensor, Pos)) Result
+inferenceRMSMCDyn = hoistClSF sampleIOGloss inferenceBehaviour @@ liftClock Busy
+ where
+  inferenceBehaviour :: (MonadDistribution m, TimeDomain td, Diff td ~ Double, MonadIO m) => BehaviourF m td (Temperature, (Sensor, Pos)) Result
+  inferenceBehaviour = proc (temperature, (measured, latent)) -> do
+    particles <- resampleMoveSequentialMonteCarloDynCl 10 1 (onlyBelowEffectiveSampleSize 5 resampleSystematic) posteriorTemperatureProcess -< measured
+    returnA -< Result{temperature, measured, latent, particles}
+
 -- | Visualize the current 'Result' at a rate controlled by the @gloss@ backend, usually 30 FPS.
 visualisationRhine :: Rhine (GlossConcT IO) (GlossClockUTC GlossSimClockIO) Result ()
 visualisationRhine = hoistClSF sampleIOGloss visualisation @@ glossClockUTC GlossSimClockIO
@@ -358,6 +367,17 @@ mainRhineMultiRateRMSMC =
             >-- keepLast Result{temperature = initialTemperature, measured = zeroVector, latent = zeroVector, particles = []} -@- glossConcurrently -->
               visualisationRhine
 
+
+mainRhineMultiRateRMSMCDyn =
+  userTemperature
+    @@ glossClockUTC GlossEventClockIO
+      >-- keepLast initialTemperature -@- glossConcurrently -->
+        modelRhine
+        >-- keepLast (initialTemperature, (zeroVector, zeroVector)) -@- glossConcurrently -->
+          inferenceRMSMCDyn
+            >-- keepLast Result{temperature = initialTemperature, measured = zeroVector, latent = zeroVector, particles = []} -@- glossConcurrently -->
+              visualisationRhine
+
 mainMultiRate :: IO ()
 mainMultiRate =
   void $
diff --git a/rhine-bayes/src/Data/MonadicStreamFunction/Bayes.hs b/rhine-bayes/src/Data/MonadicStreamFunction/Bayes.hs
index 6c9a4762..3a71edea 100644
--- a/rhine-bayes/src/Data/MonadicStreamFunction/Bayes.hs
+++ b/rhine-bayes/src/Data/MonadicStreamFunction/Bayes.hs
@@ -23,6 +23,7 @@ import Control.Monad.Bayes.Class (MonadDistribution)
 import qualified Control.Monad.Bayes.Traced.Static as Static
 import Control.Monad.Bayes.Sequential.Coroutine (hoistFirst)
 import Control.Monad.Trans.MSF (performOnFirstSample)
+import qualified Control.Monad.Bayes.Traced.Dynamic as Dynamic
 
 -- FIXME rename to sequentialMonteCarlo or smc?
 -- | Run the Sequential Monte Carlo algorithm continuously on an 'MSF'
@@ -80,12 +81,43 @@ resampleMoveSequentialMonteCarlo nParticles nMC resampler = go . Control.Monad.B
     go msfs = MSF $ \a -> do
       -- TODO This is quite different than the dunai version now. Maybe it's right nevertheless.
       -- FIXME This normalizes, which introduces bias, whatever that means
-      let bAndMSFs = composeCopies nMC mhStep
+      let bAndMSFs =  composeCopies nMC mhStep
             $ Control.Monad.Bayes.Traced.hoist resampler
             $ flip unMSF a =<< msfs
       bs <- runPopulation $ marginal $ fst <$> bAndMSFs
       return (bs, go $ snd <$> bAndMSFs)
 
+resampleMoveSequentialMonteCarloDynamic ::
+  forall m a b.
+  MonadDistribution m =>
+  -- (MonadDistribution m, HasTraced t, MonadTrans t) =>
+  -- | Number of particles
+  Int ->
+  -- | Number of MC steps
+  Int ->
+  -- | Resampler
+  (forall x. Population m x -> Population m x) ->
+  MSF (Dynamic.Traced (Population m)) a b ->
+  -- MSF (t (Population m)) a b ->
+  -- FIXME Why not MSF m a (Population b)
+  MSF m a [(b, Log Double)]
+resampleMoveSequentialMonteCarloDynamic nParticles nMC resampler = go . Dynamic.hoist (spawn nParticles >>) . pure
+  where
+    go ::
+      Monad m =>
+      Dynamic.Traced (Population m) (MSF (Dynamic.Traced (Population m)) a b) ->
+      -- Population m (MSF (t (Population m)) a b) ->
+      MSF m a [(b, Log Double)]
+    go msfs = MSF $ \a -> do
+      -- TODO This is quite different than the dunai version now. Maybe it's right nevertheless.
+      -- FIXME This normalizes, which introduces bias, whatever that means
+      let bAndMSFs = Dynamic.freeze
+            $ composeCopies nMC Dynamic.mhStep
+            $ Dynamic.hoist resampler
+            $ flip unMSF a =<< msfs
+      bs <- runPopulation $ Dynamic.marginal $ fst <$> bAndMSFs
+      return (bs, go $ snd <$> bAndMSFs)
+
 -- | Apply a function a given number of times.
 composeCopies :: Int -> (a -> a) -> (a -> a)
 composeCopies k f = foldr (.) id (replicate k f)
diff --git a/rhine-bayes/src/FRP/Rhine/Bayes.hs b/rhine-bayes/src/FRP/Rhine/Bayes.hs
index 5974a4bc..ed39fc11 100644
--- a/rhine-bayes/src/FRP/Rhine/Bayes.hs
+++ b/rhine-bayes/src/FRP/Rhine/Bayes.hs
@@ -16,6 +16,7 @@ import qualified Data.MonadicStreamFunction.Bayes as DunaiBayes
 -- rhine
 import FRP.Rhine
 import qualified Control.Monad.Bayes.Traced.Static as Static
+import qualified Control.Monad.Bayes.Traced.Dynamic as Dynamic
 
 -- * Inference methods
 
@@ -50,6 +51,22 @@ resampleMoveSequentialMonteCarloCl :: forall m cl a b . (MonadDistribution m) =>
   -> ClSF m cl a [(b, Log Double)]
 resampleMoveSequentialMonteCarloCl nParticles nMC resampler = DunaiReader.readerS . DunaiBayes.resampleMoveSequentialMonteCarlo nParticles nMC resampler . DunaiReader.runReaderS
 
+resampleMoveSequentialMonteCarloDynCl :: forall m cl a b . (MonadDistribution m) =>
+-- resampleMoveSequentialMonteCarloCl :: forall t m cl a b . (MonadDistribution m, HasTraced t, MonadTrans t) =>
+  -- | Number of particles
+  Int ->
+  -- | Number of MC steps
+  Int ->
+  -- | Resampler (see 'Control.Monad.Bayes.Population' for some standard choices)
+  (forall x . Population m x -> Population m x)
+  -- | A signal function modelling the stochastic process on which to perform inference.
+  --   @a@ represents observations upon which the model should condition, using e.g. 'score'.
+  --   It can also additionally contain hyperparameters.
+  --   @b@ is the type of estimated current state.
+  -> ClSF (Dynamic.Traced (Population m)) cl a b
+  -> ClSF m cl a [(b, Log Double)]
+resampleMoveSequentialMonteCarloDynCl nParticles nMC resampler = DunaiReader.readerS . DunaiBayes.resampleMoveSequentialMonteCarloDynamic nParticles nMC resampler . DunaiReader.runReaderS
+
 -- * Short standard library of stochastic processes
 
 -- | White noise, that is, an independent normal distribution at every time step.

From cc4f82b5dff1c3b9b9fdf0e14db96930de8993bc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Manuel=20B=C3=A4renz?= <programming@manuelbaerenz.de>
Date: Fri, 12 May 2023 16:10:16 +0200
Subject: [PATCH 8/9] WIP Vary diffusion temperature with effective sample size

---
 rhine-bayes/app/Main.hs                       | 39 +++++++++++++------
 .../src/Data/MonadicStreamFunction/Bayes.hs   | 38 ++++++++++++++++++
 rhine-bayes/src/FRP/Rhine/Bayes.hs            | 10 ++++-
 3 files changed, 75 insertions(+), 12 deletions(-)

diff --git a/rhine-bayes/app/Main.hs b/rhine-bayes/app/Main.hs
index 92c2819f..14ee5347 100644
--- a/rhine-bayes/app/Main.hs
+++ b/rhine-bayes/app/Main.hs
@@ -113,9 +113,9 @@ initialTemperature :: Temperature
 initialTemperature = 7
 
 -- | We infer the temperature by randomly moving around with a Brownian motion (Wiener process).
-temperatureProcess :: (MonadDistribution m, Diff td ~ Double) => BehaviourF m td () Temperature
-temperatureProcess = proc () -> do
-  temperatureFactor <- wienerLogDomain 20 -< ()
+temperatureProcess :: (MonadDistribution m, Diff td ~ Double) => BehaviourF m td (Diff td) Temperature
+temperatureProcess = proc dt -> do
+  temperatureFactor <- wienerVaryingLogDomain -< dt
   returnA -< runLogDomain temperatureFactor * initialTemperature
 
 -- | Auxiliary conversion function belonging to the log-domain library, see https://github.com/ekmett/log-domain/issues/38
@@ -133,16 +133,33 @@ genModelWithoutTemperature = proc temperature -> do
   sensor <- generativeModel -< latent
   returnA -< (sensor, latent)
 
+type ESS = Double
+
 {- | Given sensor data, sample a latent position and a temperature, and weight them according to the likelihood of the observed sensor position.
    Used to infer position and temperature.
 -}
-posteriorTemperatureProcess :: (MonadMeasure m, Diff td ~ Double) => BehaviourF m td Sensor (Temperature, Pos)
-posteriorTemperatureProcess = proc sensor -> do
-  temperature <- temperatureProcess -< ()
+posteriorTemperatureProcessAutoESS :: (MonadDistribution m, TimeDomain td, Diff td ~ Double) => BehaviourF (Population m) td Sensor (Temperature, Pos)
+posteriorTemperatureProcessAutoESS = withESS 20 posteriorTemperatureProcessLiveESS
+
+{- | Given sensor data, sample a latent position and a temperature, and weight them according to the likelihood of the observed sensor position.
+   Used to infer position and temperature.
+-}
+posteriorTemperatureProcessLiveESS :: (MonadMeasure m, TimeDomain td, Diff td ~ Double) => BehaviourF m td (Sensor, ESS) (Temperature, Pos)
+posteriorTemperatureProcessLiveESS = proc (sensor, ess) -> do
+  t <- sinceStart -< ()
+  temperature <- temperatureProcess -< ess / 2 + t / 20
   latent <- prior -< temperature
   arrM score -< sensorLikelihood latent sensor
   returnA -< (temperature, latent)
 
+{- | Given sensor data, sample a latent position and a temperature, and weight them according to the likelihood of the observed sensor position.
+   Used to infer position and temperature.
+-}
+posteriorTemperatureProcess :: (MonadMeasure m, TimeDomain td, Diff td ~ Double) => BehaviourF m td Sensor (Temperature, Pos)
+posteriorTemperatureProcess = proc sensor -> do
+  posteriorTemperatureProcessLiveESS -< (sensor, 20)
+
+
 -- | A collection of all displayable inference results
 data Result = Result
   { temperature :: Temperature
@@ -249,7 +266,7 @@ main = do
 {- | Given an actual temperature, simulate a latent position and measured sensor position,
    and based on the sensor data infer the latent position and the temperature.
 -}
-filtered :: Diff td ~ Double => BehaviourF App td Temperature Result
+filtered :: (TimeDomain td, Diff td ~ Double) => BehaviourF App td Temperature Result
 filtered = proc temperature -> do
   (measured, latent) <- genModelWithoutTemperature -< temperature
   particles <- runPopulationCl nParticles resampleSystematic posteriorTemperatureProcess -< measured
@@ -263,7 +280,7 @@ filtered = proc temperature -> do
         }
 
 -- | Run simulation, inference, and visualization synchronously
-mainClSF :: Diff td ~ Double => BehaviourF App td () ()
+mainClSF :: (TimeDomain td, Diff td ~ Double) => BehaviourF App td () ()
 mainClSF = proc () -> do
   output <- filtered -< initialTemperature
   visualisation -< output
@@ -318,9 +335,9 @@ userTemperature = tagS >>> arr (selector >>> fmap Product) >>> mappendS >>> arr
 inference :: Rhine (GlossConcT IO) (LiftClock IO GlossConcT Busy) (Temperature, (Sensor, Pos)) Result
 inference = hoistClSF sampleIOGloss inferenceBehaviour @@ liftClock Busy
  where
-  inferenceBehaviour :: (MonadDistribution m, Diff td ~ Double, MonadIO m) => BehaviourF m td (Temperature, (Sensor, Pos)) Result
+  inferenceBehaviour :: (MonadDistribution m, TimeDomain td, Diff td ~ Double, MonadIO m) => BehaviourF m td (Temperature, (Sensor, Pos)) Result
   inferenceBehaviour = proc (temperature, (measured, latent)) -> do
-    particles <- runPopulationCl nParticles resampleSystematic posteriorTemperatureProcess -< measured
+    particles <- runPopulationCl nParticles (onlyBelowEffectiveSampleSize 60 resampleSystematic) posteriorTemperatureProcessAutoESS -< measured
     returnA -< Result{temperature, measured, latent, particles}
 
 {- | This part performs the inference (and passes along temperature, sensor and position simulations).
@@ -329,7 +346,7 @@ inference = hoistClSF sampleIOGloss inferenceBehaviour @@ liftClock Busy
 inferenceRMSMC :: Rhine (GlossConcT IO) (LiftClock IO GlossConcT Busy) (Temperature, (Sensor, Pos)) Result
 inferenceRMSMC = hoistClSF sampleIOGloss inferenceBehaviour @@ liftClock Busy
  where
-  inferenceBehaviour :: (MonadDistribution m, Diff td ~ Double, MonadIO m) => BehaviourF m td (Temperature, (Sensor, Pos)) Result
+  inferenceBehaviour :: (MonadDistribution m, TimeDomain td, Diff td ~ Double, MonadIO m) => BehaviourF m td (Temperature, (Sensor, Pos)) Result
   inferenceBehaviour = proc (temperature, (measured, latent)) -> do
     particles <- resampleMoveSequentialMonteCarloCl 10 1 resampleSystematic posteriorTemperatureProcess -< measured
     returnA -< Result{temperature, measured, latent, particles}
diff --git a/rhine-bayes/src/Data/MonadicStreamFunction/Bayes.hs b/rhine-bayes/src/Data/MonadicStreamFunction/Bayes.hs
index 3a71edea..5997e528 100644
--- a/rhine-bayes/src/Data/MonadicStreamFunction/Bayes.hs
+++ b/rhine-bayes/src/Data/MonadicStreamFunction/Bayes.hs
@@ -131,3 +131,41 @@ tracePop _ = id
 -- FIXME see PR re-adding this to monad-bayes
 normalize :: Monad m => Population m a -> Population m a
 normalize = fromWeightedList . fmap (\particles -> second (/ (sum $ snd <$> particles)) <$> particles) . runPopulation
+
+-- FIXME See PR to monad-bayes
+
+
+-- | 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
+  (forall n . MonadDistribution n => Population n a -> Population n a) ->
+  -- | The new resampler
+  (Population m a -> Population m a)
+onlyBelowEffectiveSampleSize threshold resampler pop = do
+  (particles, ess) <- lift $ runWithEffectiveSampleSize pop
+  let newPop = fromWeightedList $ pure particles
+  -- This assumes that the resampler does not mutate the m effects, as it should
+  if ess < threshold then resampler newPop else newPop
+
+-- | Compute the effective sample size of a population from the weights.
+--
+--   See https://en.wikipedia.org/wiki/Design_effect#Effective_sample_size
+runWithEffectiveSampleSize :: Functor m => Population m a -> m ([(a, Log Double)], Double)
+runWithEffectiveSampleSize = fmap (id &&& (effectiveSampleSizeKish . map (exp . ln . snd))) . runPopulation
+  where
+    effectiveSampleSizeKish :: [Double] -> Double
+    effectiveSampleSizeKish weights = square (sum weights) / sum (square <$> weights)
+    square :: Double -> Double
+    square x = x * x
+
+measureESS :: Monad m => MSF (Population m) a b -> MSF (Population m) a (b, Double)
+measureESS = morphGS $ fmap $ \pop -> fromWeightedList $ do
+  (particles, ess) <- runWithEffectiveSampleSize pop
+  pure $ map (first (first (, ess))) particles
+
+withESS :: Monad m => Double -> MSF (Population m) (a, Double) b -> MSF (Population m) a b
+withESS initESS = feedback initESS . measureESS
diff --git a/rhine-bayes/src/FRP/Rhine/Bayes.hs b/rhine-bayes/src/FRP/Rhine/Bayes.hs
index ed39fc11..fd4374f7 100644
--- a/rhine-bayes/src/FRP/Rhine/Bayes.hs
+++ b/rhine-bayes/src/FRP/Rhine/Bayes.hs
@@ -1,4 +1,5 @@
-module FRP.Rhine.Bayes where
+module FRP.Rhine.Bayes
+  (module FRP.Rhine.Bayes, module X) where
 
 -- log-domain
 import Numeric.Log hiding (sum)
@@ -13,6 +14,8 @@ import qualified Control.Monad.Trans.MSF.Reader as DunaiReader
 -- dunai-bayes
 import qualified Data.MonadicStreamFunction.Bayes as DunaiBayes
 
+import Data.MonadicStreamFunction.Bayes as X (onlyBelowEffectiveSampleSize)
+
 -- rhine
 import FRP.Rhine
 import qualified Control.Monad.Bayes.Traced.Static as Static
@@ -118,3 +121,8 @@ wienerVaryingLogDomain ::
   (MonadDistribution m, Diff td ~ Double) =>
   BehaviourF m td (Diff td) (Log Double)
 wienerVaryingLogDomain = wienerVarying >>> arr Exp
+
+withESS :: Monad m => Double -> ClSF (Population m) cl (a, Double) b -> ClSF (Population m) cl a b
+withESS initESS = DunaiReader.readerS . DunaiBayes.withESS initESS . (arr assoc >>>) . DunaiReader.runReaderS
+  where
+    assoc ((ti, a), td) = (ti, (a, td))

From 8717ea0776c54faf377073618a48ab1bd3286c69 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Manuel=20B=C3=A4renz?= <programming@manuelbaerenz.de>
Date: Fri, 12 May 2023 16:11:08 +0200
Subject: [PATCH 9/9] CHERRY dyn

---
 rhine-bayes/app/Main.hs | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/rhine-bayes/app/Main.hs b/rhine-bayes/app/Main.hs
index 14ee5347..acbf8291 100644
--- a/rhine-bayes/app/Main.hs
+++ b/rhine-bayes/app/Main.hs
@@ -407,6 +407,12 @@ mainMultiRateRMSMC =
     launchGlossThread glossSettings $
       flow mainRhineMultiRateRMSMC
 
+mainMultiRateRMSMCDyn :: IO ()
+mainMultiRateRMSMCDyn =
+  void $
+    launchGlossThread glossSettings $
+      flow mainRhineMultiRateRMSMCDyn
+
 -- * Utilities
 
 instance MonadDistribution m => MonadDistribution (GlossConcT m) where