Free list transformer to replace deprecated ListT

docs/docs/usage.md

@@ -296,7 +296,7 @@ runWeightedT (WeightedT m) = runStateT m 1
 
 `WeightedT m` is not an instance of `MonadDistribution`, but only as instance of `MonadFactor` (and that, only when `m` is an instance of `Monad`). However, since `StateT` is a monad transformer, there is a function `lift :: m Double -> WeightedT m Double`.
 
-So if we take a `MonadDistribution` instance like `SamplerIO`, then `WeightedT SamplerIO` is an instance of both `MonadDistribution` and `MonadFactor`. Which means it is an instance of `MonadMeasure`.
+So if we take a `MonadDistribution` instance like `SamplerIO`, then `Weighted SamplerIO` is an instance of both `MonadDistribution` and `MonadFactor`. Which means it is an instance of `MonadMeasure`.
 
 So we can successfully write `(sampler . runWeightedT) sprinkler` and get a program of type `IO (Bool, Log Double)`. When run, this will draw a sample from `sprinkler` along with an **unnormalized** density for that sample.
 
@@ -328,18 +328,20 @@ Summary of key info on `PopulationT`:
 - `instance MonadFactor m => instance MonadFactor (PopulationT m)`
 
 ```haskell
-newtype PopulationT m a = PopulationT (WeightedT (ListT m) a)
+newtype PopulationT m a = PopulationT (WeightedT (FreeT [] m) a)
 ```
 
-So:
+The `FreeT []` construction is for branching our probabilistic program into different branches,
+corresponding to different choices of a random variable.
+
+It is interpreted, using `runPopulationT`, to:
 
 ```haskell
-PopulationT m a ~ m [Log Double -> (a, Log Double)]
+m [(a, Log Double)]
 ```
 
-Note that while `ListT` isn't in general a valid monad transformer, we're not requiring it to be one here.
-
-`PopulationT` is used to represent a collection of particles (in the statistical sense), along with their weights.
+This shows that `Population` is used to compute a collection of particles (in the statistical sense), along with their weights.
+Each `a` corresponds to one particle, and `Log Double` is the type of its weight.
 
 There are several useful functions associated with it:
 
@@ -360,7 +362,7 @@ gives
 [([((),0.5),((),0.5)],1.0)]
 ```
 
-Observe how here we have interpreted `(spawn 2)` as of type `PopulationT Enumerator ()`.
+Observe how here we have interpreted `(spawn 2)` as of type `Population Enumerator ()`.
 
 `resampleGeneric` takes a function to probabilistically select a set of indices from a vector, and makes a new population by selecting those indices.
 
@@ -393,8 +395,8 @@ Summary of key info on `SequentialT`:
 
 
 ```haskell
-newtype SequentialT m a =
- SequentialT {runSequentialT :: Coroutine (Await ()) m a}
+newtype Sequential m a =
+ Sequential {runSequential :: Coroutine (Await ()) m a}
 ```
 
 This is a wrapper for the `Coroutine` type applied to the `Await` constructor from `Control.Monad.Coroutine`, which is defined thus:
@@ -410,7 +412,7 @@ newtype Await x y = Await (x -> y)
 Unpacking that:
 
 ```haskell
-SequentialT m a ~ m (Either (() -> SequentialT m a) a)
+Sequential m a ~ m (Either (() -> Sequential m a) a)
 ```
 
 As usual, `m` is going to be some other probability monad, so understand `SequentialT m a` as representing a program which, after making a random choice or doing conditioning, we either obtain an `a` value, or a paused computation, which when resumed gets us back to a new `SequentialT m a`.
@@ -501,11 +503,11 @@ The latter is best understood if you're familiar with the standard use of a free
 
 ```haskell
 newtype SamF a = Random (Double -> a)
-newtype DensityT m a =
- DensityT {getDensityT :: FT SamF m a}
+newtype Density m a =
+ Density {density :: FT SamF m a}
 
-instance Monad m => MonadDistribution (DensityT m) where
- random = DensityT $ liftF (Random id)
+instance Monad m => MonadDistribution (Density m) where
+ random = Density $ liftF (Random id)
 ```
 
 The monad-bayes implementation uses a more efficient implementation of `FreeT`, namely `FT` from the `free` package, known as the *Church transformed Free monad*. This is a technique explained in https://begriffs.com/posts/2016-02-04-difference-lists-and-codennsity.html. But that only changes the operational semantics - performance aside, it works just the same as the standard `FreeT` datatype.
@@ -575,7 +577,7 @@ data Trace a = Trace
  }
 ```
 
-We also need a specification of the probabilistic program in question, free of any particular interpretation. That is precisely what `DensityT` is for.
+We also need a specification of the probabilistic program in question, free of any particular interpretation. That is precisely what `Density` is for.
 
 The simplest version of `TracedT` is in `Control.Monad.Bayes.TracedT.Basic`
 
@@ -635,13 +637,13 @@ example = do
  return x
 ```
 
-`(enumerator . runWeightedT) example` gives `[((False,0.0),0.5),((True,1.0),0.5)]`. This is quite edifying for understanding `(sampler . runWeightedT) example`. What it says is that there are precisely two ways the program will run, each with equal probability: either you get `False` with weight `0.0` or `True` with weight `1.0`.
+`(enumerator . weighted) example` gives `[((False,0.0),0.5),((True,1.0),0.5)]`. This is quite edifying for understanding `(sampler . weighted) example`. What it says is that there are precisely two ways the program will run, each with equal probability: either you get `False` with weight `0.0` or `True` with weight `1.0`.
 
 ### Quadrature
 
 As described on the section on `Integrator`, we can interpret our probabilistic program of type `MonadDistribution m => m a` as having concrete type `Integrator a`. This views our program as an integrator, allowing us to calculate expectations, probabilities and so on via quadrature (i.e. numerical approximation of an integral).
 
-This can also handle programs of type `MonadMeasure m => m a`, that is, programs with `factor` statements. For these cases, a function `normalize :: WeightedT Integrator a -> Integrator a` is employed. For example,
+This can also handle programs of type `MonadMeasure m => m a`, that is, programs with `factor` statements. For these cases, a function `normalize :: Weighted Integrator a -> Integrator a` is employed. For example,
 
 ```haskell
 model :: MonadMeasure m => m Double
@@ -652,7 +654,7 @@ model = do
  return var
 ```
 
-is really an unnormalized measure, rather than a probability distribution. `normalize` views it as of type `WeightedT Integrator Double`, which is isomorphic to `(Double -> (Double, Log Double) -> Double)`. This can be used to compute the normalization constant, and divide the integrator's output by it, all within `Integrator`.
+is really an unnormalized measure, rather than a probability distribution. `normalize` views it as of type `Weighted Integrator Double`, which is isomorphic to `(Double -> (Double, Log Double) -> Double)`. This can be used to compute the normalization constant, and divide the integrator's output by it, all within `Integrator`.
 
 ### Independent forward sampling
 
@@ -796,7 +798,7 @@ pmmh ::
 pmmh mcmcConf smcConf param model =
  (mcmc mcmcConf :: T m [(a, Log Double)] -> m [[(a, Log Double)]])
  ((param :: T m b) >>=
- (runPopulationT :: P (T m) a -> T m [(a, Log Double)])
+ (population :: P (T m) a -> T m [(a, Log Double)])
  . (pushEvidence :: P (T m) a -> P (T m) a)
  . Pop.hoist (lift :: forall x. m x -> T m x)
  . (smc smcConf :: S (P m) a -> P m a)

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
@@ -86,6 +86,7 @@ common test-deps
 library
  import: deps
  exposed-modules:
+ Control.Applicative.List
  Control.Monad.Bayes.Class
  Control.Monad.Bayes.Density.Free
  Control.Monad.Bayes.Density.State
@@ -100,6 +101,7 @@ library
  Control.Monad.Bayes.Inference.TUI
  Control.Monad.Bayes.Integrator
  Control.Monad.Bayes.Population
+ Control.Monad.Bayes.Population.Applicative
  Control.Monad.Bayes.Sampler.Lazy
  Control.Monad.Bayes.Sampler.Strict
  Control.Monad.Bayes.Sequential.Coroutine
@@ -114,8 +116,11 @@ library
  other-modules: Control.Monad.Bayes.Traced.Common
  default-language: Haskell2010
  default-extensions:
+ ApplicativeDo
  BlockArguments
+ DerivingStrategies
  FlexibleContexts
+ GeneralizedNewtypeDeriving
  ImportQualifiedPost
  LambdaCase
  OverloadedStrings

src/Control/Applicative/List.hs

@@ -0,0 +1,23 @@
+{-# LANGUAGE StandaloneDeriving #-}
+
+module Control.Applicative.List where
+
+-- base
+import Control.Applicative
+import Data.Functor.Compose
+
+-- * 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}
+ deriving newtype (Functor, Applicative, Alternative)
+
+listT :: m [a] -> ListT m a
+listT = ListT . Compose
+
+lift :: (Functor m) => m a -> ListT m a
+lift = ListT . Compose . fmap pure
+
+runListT :: ListT m a -> m [a]
+runListT = getCompose . getListT

src/Control/Monad/Bayes/Class.hs

@@ -79,9 +79,9 @@ import Control.Monad (replicateM, when)
 import Control.Monad.Cont (ContT)
 import Control.Monad.Except (ExceptT, lift)
 import Control.Monad.Identity (IdentityT)
-import Control.Monad.List (ListT)
 import Control.Monad.Reader (ReaderT)
 import Control.Monad.State (StateT)
+import Control.Monad.Trans.Free (FreeT)
 import Control.Monad.Writer (WriterT)
 import Data.Histogram qualified as H
 import Data.Histogram.Fill qualified as H
@@ -390,15 +390,15 @@ instance (MonadFactor m) => MonadFactor (StateT s m) where
 
 instance (MonadMeasure m) => MonadMeasure (StateT s m)
 
-instance (MonadDistribution m) => MonadDistribution (ListT m) where
+instance (Applicative f, (MonadDistribution m)) => MonadDistribution (FreeT f m) where
  random = lift random
  bernoulli = lift . bernoulli
  categorical = lift . categorical
 
-instance (MonadFactor m) => MonadFactor (ListT m) where
+instance (Applicative f, (MonadFactor m)) => MonadFactor (FreeT f m) where
  score = lift . score
 
-instance (MonadMeasure m) => MonadMeasure (ListT m)
+instance (Applicative f, (MonadMeasure m)) => MonadMeasure (FreeT f m)
 
 instance (MonadDistribution m) => MonadDistribution (ContT r m) where
  random = lift random

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

@@ -25,7 +25,8 @@ import Control.Monad.Bayes.Inference.MCMC (MCMCConfig (..))
 import Control.Monad.Bayes.Inference.SMC
 import Control.Monad.Bayes.Population
  ( PopulationT,
- spawn,
+ flatten,
+ single,
  withParticles,
  )
 import Control.Monad.Bayes.Sequential.Coroutine as Seq
@@ -50,8 +51,8 @@ rmsmc ::
  PopulationT m a
 rmsmc (MCMCConfig {..}) (SMCConfig {..}) =
  marginal
- . S.sequentially (composeCopies numMCMCSteps mhStep . TrStat.hoist resampler) numSteps
- . S.hoistFirst (TrStat.hoist (spawn numParticles >>))
+ . S.sequentially (composeCopies numMCMCSteps (TrStat.hoistModel (single . flatten) . TrStat.hoist (single . flatten) . mhStep) . TrStat.hoist resampler) numSteps
+ . S.hoistFirst (TrStat.hoistModel (single . flatten) . TrStat.hoist (withParticles numParticles))
 
 -- | Resample-move Sequential Monte Carlo with a more efficient
 -- tracing representation.
@@ -64,7 +65,7 @@ rmsmcBasic ::
  PopulationT m a
 rmsmcBasic (MCMCConfig {..}) (SMCConfig {..}) =
  TrBas.marginal
- . S.sequentially (composeCopies numMCMCSteps 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.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/Inference/SMC.hs

@@ -22,18 +22,38 @@ where
 
 import Control.Monad.Bayes.Class (MonadDistribution, MonadMeasure)
 import Control.Monad.Bayes.Population
- ( PopulationT,
+ ( PopulationT (..),
+ flatten,
  pushEvidence,
+ single,
  withParticles,
  )
+import Control.Monad.Bayes.Population.Applicative qualified as Applicative
 import Control.Monad.Bayes.Sequential.Coroutine as Coroutine
+import Control.Monad.Bayes.Sequential.Coroutine qualified as SequentialT
+import Control.Monad.Bayes.Weighted (WeightedT (..), weightedT)
+import Control.Monad.Coroutine
+import Control.Monad.Trans.Free (FreeF (..), FreeT (..))
 
 data SMCConfig m = SMCConfig
  { resampler :: forall x. PopulationT m x -> PopulationT m x,
  numSteps :: Int,
  numParticles :: Int
  }
 
+sequentialToPopulation :: (Monad m) => Coroutine.SequentialT (Applicative.PopulationT m) a -> PopulationT m a
+sequentialToPopulation =
+ PopulationT
+ . weightedT
+ . coroutineToFree
+ . Coroutine.runSequentialT
+ where
+ coroutineToFree =
+ FreeT
+ . fmap (Free . fmap (\(cont, p) -> either (coroutineToFree . extract) (pure . (,p)) cont))
+ . Applicative.runPopulationT
+ . resume
+
 -- | Sequential importance resampling.
 -- Basically an SMC template that takes a custom resampler.
 smc ::
@@ -42,12 +62,15 @@ smc ::
  Coroutine.SequentialT (PopulationT m) a ->
  PopulationT m a
 smc SMCConfig {..} =
- Coroutine.sequentially resampler numSteps
+ (single . flatten)
+ . Coroutine.sequentially resampler numSteps
+ . SequentialT.hoist (single . flatten)
  . Coroutine.hoistFirst (withParticles numParticles)
+ . SequentialT.hoist (single . flatten)
 
 -- | Sequential Monte Carlo with multinomial resampling at each timestep.
 -- Weights are normalized at each timestep and the total weight is pushed
 -- as a score into the transformed monad.
 smcPush ::
  (MonadMeasure m) => SMCConfig m -> Coroutine.SequentialT (PopulationT m) a -> PopulationT m a
-smcPush config = smc config {resampler = (pushEvidence . resampler config)}
+smcPush config = smc config {resampler = (single . flatten . pushEvidence . resampler config)}

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

@@ -27,8 +27,10 @@ import Control.Monad.Bayes.Class
 import Control.Monad.Bayes.Inference.MCMC
 import Control.Monad.Bayes.Inference.RMSMC (rmsmc)
 import Control.Monad.Bayes.Inference.SMC (SMCConfig (SMCConfig, numParticles, numSteps, resampler), smcPush)
-import Control.Monad.Bayes.Population as Pop (PopulationT, resampleMultinomial, runPopulationT)
+import Control.Monad.Bayes.Population as Pop (PopulationT, flatten, resampleMultinomial, runPopulationT, single)
+import Control.Monad.Bayes.Population qualified as PopulationT
 import Control.Monad.Bayes.Sequential.Coroutine (SequentialT)
+import Control.Monad.Bayes.Sequential.Coroutine qualified as SequentialT
 import Control.Monad.Bayes.Traced
 import Control.Monad.Trans (MonadTrans (..))
 import Numeric.Log (Log)
@@ -71,4 +73,10 @@ smc2 k n p t param m =
  rmsmc
  MCMCConfig {numMCMCSteps = t, proposal = SingleSiteMH, numBurnIn = 0}
  SMCConfig {numParticles = p, numSteps = k, resampler = resampleMultinomial}
- (param >>= setup . runPopulationT . smcPush (SMCConfig {numSteps = k, numParticles = n, resampler = resampleMultinomial}) . m)
+ (flattenSequentiallyTraced param >>= setup . runPopulationT . smcPush (SMCConfig {numSteps = k, numParticles = n, resampler = resampleMultinomial}) . flattenSMC2 . m)
+
+flattenSequentiallyTraced :: (Monad m) => SequentialT (TracedT (PopulationT m)) a -> SequentialT (TracedT (PopulationT m)) a
+flattenSequentiallyTraced = SequentialT.hoist $ hoistModel (single . flatten) . hoist (single . flatten)
+
+flattenSMC2 :: (Monad m) => SequentialT (PopulationT (SMC2 m)) a -> SequentialT (PopulationT (SMC2 m)) a
+ flattenSMC2 = SequentialT.hoist $ single . flatten . PopulationT.hoist (SMC2 . flattenSequentiallyTraced . setup)