Add and use applicative population transformer

Author: turion

monad-bayes.cabal

@@ -63,21 +63,21 @@ common deps
     , scientific       ^>=0.3
     , statistics       >=0.14.0  && <0.17
     , text             >=1.2     && <2.1
+    , transformers     ^>=0.5.6
     , vector           >=0.12.0  && <0.14
     , vty              ^>=5.38
 
 common test-deps
   build-depends:
     , abstract-par          ^>=0.3
-    , criterion             >=1.5    && <1.7
+    , criterion             >=1.5   && <1.7
     , directory             ^>=1.3
     , hspec                 ^>=2.11
     , monad-bayes
-    , optparse-applicative  >=0.17   && <0.19
+    , optparse-applicative  >=0.17  && <0.19
     , process               ^>=1.6
     , QuickCheck            ^>=2.14
-    , time                  >=1.9    && <1.13
-    , transformers          ^>=0.5.6
+    , time                  >=1.9   && <1.13
     , typed-process         ^>=0.2
 
   autogen-modules: Paths_monad_bayes

src/Control/Applicative/List.hs

@@ -1,17 +1,40 @@
+{-# LANGUAGE StandaloneDeriving #-}
+
 module Control.Applicative.List where
 
 -- base
-
 import Control.Applicative
 import Data.Functor.Compose
 
+-- transformers
+import Control.Monad.Trans.Writer.Strict
+
+-- log-domain
+import Numeric.Log (Log)
+
+-- * Applicative ListT
+
 -- | _Applicative_ transformer adding a list/nondeterminism/choice effect.
 --   It is not a valid monad transformer, but it is a valid 'Applicative'.
-newtype ListT m a = ListT {getListT :: Compose [] m a}
+newtype ListT m a = ListT {getListT :: Compose m [] a}
   deriving newtype (Functor, Applicative, Alternative)
 
-lift :: m a -> ListT m a
-lift = ListT . Compose . pure
+lift :: (Functor m) => m a -> ListT m a
+lift = ListT . Compose . fmap pure
 
-runListT :: ListT m a -> [m a]
+runListT :: ListT m a -> m [a]
 runListT = getCompose . getListT
+
+-- * Applicative Population transformer
+
+-- WriterT has to be used instead of WeightedT,
+-- since WeightedT uses StateT under the hood,
+-- which requires a Monad (ListT m) constraint.
+newtype PopulationT m a = PopulationT {getPopulationT :: WriterT (Log Double) (ListT m) a}
+  deriving newtype (Functor, Applicative, Alternative)
+
+runPopulationT :: PopulationT m a -> m [(a, Log Double)]
+runPopulationT = runListT . runWriterT . getPopulationT
+
+fromWeightedList :: m [(a, Log Double)] -> PopulationT m a
+fromWeightedList = PopulationT . WriterT . ListT . Compose

src/Control/Monad/Bayes/Inference/RMSMC.hs

@@ -26,6 +26,7 @@ import Control.Monad.Bayes.Inference.SMC
 import Control.Monad.Bayes.Population
   ( PopulationT,
     flatten,
+    single,
     withParticles,
   )
 import Control.Monad.Bayes.Sequential.Coroutine as Seq
@@ -50,7 +51,7 @@ rmsmc ::
   PopulationT m a
 rmsmc (MCMCConfig {..}) (SMCConfig {..}) =
   marginal
-    . S.sequentially (composeCopies numMCMCSteps (TrStat.hoist flatten . mhStep) . TrStat.hoist resampler) numSteps
+    . S.sequentially (composeCopies numMCMCSteps (TrStat.hoist (single . flatten) . mhStep) . TrStat.hoist resampler) numSteps
     . S.hoistFirst (TrStat.hoist (withParticles numParticles))
 
 -- | Resample-move Sequential Monte Carlo with a more efficient
@@ -64,7 +65,7 @@ rmsmcBasic ::
   PopulationT m a
 rmsmcBasic (MCMCConfig {..}) (SMCConfig {..}) =
   TrBas.marginal
-    . S.sequentially (TrBas.hoist flatten . composeCopies numMCMCSteps (TrBas.hoist flatten . TrBas.mhStep) . TrBas.hoist resampler) numSteps
+    . S.sequentially (TrBas.hoist (single . flatten) . composeCopies numMCMCSteps (TrBas.hoist (single . flatten) . TrBas.mhStep) . TrBas.hoist resampler) numSteps
     . S.hoistFirst (TrBas.hoist (withParticles numParticles))
 
 -- | A variant of resample-move Sequential Monte Carlo
@@ -79,7 +80,7 @@ rmsmcDynamic ::
   PopulationT m a
 rmsmcDynamic (MCMCConfig {..}) (SMCConfig {..}) =
   TrDyn.marginal
-    . S.sequentially (TrDyn.freeze . composeCopies numMCMCSteps (TrDyn.hoist flatten . TrDyn.mhStep) . TrDyn.hoist resampler) numSteps
+    . S.sequentially (TrDyn.freeze . composeCopies numMCMCSteps (TrDyn.hoist (single . flatten) . TrDyn.mhStep) . TrDyn.hoist resampler) numSteps
     . S.hoistFirst (TrDyn.hoist (withParticles numParticles))
 
 -- | Apply a function a given number of times.

src/Control/Monad/Bayes/Population.hs

@@ -35,10 +35,12 @@ module Control.Monad.Bayes.Population
     popAvg,
     withParticles,
     flatten,
+    single,
   )
 where
 
 import Control.Applicative (Alternative)
+import Control.Applicative.List qualified as ApplicativeListT
 import Control.Arrow (second)
 import Control.Monad (MonadPlus, replicateM)
 import Control.Monad.Bayes.Class
@@ -277,5 +279,11 @@ hoist ::
 hoist f = PopulationT . Weighted.hoist (hoistFreeT f) . getPopulationT
 
 -- | Flatten all layers of the free structure
-flatten :: (Monad m) => PopulationT m a -> PopulationT m a
-flatten = fromWeightedList . runPopulationT
+flatten :: (Monad m) => PopulationT m a -> ApplicativeListT.PopulationT m a
+flatten = ApplicativeListT.fromWeightedList . runPopulationT
+
+-- | Create a population from a single layer of branching computations.
+--
+-- Similar to 'fromWeightedListT'.
+single :: (Monad m) => ApplicativeListT.PopulationT m a -> PopulationT m a
+single = fromWeightedList . ApplicativeListT.runPopulationT